Update chatStream.js

server/chatStream.js

@@ -192,6 +192,20 @@ export async function websocketChatStream(body, headers, onToken, abortSignal) {
       if (payload.error) {
         console.error("[WS PAYLOAD ERROR]", payload.error);
         if (onToken) onToken(`[ERROR] ${payload.error}`);
+        
+        // Mark as finished to end the stream
+        if (!finished) {
+          finished = true;
+          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;
       }