Update chatStream.js

server/chatStream.js

@@ -13,6 +13,8 @@ const WORKER_PATH = path.join(__dirname, "searchWorker.js");
 let persistentWs = null;
 let wsAuthPromise = null;
 let requestIdCounter = 0;
+let activeStreamHandlers = new Map(); // Track active stream handlers by request ID
+let errorHandlers = new Map(); // Track error handlers by request ID
 
 async function getSafeWebSocket() {
   // If we have a valid persistent connection, return it
@@ -93,6 +95,30 @@ async function getSafeWebSocket() {
       });
     });
 
+    // Set up the global message and error routing after authentication
+    const globalMessageHandler = (data) => {
+      const line = data.toString();
+      // Route to all active stream handlers
+      for (const [id, handler] of activeStreamHandlers.entries()) {
+        if (!id.startsWith('__')) { // Skip metadata keys
+          handler(line);
+        }
+      }
+    };
+
+    const globalErrorHandler = (err) => {
+      console.error("[WS ERROR]", err);
+      // Notify all active streams
+      for (const [id, handler] of errorHandlers.entries()) {
+        handler(err);
+      }
+    };
+
+    persistentWs.on("message", globalMessageHandler);
+    persistentWs.on("error", globalErrorHandler);
+    activeStreamHandlers.set("__messageListener__", globalMessageHandler);
+    activeStreamHandlers.set("__errorHandler__", globalErrorHandler);
+
     wsAuthPromise = null;
     return persistentWs;
   })();
@@ -166,9 +192,15 @@ export async function websocketChatStream(body, headers, onToken, abortSignal) {
   let finished = false;
 
   return new Promise((resolve, reject) => {
-    const messageHandler = (data) => {
-      const line = data.toString();
-      
+    const cleanup = () => {
+      activeStreamHandlers.delete(currentRequestId);
+      errorHandlers.delete(currentRequestId);
+      if (abortSignal) {
+        abortSignal.removeEventListener("abort", abortHandler);
+      }
+    };
+
+    const messageHandler = (line) => {
       // Parse request ID prefix (format: "id:payload")
       const colonIdx = line.indexOf(':');
       if (colonIdx === -1) {
@@ -196,14 +228,12 @@ export async function websocketChatStream(body, headers, onToken, abortSignal) {
         // Mark as finished to end the stream
         if (!finished) {
           finished = true;
+          cleanup();
           const toolCalls = [...toolCallBuffer.values()].map((t) => ({
             id: t.id || `call_${crypto.randomUUID()}`,
             type: "function",
             function: { name: t.name, arguments: t.arguments },
           }));
-          ws.removeListener("message", messageHandler);
-          ws.removeListener("error", errorHandler);
-          if (abortSignal) abortSignal.removeEventListener("abort", abortHandler);
           resolve({ assistantText, toolCalls });
         }
         return;
@@ -226,28 +256,23 @@ export async function websocketChatStream(body, headers, onToken, abortSignal) {
       }
 
       if (payload.choices?.[0]?.finish_reason) {
-        finished = true;
-        const toolCalls = [...toolCallBuffer.values()].map((t) => ({
-          id: t.id || `call_${crypto.randomUUID()}`,
-          type: "function",
-          function: { name: t.name, arguments: t.arguments },
-        }));
-        
-        // Remove listener and resolve
-        ws.removeListener("message", messageHandler);
-        ws.removeListener("error", errorHandler);
-        if (abortSignal) abortSignal.removeEventListener("abort", abortHandler);
-        resolve({ assistantText, toolCalls });
+        if (!finished) {
+          finished = true;
+          cleanup();
+          const toolCalls = [...toolCallBuffer.values()].map((t) => ({
+            id: t.id || `call_${crypto.randomUUID()}`,
+            type: "function",
+            function: { name: t.name, arguments: t.arguments },
+          }));
+          resolve({ assistantText, toolCalls });
+        }
       }
     };
 
     const errorHandler = (err) => {
-      console.error("[WS ERROR]", err);
       if (!finished) {
         finished = true;
-        ws.removeListener("message", messageHandler);
-        ws.removeListener("error", errorHandler);
-        if (abortSignal) abortSignal.removeEventListener("abort", abortHandler);
+        cleanup();
         reject(err);
       }
     };
@@ -255,14 +280,14 @@ export async function websocketChatStream(body, headers, onToken, abortSignal) {
     const abortHandler = () => {
       if (!finished) {
         finished = true;
-        ws.removeListener("message", messageHandler);
-        ws.removeListener("error", errorHandler);
+        cleanup();
         reject(new Error("AbortError"));
       }
     };
 
-    ws.on("message", messageHandler);
-    ws.on("error", errorHandler);
+    // Register handlers for this request
+    activeStreamHandlers.set(currentRequestId, messageHandler);
+    errorHandlers.set(currentRequestId, errorHandler);
 
     if (abortSignal) {
       abortSignal.addEventListener("abort", abortHandler);