container lifecycle management

gateway/src/channels/websocket-handler.ts

@@ -0,0 +1,161 @@
+import type { FastifyInstance, FastifyRequest } from 'fastify';
+import type { WebSocket } from '@fastify/websocket';
+import type { Authenticator } from '../auth/authenticator.js';
+import { AgentHarness } from '../harness/agent-harness.js';
+import type { InboundMessage } from '../types/messages.js';
+import { randomUUID } from 'crypto';
+
+import type { ProviderConfig } from '../llm/provider.js';
+
+export interface WebSocketHandlerConfig {
+  authenticator: Authenticator;
+  providerConfig: ProviderConfig;
+}
+
+/**
+ * WebSocket channel handler
+ */
+export class WebSocketHandler {
+  private config: WebSocketHandlerConfig;
+  private sessions = new Map<string, AgentHarness>();
+
+  constructor(config: WebSocketHandlerConfig) {
+    this.config = config;
+  }
+
+  /**
+   * Register WebSocket routes
+   */
+  register(app: FastifyInstance): void {
+    app.get(
+      '/ws/chat',
+      { websocket: true },
+      async (socket: WebSocket, request: FastifyRequest) => {
+        await this.handleConnection(socket, request, app);
+      }
+    );
+  }
+
+  /**
+   * Handle WebSocket connection
+   */
+  private async handleConnection(
+    socket: WebSocket,
+    request: FastifyRequest,
+    app: FastifyInstance
+  ): Promise<void> {
+    const logger = app.log;
+
+    // Send initial connecting message
+    socket.send(
+      JSON.stringify({
+        type: 'status',
+        status: 'authenticating',
+        message: 'Authenticating...',
+      })
+    );
+
+    // Authenticate (this may take time if creating container)
+    const authContext = await this.config.authenticator.authenticateWebSocket(request);
+    if (!authContext) {
+      logger.warn('WebSocket authentication failed');
+      socket.send(
+        JSON.stringify({
+          type: 'error',
+          message: 'Authentication failed',
+        })
+      );
+      socket.close(1008, 'Authentication failed');
+      return;
+    }
+
+    logger.info(
+      { userId: authContext.userId, sessionId: authContext.sessionId },
+      'WebSocket connection authenticated'
+    );
+
+    // Send workspace starting message
+    socket.send(
+      JSON.stringify({
+        type: 'status',
+        status: 'initializing',
+        message: 'Starting your workspace...',
+      })
+    );
+
+    // Create agent harness
+    const harness = new AgentHarness({
+      userId: authContext.userId,
+      sessionId: authContext.sessionId,
+      license: authContext.license,
+      providerConfig: this.config.providerConfig,
+      logger,
+    });
+
+    try {
+      await harness.initialize();
+      this.sessions.set(authContext.sessionId, harness);
+
+      // Send connected message
+      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 {
+          const payload = JSON.parse(data.toString());
+
+          if (payload.type === 'message') {
+            const inboundMessage: InboundMessage = {
+              messageId: randomUUID(),
+              userId: authContext.userId,
+              sessionId: authContext.sessionId,
+              content: payload.content,
+              attachments: payload.attachments,
+              timestamp: new Date(),
+            };
+
+            const response = await harness.handleMessage(inboundMessage);
+
+            socket.send(
+              JSON.stringify({
+                type: 'message',
+                ...response,
+              })
+            );
+          }
+        } catch (error) {
+          logger.error({ error }, 'Error handling WebSocket message');
+          socket.send(
+            JSON.stringify({
+              type: 'error',
+              message: 'Failed to process message',
+            })
+          );
+        }
+      });
+
+      // Handle disconnection
+      socket.on('close', async () => {
+        logger.info({ sessionId: authContext.sessionId }, 'WebSocket disconnected');
+        await harness.cleanup();
+        this.sessions.delete(authContext.sessionId);
+      });
+
+      socket.on('error', (error) => {
+        logger.error({ error, sessionId: authContext.sessionId }, 'WebSocket error');
+      });
+    } catch (error) {
+      logger.error({ error }, 'Failed to initialize agent harness');
+      socket.close(1011, 'Internal server error');
+      await harness.cleanup();
+    }
+  }
+}