Update chatStream.js

server/chatStream.js

@@ -55,9 +55,9 @@ function makeClient(accessToken, clientId) {
 
 export async function websocketChatStream(body, headers, onToken, abortSignal) {
   const wsURL =
-    (LIGHTNING_BASE.startsWith("https")
-      ? LIGHTNING_BASE.replace("https", "wss")
-      : LIGHTNING_BASE.replace("http", "ws")) +
+    (process.env.LIGHTNING_BASE.startsWith("https")
+      ? process.env.LIGHTNING_BASE.replace("https", "wss")
+      : process.env.LIGHTNING_BASE.replace("http", "ws")) +
     "/ws/chat";
 
   const ws = new WebSocket(wsURL);
@@ -76,12 +76,15 @@ export async function websocketChatStream(body, headers, onToken, abortSignal) {
       try {
         const msg = JSON.parse(data.toString());
         if (msg.type === "auth" && msg.status === "ok") resolve();
-      } catch {}
+        if (msg.error) reject(new Error(`WebSocket auth error: ${msg.error}`));
+      } catch (e) {
+        reject(e);
+      }
     });
     ws.on("error", reject);
   });
 
-  // Send chat body
+  // Send body
   ws.send(JSON.stringify({ body, headers }));
 
   let assistantText = "";
@@ -97,31 +100,28 @@ export async function websocketChatStream(body, headers, onToken, abortSignal) {
 
     ws.on("message", (data) => {
       const line = data.toString();
-
-      if (!line.startsWith("data:")) return;
-
       let payload;
+
       try {
-        payload = JSON.parse(line.slice(5));
+        payload = JSON.parse(line);
       } catch {
+        // ignore non-JSON lines
         return;
       }
 
-      // ⚠️ Log any server error
+      // Handle errors
       if (payload.error) {
-        console.error("Received error from server:", payload.error);
+        onToken(`[ERROR] ${payload.error}`);
+        return;
       }
 
-      // Optional: log all payloads for debugging
-      // console.log("Server payload:", payload);
-
       const delta = payload.choices?.[0]?.delta;
       if (!delta) return;
 
-      // Append assistant text
+      // Stream text content
       if (delta.content) {
         assistantText += delta.content;
-        if (onToken) onToken(delta.content);
+        onToken(delta.content);
       }
 
       // Handle tool calls
@@ -135,7 +135,7 @@ export async function websocketChatStream(body, headers, onToken, abortSignal) {
         }
       }
 
-      // Finish reason → resolve promise
+      // Finish
       if (payload.choices?.[0]?.finish_reason) {
         ws.close();
 
@@ -155,9 +155,10 @@ export async function websocketChatStream(body, headers, onToken, abortSignal) {
       }
     });
 
-    ws.on("error", (err) => {
-      console.error("WebSocket error:", err);
-      reject(err);
+    ws.on("error", reject);
+    ws.on("close", () => {
+      // If close happens without finish_reason
+      if (!assistantText) reject(new Error("WebSocket closed prematurely"));
     });
   });
 }