Update chatStream.js

server/chatStream.js

@@ -54,38 +54,63 @@ 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);
+
   const ws = new WebSocket(wsURL);
 
-  // Wait for connection open
   await new Promise((resolve, reject) => {
-    ws.on("open", resolve);
-    ws.on("error", reject);
+    ws.on("open", () => {
+      console.log("[WS] Connection opened");
+      resolve();
+    });
+    ws.on("error", (err) => {
+      console.error("[WS] Connection error:", err);
+      reject(err);
+    });
   });
 
   // Authenticate
   ws.send(JSON.stringify({ key: process.env.WEBSOCKET_KEY }));
+  console.log("[WS] Auth sent");
 
   await new Promise((resolve, reject) => {
     ws.on("message", (data) => {
+      const msgStr = data.toString();
+      console.log("[WS] Auth message:", msgStr);
       try {
-        const msg = JSON.parse(data.toString());
-        if (msg.type === "auth" && msg.status === "ok") resolve();
-        if (msg.error) reject(new Error(`WebSocket auth error: ${msg.error}`));
+        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);
       }
     });
-    ws.on("error", reject);
+    ws.on("error", (err) => {
+      console.error("[WS] Auth error event:", err);
+      reject(err);
+    });
   });
 
-  // Send body
+  // Send request body
   ws.send(JSON.stringify({ body, headers }));
+  console.log("[WS] Body sent");
 
   let assistantText = "";
   const toolCallBuffer = new Map();
@@ -93,6 +118,7 @@ export async function websocketChatStream(body, headers, onToken, abortSignal) {
   return new Promise((resolve, reject) => {
     if (abortSignal) {
       abortSignal.addEventListener("abort", () => {
+        console.log("[WS] Aborted by signal");
         ws.close();
         reject(new Error("AbortError"));
       });
@@ -105,20 +131,24 @@ export async function websocketChatStream(body, headers, onToken, abortSignal) {
       try {
         payload = JSON.parse(line);
       } catch {
-        // ignore non-JSON lines
+        console.log("[WS] Non-JSON message received:", line);
         return;
       }
 
+      // Log every event
+      console.log("[WS] Payload received:", JSON.stringify(payload));
+
       // Handle errors
       if (payload.error) {
-        onToken(`[ERROR] ${payload.error}`);
+        console.error("[WS] Payload error:", payload.error);
+        onToken(`[ERROR] ${payload.error}`); // still sent to client
         return;
       }
 
       const delta = payload.choices?.[0]?.delta;
       if (!delta) return;
 
-      // Stream text content
+      // Stream content only through onToken, but do not log the token
       if (delta.content) {
         assistantText += delta.content;
         onToken(delta.content);
@@ -132,33 +162,36 @@ 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
       if (payload.choices?.[0]?.finish_reason) {
+        console.log("[WS] Finish reason:", payload.choices[0].finish_reason);
         ws.close();
 
         const toolCalls = [...toolCallBuffer.values()].map((t) => ({
           id: t.id || `call_${crypto.randomUUID()}`,
           type: "function",
-          function: {
-            name: t.name,
-            arguments: t.arguments,
-          },
+          function: { name: t.name, arguments: t.arguments },
         }));
 
-        resolve({
-          assistantText,
-          toolCalls,
-        });
+        resolve({ assistantText, toolCalls });
       }
     });
 
-    ws.on("error", reject);
-    ws.on("close", () => {
-      // If close happens without finish_reason
-      if (!assistantText) reject(new Error("WebSocket closed prematurely"));
+    ws.on("error", (err) => {
+      console.error("[WS] Error event:", err);
+      // still propagate error to caller
+      reject(err);
+    });
+
+    ws.on("close", (code, reason) => {
+      console.log(`[WS] Connection closed (code: ${code}, reason: ${reason})`);
+      if (!assistantText) {
+        reject(new Error("WebSocket closed prematurely"));
+      }
     });
   });
 }