client-py connected

gateway/src/channels/websocket-handler.ts

@@ -158,21 +158,25 @@ export class WebSocketHandler {
       }),
     };
 
-    // Create agent harness
-    const harness = new AgentHarness({
-      userId: authContext.userId,
-      sessionId: authContext.sessionId,
-      license: authContext.license,
-      providerConfig: this.config.providerConfig,
-      logger,
-    });
+    // Declare harness outside try block so it's available in catch
+    let harness: AgentHarness | undefined;
 
     try {
-      // Initialize workspace and harness
+      // Initialize workspace first
       await workspace.initialize();
       workspace.setAdapter(wsAdapter);
       this.workspaces.set(authContext.sessionId, workspace);
 
+      // Create agent harness with workspace manager
+      harness = new AgentHarness({
+        userId: authContext.userId,
+        sessionId: authContext.sessionId,
+        license: authContext.license,
+        providerConfig: this.config.providerConfig,
+        logger,
+        workspaceManager: workspace,
+      });
+
       await harness.initialize();
       this.harnesses.set(authContext.sessionId, harness);
 
@@ -220,8 +224,8 @@ export class WebSocketHandler {
           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
+          if (payload.type === 'message' || payload.type === 'agent_user_message') {
+            // Chat message - send to agent harness with streaming
             const inboundMessage: InboundMessage = {
               messageId: randomUUID(),
               userId: authContext.userId,
@@ -231,14 +235,41 @@ export class WebSocketHandler {
               timestamp: new Date(),
             };
 
-            const response = await harness.handleMessage(inboundMessage);
+            if (!harness) {
+              logger.error('Harness not initialized');
+              socket.send(JSON.stringify({ type: 'error', message: 'Session not ready' }));
+              return;
+            }
 
-            socket.send(
-              JSON.stringify({
-                type: 'message',
-                ...response,
-              })
-            );
+            // Stream response chunks to client
+            try {
+              for await (const chunk of harness.streamMessage(inboundMessage)) {
+                socket.send(
+                  JSON.stringify({
+                    type: 'agent_chunk',
+                    content: chunk,
+                    done: false,
+                  })
+                );
+              }
+
+              // Send final chunk with done flag
+              socket.send(
+                JSON.stringify({
+                  type: 'agent_chunk',
+                  content: '',
+                  done: true,
+                })
+              );
+            } catch (error) {
+              logger.error({ error }, 'Error streaming response');
+              socket.send(
+                JSON.stringify({
+                  type: 'error',
+                  message: 'Failed to generate response',
+                })
+              );
+            }
           } else if (payload.type === 'hello') {
             // Workspace sync: hello message
             logger.debug({ seqs: payload.seqs }, 'Handling workspace hello');
@@ -280,8 +311,10 @@ export class WebSocketHandler {
         this.workspaces.delete(authContext.sessionId);
 
         // Cleanup harness
-        await harness.cleanup();
-        this.harnesses.delete(authContext.sessionId);
+        if (harness) {
+          await harness.cleanup();
+          this.harnesses.delete(authContext.sessionId);
+        }
       });
 
       socket.on('error', (error: any) => {
@@ -292,7 +325,9 @@ export class WebSocketHandler {
       socket.close(1011, 'Internal server error');
       await workspace.shutdown();
       this.workspaces.delete(authContext.sessionId);
-      await harness.cleanup();
+      if (harness) {
+        await harness.cleanup();
+      }
     }
   }