Server stuff

server/chatStream.js

@@ -18,21 +18,71 @@ const WORKER_PATH = path.join(__dirname, "searchWorker.js");
 // Persistent WebSocket pool
 let persistentWs = null;
 let wsAuthPromise = null;
-let requestIdCounter = 0;
+let requestIdCounter = 0; // Upstream request IDs are scoped to the current websocket connection.
 let activeStreamHandlers = new Map(); // Track active stream handlers by request ID
 let errorHandlers = new Map(); // Track error handlers by request ID
 
+function buildUpstreamSocketMessage(reason, err = null) {
+  return [reason, err?.message].filter(Boolean).join(": ");
+}
+
+function invalidatePersistentWebSocket(ws, reason, err = null) {
+  if (!ws || persistentWs !== ws) return;
+
+  persistentWs = null;
+  wsAuthPromise = null;
+  requestIdCounter = 0;
+  activeStreamHandlers.delete("__messageListener__");
+  activeStreamHandlers.delete("__errorHandler__");
+  activeStreamHandlers.delete("__closeHandler__");
+
+  const failure =
+    err?.name === "RetryableUpstreamConnectionError"
+      ? err
+      : new RetryableUpstreamConnectionError(
+          "The model connection was interrupted. Reconnecting automatically.",
+          buildUpstreamSocketMessage(reason, err)
+        );
+
+  const pendingErrorHandlers = [...errorHandlers.values()];
+  for (const handler of pendingErrorHandlers) {
+    try {
+      handler(failure);
+    } catch {
+      // Ignore request-specific cleanup errors while invalidating the shared socket.
+    }
+  }
+
+  try {
+    ws.removeAllListeners("message");
+    ws.removeAllListeners("error");
+    ws.removeAllListeners("close");
+  } catch {
+    // Ignore listener cleanup failures.
+  }
+
+  if (ws.readyState === WebSocket.OPEN || ws.readyState === WebSocket.CONNECTING) {
+    try {
+      ws.terminate();
+    } catch {
+      // Ignore termination failures while resetting the shared socket.
+    }
+  }
+}
+
 async function getSafeWebSocket() {
   if (persistentWs && persistentWs.readyState === WebSocket.OPEN) return persistentWs;
   if (wsAuthPromise) return wsAuthPromise;
 
   wsAuthPromise = (async () => {
+    const lightningBase = process.env.LIGHTNING_BASE || LIGHTNING_BASE;
     const wsURL =
-      (process.env.LIGHTNING_BASE.startsWith("https")
-        ? process.env.LIGHTNING_BASE.replace("https", "wss")
-        : process.env.LIGHTNING_BASE.replace("http", "ws")) + "/ws/chat";
+      (lightningBase.startsWith("https")
+        ? lightningBase.replace("https", "wss")
+        : lightningBase.replace("http", "ws")) + "/ws/chat";
 
-    persistentWs = new WebSocket(wsURL);
+    const ws = new WebSocket(wsURL);
+    persistentWs = ws;
 
     const safeParse = (str) => {
       try {
@@ -45,50 +95,112 @@ async function getSafeWebSocket() {
     };
 
     await new Promise((resolve, reject) => {
-      const timer = setTimeout(() => reject(new Error("WS connection timeout")), 5000);
-      persistentWs.on("open", () => {
+      const cleanup = () => {
         clearTimeout(timer);
+        ws.removeListener("open", onOpen);
+        ws.removeListener("error", onError);
+        ws.removeListener("close", onClose);
+      };
+      const timer = setTimeout(() => {
+        cleanup();
+        if (persistentWs === ws) {
+          persistentWs = null;
+          wsAuthPromise = null;
+        }
+        try {
+          ws.terminate();
+        } catch {
+          // Ignore termination failures while timing out the handshake.
+        }
+        reject(new Error("WS connection timeout"));
+      }, 5000);
+      const onOpen = () => {
+        cleanup();
         resolve();
-      });
-      persistentWs.on("error", (err) => {
+      };
+      const onError = (err) => {
+        cleanup();
         console.error("[WS] Connection error", err);
-        clearTimeout(timer);
-        persistentWs = null;
-        wsAuthPromise = null;
+        if (persistentWs === ws) {
+          persistentWs = null;
+          wsAuthPromise = null;
+        }
         reject(err);
-      });
+      };
+      const onClose = (code, reasonBuffer) => {
+        cleanup();
+        const reason = reasonBuffer?.toString?.() || "";
+        if (persistentWs === ws) {
+          persistentWs = null;
+          wsAuthPromise = null;
+        }
+        reject(new Error(`WS connection closed (${code})${reason ? `: ${reason}` : ""}`));
+      };
+      ws.on("open", onOpen);
+      ws.on("error", onError);
+      ws.on("close", onClose);
     });
 
-    persistentWs.send(JSON.stringify({ key: process.env.WEBSOCKET_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);
+      const cleanup = () => {
+        clearTimeout(timer);
+        ws.removeListener("message", authHandler);
+        ws.removeListener("error", onError);
+        ws.removeListener("close", onClose);
+      };
+      const timer = setTimeout(() => {
+        cleanup();
+        if (persistentWs === ws) {
+          persistentWs = null;
+          wsAuthPromise = null;
+        }
+        try {
+          ws.terminate();
+        } catch {
+          // Ignore termination failures while timing out auth.
+        }
+        reject(new Error("WS auth timeout"));
+      }, 5000);
       const authHandler = (data) => {
         const msg = safeParse(data.toString());
         if (!msg) return;
         if (msg.type === "auth" && msg.status === "ok") {
-          persistentWs.removeListener("message", authHandler);
-          clearTimeout(timer);
+          cleanup();
           resolve();
         }
         if (msg.error) {
+          cleanup();
           console.error("[WS] Auth error", msg.error);
-          persistentWs.removeListener("message", authHandler);
-          clearTimeout(timer);
-          persistentWs = null;
-          wsAuthPromise = null;
+          if (persistentWs === ws) {
+            persistentWs = null;
+            wsAuthPromise = null;
+          }
           reject(new Error(`WS auth error: ${msg.error}`));
         }
       };
-      persistentWs.on("message", authHandler);
-      persistentWs.on("error", (err) => {
+      const onError = (err) => {
+        cleanup();
         console.error("[WS] Auth error event", err);
-        persistentWs.removeListener("message", authHandler);
-        clearTimeout(timer);
-        persistentWs = null;
-        wsAuthPromise = null;
+        if (persistentWs === ws) {
+          persistentWs = null;
+          wsAuthPromise = null;
+        }
         reject(err);
-      });
+      };
+      const onClose = (code, reasonBuffer) => {
+        cleanup();
+        const reason = reasonBuffer?.toString?.() || "";
+        if (persistentWs === ws) {
+          persistentWs = null;
+          wsAuthPromise = null;
+        }
+        reject(new Error(`WS auth closed (${code})${reason ? `: ${reason}` : ""}`));
+      };
+      ws.on("message", authHandler);
+      ws.on("error", onError);
+      ws.on("close", onClose);
     });
 
     const globalMessageHandler = (data) => {
@@ -100,18 +212,28 @@ async function getSafeWebSocket() {
 
     const globalErrorHandler = (err) => {
       console.error("[WS ERROR]", err);
-      for (const [id, handler] of errorHandlers.entries()) {
-        handler(err);
-      }
+      invalidatePersistentWebSocket(ws, "Upstream websocket error", err);
     };
 
-    persistentWs.on("message", globalMessageHandler);
-    persistentWs.on("error", globalErrorHandler);
+    const globalCloseHandler = (code, reasonBuffer) => {
+      const reason = reasonBuffer?.toString?.() || "";
+      console.warn(`[WS CLOSE] ${code}${reason ? `: ${reason}` : ""}`);
+      invalidatePersistentWebSocket(
+        ws,
+        `Upstream websocket closed (${code})${reason ? `: ${reason}` : ""}`
+      );
+    };
+
+    ws.on("message", globalMessageHandler);
+    ws.on("error", globalErrorHandler);
+    ws.on("close", globalCloseHandler);
+    requestIdCounter = 0;
     activeStreamHandlers.set("__messageListener__", globalMessageHandler);
     activeStreamHandlers.set("__errorHandler__", globalErrorHandler);
+    activeStreamHandlers.set("__closeHandler__", globalCloseHandler);
 
     wsAuthPromise = null;
-    return persistentWs;
+    return ws;
   })();
 
   return wsAuthPromise;
@@ -182,9 +304,13 @@ const HISTORY_SUMMARY_TOKEN_BUDGET = 600;
 const NOTES_TOKEN_BUDGET = 900;
 const MAX_DYNAMIC_MESSAGES = 10;
 const MAX_UPSTREAM_RATE_LIMIT_RETRIES = 4;
+const MAX_UPSTREAM_CONNECTION_RETRIES = 2;
 const DEFAULT_UPSTREAM_RETRY_MS = 4000;
 const MAX_UPSTREAM_RETRY_MS = 15000;
 const UPSTREAM_RETRY_BUFFER_MS = 350;
+const UPSTREAM_FIRST_RESPONSE_TIMEOUT_MS = 15000;
+const UPSTREAM_IDLE_TIMEOUT_MS = 45000;
+const UPSTREAM_STREAM_TIMEOUT_MS = 120000;
 
 // In-memory stores for staged prompt resources and assistant notes
 const promptContextStore = new Map();     // sessionId -> { resources, resourcesById }
@@ -892,6 +1018,15 @@ class RetryableRateLimitError extends Error {
   }
 }
 
+class RetryableUpstreamConnectionError extends Error {
+  constructor(publicMessage, internalMessage = null) {
+    super(publicMessage);
+    this.name = "RetryableUpstreamConnectionError";
+    this.publicMessage = publicMessage;
+    this.internalMessage = internalMessage || publicMessage;
+  }
+}
+
 class UpstreamProviderError extends Error {
   constructor(publicMessage, internalMessage = null) {
     super(publicMessage);
@@ -1038,6 +1173,9 @@ function getPublicErrorMessage(err) {
   if (err.name === "RetryableRateLimitError") {
     return "The model provider is temporarily rate limited. Please try again in a few seconds.";
   }
+  if (err.name === "RetryableUpstreamConnectionError") {
+    return "The model connection was interrupted. Please try again.";
+  }
   if (err.publicMessage) return err.publicMessage;
   return String(err);
 }
@@ -1049,14 +1187,23 @@ async function websocketChatStreamWithRetry(body, headers, onToken, abortSignal)
     } catch (err) {
       if (err?.name === "AbortError") throw err;
 
-      const retryable = err?.name === "RetryableRateLimitError";
-      if (!retryable || retryIndex >= MAX_UPSTREAM_RATE_LIMIT_RETRIES) {
+      const retryable =
+        err?.name === "RetryableRateLimitError" ||
+        err?.name === "RetryableUpstreamConnectionError";
+      const maxRetries =
+        err?.name === "RetryableUpstreamConnectionError"
+          ? MAX_UPSTREAM_CONNECTION_RETRIES
+          : MAX_UPSTREAM_RATE_LIMIT_RETRIES;
+
+      if (!retryable || retryIndex >= maxRetries) {
         throw err;
       }
 
       const waitMs = getRetryDelayMs(err, retryIndex);
       console.warn(
-        `[streamChat] Upstream rate limited, retrying in ${waitMs}ms (${retryIndex + 1}/${MAX_UPSTREAM_RATE_LIMIT_RETRIES})`
+        err?.name === "RetryableUpstreamConnectionError"
+          ? `[streamChat] Upstream websocket interrupted, retrying in ${waitMs}ms (${retryIndex + 1}/${maxRetries})`
+          : `[streamChat] Upstream rate limited, retrying in ${waitMs}ms (${retryIndex + 1}/${maxRetries})`
       );
       await sleepWithAbort(waitMs, abortSignal);
     }
@@ -1080,28 +1227,79 @@ export async function websocketChatStream(body, headers, onToken, abortSignal) {
   const safeParse = (str) => {
     try { return JSON.parse(str.startsWith("data: ") ? str.slice(6) : str); } catch { return null; }
   };
-  ws.send(JSON.stringify({ body, headers }));
   let assistantText = "";
   const toolCallBuffer = new Map();
   let finished = false;
+  let sawAnyPayload = false;
 
   return new Promise((resolve, reject) => {
-    const timeoutId = setTimeout(() => {
+    let inactivityTimeoutId = null;
+    const overallTimeoutId = setTimeout(() => {
       if (!finished) {
         finished = true;
         cleanup();
         const toolCalls = serializeToolCalls(toolCallBuffer);
         resolve({ assistantText, toolCalls });
       }
-    }, 120000);
+    }, UPSTREAM_STREAM_TIMEOUT_MS);
 
     const cleanup = () => {
       activeStreamHandlers.delete(currentRequestId);
       errorHandlers.delete(currentRequestId);
-      clearTimeout(timeoutId);
+      clearTimeout(overallTimeoutId);
+      clearTimeout(inactivityTimeoutId);
       if (abortSignal) abortSignal.removeEventListener("abort", abortHandler);
     };
 
+    const rejectWithSocketReset = (err, reason) => {
+      if (finished) return;
+      finished = true;
+      cleanup();
+      invalidatePersistentWebSocket(ws, reason, err);
+      reject(err);
+    };
+
+    const rejectUnexpectedClose = (publicMessage, internalMessage) => {
+      const hasPartialOutput = assistantText.trim().length > 0 || toolCallBuffer.size > 0;
+      if (!hasPartialOutput) {
+        rejectWithSocketReset(
+          new RetryableUpstreamConnectionError(
+            "The model connection was interrupted. Reconnecting automatically.",
+            internalMessage
+          ),
+          internalMessage
+        );
+        return;
+      }
+
+      if (finished) return;
+      finished = true;
+      cleanup();
+      invalidatePersistentWebSocket(ws, internalMessage);
+      reject(new UpstreamProviderError(publicMessage, internalMessage));
+    };
+
+    const refreshInactivityTimeout = () => {
+      clearTimeout(inactivityTimeoutId);
+      inactivityTimeoutId = setTimeout(() => {
+        if (sawAnyPayload) {
+          rejectUnexpectedClose(
+            "The model provider interrupted the response. Please try again.",
+            "Upstream websocket became idle before the response finished."
+          );
+          return;
+        }
+
+        rejectWithSocketReset(
+          new RetryableUpstreamConnectionError(
+            "The model connection did not respond in time. Reconnecting automatically.",
+            "Upstream websocket produced no response before timeout."
+          ),
+          "Upstream websocket produced no response before timeout."
+        );
+      }, sawAnyPayload ? UPSTREAM_IDLE_TIMEOUT_MS : UPSTREAM_FIRST_RESPONSE_TIMEOUT_MS);
+    };
+
     const messageHandler = (line) => {
       const colonIdx = line.indexOf(':');
       if (colonIdx === -1) return;
@@ -1110,6 +1308,9 @@ export async function websocketChatStream(body, headers, onToken, abortSignal) {
       if (msgRequestId !== String(currentRequestId)) return;
       if (!payload) return;
 
+      sawAnyPayload = true;
+      refreshInactivityTimeout();
+
       if (payload.error && !payload.choices) {
         if (!finished) {
           finished = true;
@@ -1198,12 +1399,53 @@ export async function websocketChatStream(body, headers, onToken, abortSignal) {
       }
     };
 
-    const errorHandler = (err) => { if (!finished) { finished = true; cleanup(); reject(err); } };
-    const abortHandler = () => { if (!finished) { finished = true; cleanup(); reject(new Error("AbortError")); } };
+    const errorHandler = (err) => {
+      if (finished) return;
+      finished = true;
+      cleanup();
+
+      const hasPartialOutput = assistantText.trim().length > 0 || toolCallBuffer.size > 0;
+      if (hasPartialOutput && err?.name === "RetryableUpstreamConnectionError") {
+        reject(
+          new UpstreamProviderError(
+            "The model provider interrupted the response. Please try again.",
+            err?.internalMessage || err?.message || "Upstream websocket interrupted after partial output."
+          )
+        );
+        return;
+      }
+
+      reject(err);
+    };
+    const abortHandler = () => { if (!finished) { finished = true; cleanup(); reject(createAbortError()); } };
 
     activeStreamHandlers.set(currentRequestId, messageHandler);
     errorHandlers.set(currentRequestId, errorHandler);
     if (abortSignal) abortSignal.addEventListener("abort", abortHandler);
+
+    refreshInactivityTimeout();
+
+    try {
+      ws.send(JSON.stringify({ body, headers }), (err) => {
+        if (err) {
+          rejectWithSocketReset(
+            new RetryableUpstreamConnectionError(
+              "The model connection was interrupted. Reconnecting automatically.",
+              buildUpstreamSocketMessage("Failed to send upstream websocket request", err)
+            ),
+            "Failed to send upstream websocket request"
+          );
+        }
+      });
+    } catch (err) {
+      rejectWithSocketReset(
+        new RetryableUpstreamConnectionError(
+          "The model connection was interrupted. Reconnecting automatically.",
+          buildUpstreamSocketMessage("Failed to send upstream websocket request", err)
+        ),
+        "Failed to send upstream websocket request"
+      );
+    }
   });
 }
 
@@ -1347,7 +1589,7 @@ export async function streamChat({
         ...buildModelMessages(baseMessages, workingMessages, sessionId),
         {
           role: "system",
-          content: "Tool-use budget is exhausted for this response. Do not call tools. Answer directly using the information already gathered. If something is still missing, briefly say what is missing without calling tools.",
+          content: "Tool-use budget is exhausted for this response. Do not call tools. Answer directly using the information already gathered. If something is still missing, briefly say what is missing without calling tools. You may ask the user to let you continue your response.",
         },
       ];
 
@@ -1378,18 +1620,18 @@ export async function streamChat({
     const sessionName = extractSessionName(assistantText);
 
     if (typeof onDone === "function") {
-      onDone(assistantText, allToolCalls, false, sessionName, responseEdits, responseSegments);
+      await onDone(assistantText, allToolCalls, false, sessionName, responseEdits, responseSegments);
     }
 
     clearPromptState(sessionId);
 
   } catch (err) {
-    console.error("streamChat error:", err?.internalMessage || err);
     clearPromptState(sessionId);
     if (err.name === "AbortError" || err.message === "AbortError") {
-      if (typeof onDone === "function") onDone(null, null, true, null);
+      if (typeof onDone === "function") await onDone(null, null, true, null);
     } else {
-      if (typeof onError === "function") onError(getPublicErrorMessage(err));
+      console.error("streamChat error:", err?.internalMessage || err);
+      if (typeof onError === "function") await onError(getPublicErrorMessage(err));
     }
   }
 }

server/index.js

@@ -9,7 +9,7 @@ import fetch from 'node-fetch';
 import rateLimit from 'express-rate-limit';
 import fs from 'fs';
 import { registerFeedbackRoutes } from './handleFeedback.js';
-import { handleWsMessage } from './wsHandler.js';
+import { abortActiveStream, handleWsMessage } from './wsHandler.js';
 import { sessionStore, initStoreConfig } from './sessionStore.js';
 import { SUPABASE_URL, SUPABASE_ANON_KEY } from './config.js';
 import { safeSend } from './helpers.js';
@@ -732,8 +732,16 @@ wss.on('connection',(ws,req)=>{
     catch(ex){ console.error("Invalid message error:",ex.message,"\nStack:",ex.stack); safeSend(ws,{type:'error',message:'Invalid message: '+ex.message}); }
   });
 
-  ws.on('close',()=>{ const c=wsClients.get(ws); if(c?.userId) sessionStore.markOffline(c.userId,ws); wsClients.delete(ws); });
-  ws.on('error',()=>wsClients.delete(ws));
+  ws.on('close',()=>{
+    abortActiveStream(ws);
+    const c=wsClients.get(ws);
+    if(c?.userId) sessionStore.markOffline(c.userId,ws);
+    wsClients.delete(ws);
+  });
+  ws.on('error',()=>{
+    abortActiveStream(ws);
+    wsClients.delete(ws);
+  });
   safeSend(ws,{type:'connected', tempId:wsClients.get(ws)?.tempId});
 });
 

server/wsHandler.js

@@ -51,6 +51,12 @@ const CONTINUE_ASSISTANT_PROMPT =
   'Continue your previous response exactly where it left off. Do not restart, summarize, or repeat the opening. Preserve the same formatting and only add the missing continuation.';
 const FREE_WEB_SEARCH_LIMIT = 15;
 
+export function abortActiveStream(ws) {
+  if (!activeStreams.has(ws)) return;
+  activeStreams.get(ws).abort();
+  activeStreams.delete(ws);
+}
+
 initGuestRequestLimiter().catch(err => console.error('Failed to initialize guest request limiter:', err));
 
 function usageOwnerKey(client, clientId = '') {
@@ -306,7 +312,7 @@ const handlers = {
       : sessionStore.getTempSession(client.tempId, sessionId);
     if (!session) return safeSend(ws, { type: 'error', message: 'Session not found' });
 
-    if (activeStreams.has(ws)) activeStreams.get(ws).abort();
+    abortActiveStream(ws);
     const abort = new AbortController();
     activeStreams.set(ws, abort);
     safeSend(ws, { type: 'chat:start', sessionId });
@@ -433,11 +439,12 @@ const handlers = {
         activeStreams.delete(ws);
         console.error('streamChat error:', err);
         safeSend(ws, { type: 'chat:error', error: String(err), sessionId });
+        safeSend(ws, { type: 'chat:aborted', sessionId, reason: 'error' });
       },
     });
   },
 
-  'chat:stop': (ws) => { if (activeStreams.has(ws)) { activeStreams.get(ws).abort(); activeStreams.delete(ws); } },
+  'chat:stop': (ws) => { abortActiveStream(ws); },
 
   'chat:editMessage': async (ws, msg, client) => {
     const { sessionId, messageIndex, newContent } = msg;
@@ -552,7 +559,7 @@ const handlers = {
       return safeSend(ws, { type: 'error', message: 'Assistant message not found' });
     }
 
-    if (activeStreams.has(ws)) activeStreams.get(ws).abort();
+    abortActiveStream(ws);
     const abort = new AbortController();
     activeStreams.set(ws, abort);
 
@@ -687,6 +694,7 @@ const handlers = {
         activeStreams.delete(ws);
         console.error('assistant action streamChat error:', err);
         safeSend(ws, { type: 'chat:error', error: String(err), sessionId });
+        safeSend(ws, { type: 'chat:aborted', sessionId, reason: 'error' });
       },
     });
   },