Update chatStream.js

server/chatStream.js

@@ -4,6 +4,7 @@ import { fileURLToPath } from "url";
 import path from "path";
 import { LIGHTNING_BASE } from "./config.js";
 import WebSocket from "ws";
+import crypto from "crypto";
 
 const __dirname = path.dirname(fileURLToPath(import.meta.url));
 const WORKER_PATH = path.join(__dirname, "searchWorker.js");
@@ -53,61 +54,64 @@ function makeClient(accessToken, clientId) {
   });
 }
 
-export async function websocketChatStream(body, headers, onToken, abortSignal) {
+export async function websocketChatStreamVerbose(body, headers, onToken, abortSignal) {
   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);
 
-  // Utility to safely parse JSON
   const safeParse = (str) => {
-    try { return JSON.parse(str); } 
-    catch { return null; }
+    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 open");
       clearTimeout(timer);
-      console.log("[WS] Opened connection");
       resolve();
     });
     ws.on("error", (err) => {
+      console.error("[WS] Connection error", err);
       clearTimeout(timer);
-      console.error("[WS] Connection error:", err);
       reject(err);
     });
   });
 
-  // Authenticate with timeout
+  // Authenticate
+  console.log("[WS] Sending auth key");
   ws.send(JSON.stringify({ key: process.env.WEBSOCKET_KEY }));
+
   await new Promise((resolve, reject) => {
     const timer = setTimeout(() => reject(new Error("WS auth timeout")), 5000);
     ws.on("message", (data) => {
-      const msg = safeParse(data.toString());
       console.log("[WS RAW MESSAGE]", data.toString());
+      const msg = safeParse(data.toString());
       if (!msg) return;
+
       if (msg.type === "auth" && msg.status === "ok") {
-        clearTimeout(timer);
         console.log("[WS] Auth successful");
+        clearTimeout(timer);
         resolve();
       }
       if (msg.error) {
+        console.error("[WS] Auth error", 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 the chat body
+  console.log("[WS] Sending chat body", JSON.stringify(body, null, 2));
   ws.send(JSON.stringify({ body, headers }));
 
   let assistantText = "";
@@ -115,7 +119,6 @@ export async function websocketChatStream(body, headers, onToken, abortSignal) {
   let finished = false;
 
   return new Promise((resolve, reject) => {
-    // Abort handling
     if (abortSignal) {
       abortSignal.addEventListener("abort", () => {
         console.log("[WS] Aborted by signal");
@@ -124,31 +127,32 @@ export async function websocketChatStream(body, headers, onToken, abortSignal) {
       });
     }
 
-    // Handle messages
     ws.on("message", (data) => {
       const line = data.toString();
       console.log("[WS MESSAGE RECEIVED]", line);
-      let payload = safeParse(line);
-      if (!payload) return;
 
-      // Handle errors from server
+      const payload = safeParse(line);
+      if (!payload) {
+        console.warn("[WS] Failed to parse JSON:", line);
+        return;
+      }
+
+      // Log server-side errors
       if (payload.error) {
         console.error("[WS PAYLOAD ERROR]", payload.error);
-        if (typeof onToken === "function") onToken(`[ERROR] ${payload.error}`);
+        if (onToken) onToken(`[ERROR] ${payload.error}`);
         return;
       }
 
       const delta = payload.choices?.[0]?.delta;
-      if (!delta) return;
-
-      // Stream content tokens
-      if (delta.content) {
+      if (delta?.content) {
         assistantText += delta.content;
-        if (typeof onToken === "function") onToken(delta.content);
+        console.log("[WS TOKEN STREAM]", delta.content);
+        if (onToken) onToken(delta.content);
       }
 
-      // Collect tool calls
-      if (delta.tool_calls) {
+      if (delta?.tool_calls) {
+        console.log("[WS TOOL CALLS]", delta.tool_calls);
         for (const call of delta.tool_calls) {
           const entry = toolCallBuffer.get(call.index) ?? { arguments: "" };
           if (call.id) entry.id = call.id;
@@ -158,7 +162,6 @@ export async function websocketChatStream(body, headers, onToken, abortSignal) {
         }
       }
 
-      // Finish detection
       if (payload.choices?.[0]?.finish_reason) {
         finished = true;
         ws.close();
@@ -167,7 +170,7 @@ export async function websocketChatStream(body, headers, onToken, abortSignal) {
           type: "function",
           function: { name: t.name, arguments: t.arguments },
         }));
-        console.log("[WS] Finished streaming");
+        console.log("[WS] Finished streaming", { assistantText, toolCalls });
         resolve({ assistantText, toolCalls });
       }
     });
@@ -181,6 +184,9 @@ export async function websocketChatStream(body, headers, onToken, abortSignal) {
       console.log(`[WS CLOSED] Code: ${code}, Reason: ${reason}`);
       if (!finished) reject(new Error("WebSocket closed prematurely"));
     });
+
+    ws.on("ping", () => console.log("[WS PING] Received ping"));
+    ws.on("pong", () => console.log("[WS PONG] Received pong"));
   });
 }