Update chatStream.js

server/chatStream.js

@@ -53,55 +53,107 @@ function makeClient(accessToken, clientId) {
   });
 }
 
-async function websocketChatStream(body, headers, onToken) {
+async function websocketChatStream(body, headers, onToken, abortSignal) {
 
   const wsURL =
     (LIGHTNING_BASE.startsWith("https")
       ? LIGHTNING_BASE.replace("https", "wss")
-      : LIGHTNING_BASE.replace("http", "ws")
-    ) + "/ws/chat";
-  await new Promise(resolve => {
+      : LIGHTNING_BASE.replace("http", "ws")) +
+    "/ws/chat";
+
+  const ws = new WebSocket(wsURL);
+
+  await new Promise((resolve, reject) => {
     ws.on("open", resolve);
+    ws.on("error", reject);
   });
 
-  ws.send(JSON.stringify({
-    key: process.env.WEBSOCKET_KEY
-  }));
+  ws.send(
+    JSON.stringify({
+      key: process.env.WEBSOCKET_KEY,
+    })
+  );
 
-  await new Promise(resolve => {
-    ws.onmessage = e => {
-      const msg = JSON.parse(e.data);
-      if (msg.type === "auth") resolve();
-    };
+  await new Promise((resolve, reject) => {
+    ws.on("message", (data) => {
+      try {
+        const msg = JSON.parse(data.toString());
+        if (msg.type === "auth" && msg.status === "ok") resolve();
+      } catch {}
+    });
+    ws.on("error", reject);
   });
 
-  ws.send(JSON.stringify({
-    body,
-    headers
-  }));
+  ws.send(
+    JSON.stringify({
+      body,
+      headers,
+    })
+  );
 
   let assistantText = "";
+  const toolCallBuffer = new Map();
 
   return new Promise((resolve, reject) => {
 
+    if (abortSignal) {
+      abortSignal.addEventListener("abort", () => {
+        ws.close();
+        reject(new Error("AbortError"));
+      });
+    }
+
     ws.on("message", (data) => {
 
       const line = data.toString();
 
       if (!line.startsWith("data:")) return;
 
-      const payload = JSON.parse(line.slice(5));
+      let payload;
+
+      try {
+        payload = JSON.parse(line.slice(5));
+      } catch {
+        return;
+      }
 
       const delta = payload.choices?.[0]?.delta;
+      if (!delta) return;
 
-      if (delta?.content) {
+      if (delta.content) {
         assistantText += delta.content;
         onToken(delta.content);
       }
 
+      if (delta.tool_calls) {
+        for (const call of delta.tool_calls) {
+          const entry = toolCallBuffer.get(call.index) ?? { arguments: "" };
+
+          if (call.id) entry.id = call.id;
+          if (call.function?.name) entry.name = call.function.name;
+          if (call.function?.arguments) entry.arguments += call.function.arguments;
+
+          toolCallBuffer.set(call.index, entry);
+        }
+      }
+
       if (payload.choices?.[0]?.finish_reason) {
+
         ws.close();
-        resolve({ assistantText });
+
+        const toolCalls = [...toolCallBuffer.values()].map((t) => ({
+          id: t.id || `call_${crypto.randomUUID()}`,
+          type: "function",
+          function: {
+            name: t.name,
+            arguments: t.arguments,
+          },
+        }));
+
+        resolve({
+          assistantText,
+          toolCalls,
+        });
       }
     });
 
@@ -128,7 +180,7 @@ export function extractSessionName(text) {
   return name;
 }
 
-export async function streamChat(ws, {
+export async function streamChat({
   sessionId,
   model,
   history,
@@ -175,7 +227,6 @@ export async function streamChat(ws, {
           .trim() || "";
 
     }
-
   }
 
   const hasUserMessage =
@@ -190,12 +241,15 @@ export async function streamChat(ws, {
   ];
 
   if (hasUserMessage) {
-    messages.push({ role: "user", content: normalizedUserMessage });
+    messages.push({
+      role: "user",
+      content: normalizedUserMessage,
+    });
   }
 
   const headers = {
     ...(accessToken ? { Authorization: `Bearer ${accessToken}` } : {}),
-    ...(clientId ? { "X-Client-ID": clientId } : {})
+    ...(clientId ? { "X-Client-ID": clientId } : {}),
   };
 
   try {
@@ -204,7 +258,7 @@ export async function streamChat(ws, {
       model: model || "lightning",
       messages,
       tools: enabledTools.length ? enabledTools : undefined,
-      stream: true
+      stream: true,
     };
 
     let { assistantText, toolCalls } =
@@ -213,7 +267,7 @@ export async function streamChat(ws, {
     if (toolCalls.length > 0) {
 
       const toolResults = await processToolCalls(
-        ws,
+        null,
         toolCalls,
         tools,
         accessToken,
@@ -231,7 +285,7 @@ export async function streamChat(ws, {
       if (hasUserMessage) {
         followUpMessages.push({
           role: "user",
-          content: normalizedUserMessage
+          content: normalizedUserMessage,
         });
       }
 
@@ -239,7 +293,7 @@ export async function streamChat(ws, {
         {
           role: "assistant",
           content: assistantText || "",
-          tool_calls: toolCalls
+          tool_calls: toolCalls,
         },
         ...toolResults
       );
@@ -247,14 +301,13 @@ export async function streamChat(ws, {
       const followUpBody = {
         model: model || "lightning",
         messages: followUpMessages,
-        stream: true
+        stream: true,
       };
 
       const followUp =
         await websocketChatStream(followUpBody, headers, onToken, abortSignal);
 
       assistantText += followUp.assistantText;
-
     }
 
     const sessionName = extractSessionName(assistantText);
@@ -269,11 +322,10 @@ export async function streamChat(ws, {
     ) {
       onDone(null, null, true, null);
     } else {
+      console.error("streamChat error:", err);
       onError(String(err));
     }
-
   }
-
 }
 
 const VALID_ROLES = new Set(["system", "user", "assistant", "tool"]);