chart data loading

gateway/src/channels/websocket-handler.ts

@@ -7,12 +7,36 @@ import { randomUUID } from 'crypto';
 
 import type { ProviderConfig } from '../llm/provider.js';
 import type { SessionRegistry, EventSubscriber, Session } from '../events/index.js';
+import type { OHLCService } from '../services/ohlc-service.js';
+import type { SymbolIndexService } from '../services/symbol-index-service.js';
+import type { ContainerManager } from '../k8s/container-manager.js';
+import {
+  WorkspaceManager,
+  DEFAULT_STORES,
+  type ChannelAdapter,
+  type ChannelCapabilities,
+  type SnapshotMessage,
+  type PatchMessage,
+} from '../workspace/index.js';
+
+/**
+ * Safe JSON stringifier that handles BigInt values
+ * Converts BigInt to Number (safe for timestamps and other integer values)
+ */
+function jsonStringifySafe(obj: any): string {
+  return JSON.stringify(obj, (_key, value) =>
+    typeof value === 'bigint' ? Number(value) : value
+  );
+}
 
 export interface WebSocketHandlerConfig {
   authenticator: Authenticator;
+  containerManager: ContainerManager;
   providerConfig: ProviderConfig;
   sessionRegistry: SessionRegistry;
   eventSubscriber: EventSubscriber;
+  ohlcService?: OHLCService;  // Optional for historical data support
+  symbolIndexService?: SymbolIndexService;  // Optional for symbol search
 }
 
 /**
@@ -24,6 +48,7 @@ export interface WebSocketHandlerConfig {
 export class WebSocketHandler {
   private config: WebSocketHandlerConfig;
   private harnesses = new Map<string, AgentHarness>();
+  private workspaces = new Map<string, WorkspaceManager>();
 
   constructor(config: WebSocketHandlerConfig) {
     this.config = config;
@@ -61,8 +86,8 @@ export class WebSocketHandler {
       })
     );
 
-    // Authenticate (this may take time if creating container)
-    const authContext = await this.config.authenticator.authenticateWebSocket(request);
+    // Authenticate (returns immediately if container is spinning up)
+    const { authContext, isSpinningUp } = await this.config.authenticator.authenticateWebSocket(request);
     if (!authContext) {
       logger.warn('WebSocket authentication failed');
       socket.send(
@@ -76,18 +101,62 @@ export class WebSocketHandler {
     }
 
     logger.info(
-      { userId: authContext.userId, sessionId: authContext.sessionId },
+      { userId: authContext.userId, sessionId: authContext.sessionId, isSpinningUp },
       'WebSocket connection authenticated'
     );
 
-    // Send workspace starting message
-    socket.send(
-      JSON.stringify({
-        type: 'status',
-        status: 'initializing',
-        message: 'Starting your workspace...',
-      })
-    );
+    // If container is spinning up, send status and start background polling
+    if (isSpinningUp) {
+      socket.send(
+        JSON.stringify({
+          type: 'status',
+          status: 'spinning_up',
+          message: 'Your workspace is starting up, please wait...',
+        })
+      );
+
+      // Start background polling for container readiness
+      this.pollContainerReadiness(socket, authContext, app).catch((error) => {
+        logger.error({ error, userId: authContext.userId }, 'Error polling container readiness');
+      });
+
+      // Don't return - continue with session setup so we can receive messages once ready
+    } else {
+      // Send workspace starting message
+      socket.send(
+        JSON.stringify({
+          type: 'status',
+          status: 'initializing',
+          message: 'Starting your workspace...',
+        })
+      );
+    }
+
+    // Create workspace manager for this session
+    const workspace = new WorkspaceManager({
+      userId: authContext.userId,
+      sessionId: authContext.sessionId,
+      stores: DEFAULT_STORES,
+      // containerSync will be added when MCP client is implemented
+      logger,
+    });
+
+    // Create WebSocket channel adapter
+    const wsAdapter: ChannelAdapter = {
+      sendSnapshot: (msg: SnapshotMessage) => {
+        socket.send(JSON.stringify(msg));
+      },
+      sendPatch: (msg: PatchMessage) => {
+        socket.send(JSON.stringify(msg));
+      },
+      getCapabilities: (): ChannelCapabilities => ({
+        supportsSync: true,
+        supportsImages: true,
+        supportsMarkdown: true,
+        supportsStreaming: true,
+        supportsTradingViewEmbed: true,
+      }),
+    };
 
     // Create agent harness
     const harness = new AgentHarness({
@@ -99,6 +168,11 @@ export class WebSocketHandler {
     });
 
     try {
+      // Initialize workspace and harness
+      await workspace.initialize();
+      workspace.setAdapter(wsAdapter);
+      this.workspaces.set(authContext.sessionId, workspace);
+
       await harness.initialize();
       this.harnesses.set(authContext.sessionId, harness);
 
@@ -125,23 +199,29 @@ export class WebSocketHandler {
         'Session registered for events'
       );
 
-      // Send connected message
-      socket.send(
-        JSON.stringify({
-          type: 'connected',
-          sessionId: authContext.sessionId,
-          userId: authContext.userId,
-          licenseType: authContext.license.licenseType,
-          message: 'Connected to Dexorder AI',
-        })
-      );
+      // Send connected message (only if not spinning up - otherwise sent by pollContainerReadiness)
+      if (!isSpinningUp) {
+        socket.send(
+          JSON.stringify({
+            type: 'connected',
+            sessionId: authContext.sessionId,
+            userId: authContext.userId,
+            licenseType: authContext.license.licenseType,
+            message: 'Connected to Dexorder AI',
+          })
+        );
+      }
 
       // Handle messages
       socket.on('message', async (data: Buffer) => {
         try {
+          logger.info({ rawMessage: data.toString().substring(0, 500) }, 'WebSocket message received');
           const payload = JSON.parse(data.toString());
+          logger.info({ type: payload.type, request_id: payload.request_id }, 'WebSocket message parsed');
 
+          // Route based on message type
           if (payload.type === 'message') {
+            // Chat message - send to agent harness
             const inboundMessage: InboundMessage = {
               messageId: randomUUID(),
               userId: authContext.userId,
@@ -159,6 +239,20 @@ export class WebSocketHandler {
                 ...response,
               })
             );
+          } else if (payload.type === 'hello') {
+            // Workspace sync: hello message
+            logger.debug({ seqs: payload.seqs }, 'Handling workspace hello');
+            await workspace.handleHello(payload.seqs || {});
+          } else if (payload.type === 'patch') {
+            // Workspace sync: patch message
+            logger.debug({ store: payload.store, seq: payload.seq }, 'Handling workspace patch');
+            await workspace.handlePatch(payload.store, payload.seq, payload.patch || []);
+          } else if (this.isDatafeedMessage(payload)) {
+            // Historical data request - send to OHLC service
+            logger.info({ type: payload.type }, 'Routing to datafeed handler');
+            await this.handleDatafeedMessage(socket, payload, logger);
+          } else {
+            logger.warn({ type: payload.type }, 'Unknown message type received');
           }
         } catch (error) {
           logger.error({ error }, 'Error handling WebSocket message');
@@ -181,6 +275,10 @@ export class WebSocketHandler {
           await this.config.eventSubscriber.onSessionDisconnect(removedSession);
         }
 
+        // Cleanup workspace
+        await workspace.shutdown();
+        this.workspaces.delete(authContext.sessionId);
+
         // Cleanup harness
         await harness.cleanup();
         this.harnesses.delete(authContext.sessionId);
@@ -190,12 +288,76 @@ export class WebSocketHandler {
         logger.error({ error, sessionId: authContext.sessionId }, 'WebSocket error');
       });
     } catch (error) {
-      logger.error({ error }, 'Failed to initialize agent harness');
+      logger.error({ error }, 'Failed to initialize session');
       socket.close(1011, 'Internal server error');
+      await workspace.shutdown();
+      this.workspaces.delete(authContext.sessionId);
       await harness.cleanup();
     }
   }
 
+  /**
+   * Poll for container readiness in the background
+   * Sends notification to client when container is ready
+   */
+  private async pollContainerReadiness(
+    socket: WebSocket,
+    authContext: any,
+    app: FastifyInstance
+  ): Promise<void> {
+    const logger = app.log;
+    const userId = authContext.userId;
+
+    logger.info({ userId }, 'Starting background poll for container readiness');
+
+    try {
+      // Wait for container to become ready (2 minute timeout)
+      const ready = await this.config.containerManager.waitForContainerReady(userId, 120000);
+
+      if (ready) {
+        logger.info({ userId }, 'Container is now ready, notifying client');
+
+        // Send ready notification
+        socket.send(
+          JSON.stringify({
+            type: 'status',
+            status: 'ready',
+            message: 'Your workspace is ready!',
+          })
+        );
+
+        // Also send the 'connected' message
+        socket.send(
+          JSON.stringify({
+            type: 'connected',
+            sessionId: authContext.sessionId,
+            userId: authContext.userId,
+            licenseType: authContext.license.licenseType,
+            message: 'Connected to Dexorder AI',
+          })
+        );
+      } else {
+        logger.warn({ userId }, 'Container failed to become ready within timeout');
+
+        socket.send(
+          JSON.stringify({
+            type: 'error',
+            message: 'Workspace failed to start. Please try again later.',
+          })
+        );
+      }
+    } catch (error) {
+      logger.error({ error, userId }, 'Error waiting for container readiness');
+
+      socket.send(
+        JSON.stringify({
+          type: 'error',
+          message: 'Error starting workspace. Please try again later.',
+        })
+      );
+    }
+  }
+
   /**
    * Derive the container's XPUB event endpoint from the MCP server URL.
    *
@@ -212,4 +374,173 @@ export class WebSocketHandler {
       return mcpServerUrl.replace('http://', 'tcp://').replace(':3000', ':5570');
     }
   }
+
+  /**
+   * Check if message is a datafeed message (TradingView protocol)
+   */
+  private isDatafeedMessage(payload: any): boolean {
+    const datafeedTypes = [
+      'get_config',
+      'search_symbols',
+      'resolve_symbol',
+      'get_bars',
+      'subscribe_bars',
+      'unsubscribe_bars',
+    ];
+    return datafeedTypes.includes(payload.type);
+  }
+
+  /**
+   * Handle datafeed messages (TradingView protocol)
+   */
+  private async handleDatafeedMessage(
+    socket: WebSocket,
+    payload: any,
+    logger: any
+  ): Promise<void> {
+    logger.info({ type: payload.type, payload }, 'handleDatafeedMessage called');
+    const ohlcService = this.config.ohlcService;
+    const symbolIndexService = this.config.symbolIndexService;
+
+    logger.info({
+      hasOhlcService: !!ohlcService,
+      hasSymbolIndexService: !!symbolIndexService
+    }, 'Service availability');
+
+    if (!ohlcService && !symbolIndexService) {
+      logger.warn('No datafeed services available');
+      return;
+    }
+
+    const requestId = payload.request_id || randomUUID();
+
+    try {
+      switch (payload.type) {
+        case 'get_config': {
+          const config = ohlcService ? await ohlcService.getConfig() : { supported_resolutions: ['1', '5', '15', '60', '1D'] };
+          socket.send(
+            JSON.stringify({
+              type: 'get_config_response',
+              request_id: requestId,
+              config,
+            })
+          );
+          break;
+        }
+
+        case 'search_symbols': {
+          logger.info({ query: payload.query, limit: payload.limit }, 'Handling search_symbols');
+          // Use SymbolIndexService if available, otherwise fallback to OHLCService stub
+          const symbolIndexService = this.config.symbolIndexService;
+          logger.info({ hasSymbolIndexService: !!symbolIndexService }, 'Service check for search');
+
+          const results = symbolIndexService
+            ? await symbolIndexService.search(payload.query, payload.limit || 30)
+            : (ohlcService ? await ohlcService.searchSymbols(
+                payload.query,
+                payload.symbol_type,
+                payload.exchange,
+                payload.limit || 30
+              ) : []);
+
+          logger.info({ resultsCount: results.length }, 'Search complete');
+          socket.send(
+            JSON.stringify({
+              type: 'search_symbols_response',
+              request_id: requestId,
+              results,
+            })
+          );
+          break;
+        }
+
+        case 'resolve_symbol': {
+          logger.info({ symbol: payload.symbol }, 'Handling resolve_symbol');
+          // Use SymbolIndexService if available, otherwise fallback to OHLCService stub
+          const symbolIndexService = this.config.symbolIndexService;
+          logger.info({ hasSymbolIndexService: !!symbolIndexService }, 'Service check for resolve');
+
+          const symbolInfo = symbolIndexService
+            ? await symbolIndexService.resolveSymbol(payload.symbol)
+            : (ohlcService ? await ohlcService.resolveSymbol(payload.symbol) : null);
+
+          logger.info({ found: !!symbolInfo }, 'Symbol resolution complete');
+
+          if (!symbolInfo) {
+            logger.warn({ symbol: payload.symbol }, 'Symbol not found');
+            socket.send(
+              JSON.stringify({
+                type: 'error',
+                request_id: requestId,
+                error_message: `Symbol not found: ${payload.symbol}`,
+              })
+            );
+          } else {
+            logger.info({ symbolInfo }, 'Sending symbol_info response');
+            socket.send(
+              JSON.stringify({
+                type: 'resolve_symbol_response',
+                request_id: requestId,
+                symbol_info: symbolInfo,
+              })
+            );
+          }
+          break;
+        }
+
+        case 'get_bars': {
+          if (!ohlcService) {
+            socket.send(JSON.stringify({
+              type: 'error',
+              request_id: requestId,
+              error_message: 'OHLC service not available'
+            }));
+            break;
+          }
+          const history = await ohlcService.fetchOHLC(
+            payload.symbol,
+            payload.resolution,
+            payload.from_time,
+            payload.to_time,
+            payload.countback
+          );
+          socket.send(
+            jsonStringifySafe({
+              type: 'get_bars_response',
+              request_id: requestId,
+              history,
+            })
+          );
+          break;
+        }
+
+        case 'subscribe_bars':
+        case 'unsubscribe_bars':
+          // TODO: Implement real-time subscriptions
+          socket.send(
+            JSON.stringify({
+              type: `${payload.type}_response`,
+              request_id: requestId,
+              subscription_id: payload.subscription_id,
+              success: false,
+              message: 'Real-time subscriptions not yet implemented',
+            })
+          );
+          break;
+
+        default:
+          logger.warn({ type: payload.type }, 'Unknown datafeed message type');
+      }
+    } catch (error: any) {
+      logger.error({ error, type: payload.type }, 'Error handling datafeed message');
+      socket.send(
+        jsonStringifySafe({
+          type: 'error',
+          request_id: requestId,
+          error_code: 'INTERNAL_ERROR',
+          error_message: error.message || 'Internal server error',
+        })
+      );
+    }
+  }
 }