Server restart issues

server/chatStream.js

@@ -22,17 +22,66 @@ let requestIdCounter = 0;
 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;
+  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 +94,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 +211,27 @@ 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);
     activeStreamHandlers.set("__messageListener__", globalMessageHandler);
     activeStreamHandlers.set("__errorHandler__", globalErrorHandler);
+    activeStreamHandlers.set("__closeHandler__", globalCloseHandler);
 
     wsAuthPromise = null;
-    return persistentWs;
+    return ws;
   })();
 
   return wsAuthPromise;
@@ -182,9 +302,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 +1016,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 +1171,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 +1185,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 +1225,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 +1306,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 +1397,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"
+      );
+    }
   });
 }
 
@@ -1384,11 +1624,11 @@ export async function streamChat({
     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);
     } else {
+      console.error("streamChat error:", err?.internalMessage || err);
       if (typeof onError === "function") 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 });
@@ -437,7 +443,7 @@ const handlers = {
     });
   },
 
-  '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 +558,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);