Update chatStream.js

server/chatStream.js

@@ -54,107 +54,100 @@ function makeClient(accessToken, clientId) {
 }
 
 export async function websocketChatStream(body, headers, onToken, abortSignal) {
-  if (!process.env.LIGHTNING_BASE) {
-    throw new Error("LIGHTNING_BASE environment variable is not defined");
-  }
-
   const wsURL =
     (process.env.LIGHTNING_BASE.startsWith("https")
       ? process.env.LIGHTNING_BASE.replace("https", "wss")
-      : process.env.LIGHTNING_BASE.replace("http", "ws")) +
-    "/ws/chat";
-
-  console.log("[WS] Connecting to", wsURL);
+      : process.env.LIGHTNING_BASE.replace("http", "ws")) + "/ws/chat";
 
   const ws = new WebSocket(wsURL);
 
+  // Utility to safely parse JSON
+  const safeParse = (str) => {
+    try { return JSON.parse(str); } 
+    catch { return null; }
+  };
+
+  // Wait for open with timeout
   await new Promise((resolve, reject) => {
+    const timer = setTimeout(() => reject(new Error("WS connection timeout")), 5000);
     ws.on("open", () => {
-      console.log("[WS] Connection opened");
+      clearTimeout(timer);
+      console.log("[WS] Opened connection");
       resolve();
     });
     ws.on("error", (err) => {
+      clearTimeout(timer);
       console.error("[WS] Connection error:", err);
       reject(err);
     });
   });
 
-  // Authenticate
+  // Authenticate with timeout
   ws.send(JSON.stringify({ key: process.env.WEBSOCKET_KEY }));
-  console.log("[WS] Auth sent");
-
   await new Promise((resolve, reject) => {
+    const timer = setTimeout(() => reject(new Error("WS auth timeout")), 5000);
     ws.on("message", (data) => {
-      const msgStr = data.toString();
-      console.log("[WS] Auth message:", msgStr);
-      try {
-        const msg = JSON.parse(msgStr);
-        if (msg.type === "auth" && msg.status === "ok") {
-          console.log("[WS] Auth successful");
-          resolve();
-        }
-        if (msg.error) {
-          console.error("[WS] Auth error:", msg.error);
-          reject(new Error(`WebSocket auth error: ${msg.error}`));
-        }
-      } catch (e) {
-        console.error("[WS] Auth message parse error:", e);
-        reject(e);
+      const msg = safeParse(data.toString());
+      console.log("[WS RAW MESSAGE]", data.toString());
+      if (!msg) return;
+      if (msg.type === "auth" && msg.status === "ok") {
+        clearTimeout(timer);
+        console.log("[WS] Auth successful");
+        resolve();
+      }
+      if (msg.error) {
+        clearTimeout(timer);
+        console.error("[WS] Auth error:", msg.error);
+        reject(new Error(`WS auth error: ${msg.error}`));
       }
     });
     ws.on("error", (err) => {
-      console.error("[WS] Auth error event:", err);
+      clearTimeout(timer);
       reject(err);
     });
   });
 
-  // Send request body
+  // Send the chat body
   ws.send(JSON.stringify({ body, headers }));
-  console.log("[WS] Body sent");
 
   let assistantText = "";
   const toolCallBuffer = new Map();
+  let finished = false;
 
   return new Promise((resolve, reject) => {
+    // Abort handling
     if (abortSignal) {
       abortSignal.addEventListener("abort", () => {
         console.log("[WS] Aborted by signal");
         ws.close();
-        reject(new Error("AbortError"));
+        if (!finished) reject(new Error("AbortError"));
       });
     }
 
+    // Handle messages
     ws.on("message", (data) => {
       const line = data.toString();
-      let payload;
+      console.log("[WS MESSAGE RECEIVED]", line);
+      let payload = safeParse(line);
+      if (!payload) return;
 
-      try {
-        payload = JSON.parse(line);
-      } catch {
-        console.log("[WS] Non-JSON message received:", line);
-        return;
-      }
-
-      // Log every event
-      console.log("[WS] Payload received:", JSON.stringify(payload));
-
-      // Handle errors
+      // Handle errors from server
       if (payload.error) {
-        console.error("[WS] Payload error:", payload.error);
-        onToken(`[ERROR] ${payload.error}`); // still sent to client
+        console.error("[WS PAYLOAD ERROR]", payload.error);
+        if (typeof onToken === "function") onToken(`[ERROR] ${payload.error}`);
         return;
       }
 
       const delta = payload.choices?.[0]?.delta;
       if (!delta) return;
 
-      // Stream content only through onToken, but do not log the token
+      // Stream content tokens
       if (delta.content) {
         assistantText += delta.content;
-        onToken(delta.content);
+        if (typeof onToken === "function") onToken(delta.content);
       }
 
-      // Handle tool calls
+      // Collect tool calls
       if (delta.tool_calls) {
         for (const call of delta.tool_calls) {
           const entry = toolCallBuffer.get(call.index) ?? { arguments: "" };
@@ -162,36 +155,31 @@ export async function websocketChatStream(body, headers, onToken, abortSignal) {
           if (call.function?.name) entry.name = call.function.name;
           if (call.function?.arguments) entry.arguments += call.function.arguments;
           toolCallBuffer.set(call.index, entry);
-          console.log("[WS] Tool call received:", JSON.stringify(entry));
         }
       }
 
-      // Finish
+      // Finish detection
       if (payload.choices?.[0]?.finish_reason) {
-        console.log("[WS] Finish reason:", payload.choices[0].finish_reason);
+        finished = true;
         ws.close();
-
         const toolCalls = [...toolCallBuffer.values()].map((t) => ({
           id: t.id || `call_${crypto.randomUUID()}`,
           type: "function",
           function: { name: t.name, arguments: t.arguments },
         }));
-
+        console.log("[WS] Finished streaming");
         resolve({ assistantText, toolCalls });
       }
     });
 
     ws.on("error", (err) => {
-      console.error("[WS] Error event:", err);
-      // still propagate error to caller
-      reject(err);
+      console.error("[WS ERROR]", err);
+      if (!finished) reject(err);
     });
 
     ws.on("close", (code, reason) => {
-      console.log(`[WS] Connection closed (code: ${code}, reason: ${reason})`);
-      if (!assistantText) {
-        reject(new Error("WebSocket closed prematurely"));
-      }
+      console.log(`[WS CLOSED] Code: ${code}, Reason: ${reason}`);
+      if (!finished) reject(new Error("WebSocket closed prematurely"));
     });
   });
 }