Harden chat WebSocket for mobile lifecycle (#23) (#25)

## Summary
- Closes #23
- Add WebSocket reconnect with exponential backoff + jitter (capped attempts), heartbeat keepalives, and reconnect on browser visibility/`online` plus Capacitor `App`/`Network` plugins when present
- Refresh JWT before reconnect (coordinates with #22) so conversation refetch after a drop still works
- Mid-stream drop recovery: mark interrupted + refetch conversation; surface connecting/reconnecting/disconnected state and disable send while down

## Test plan
- [x] Unit tests: `websocketReconnect` helpers + `WebSocketProvider` reconnect/resume/send (`npm run test:ci`)
- [ ] Background app ~30s and ~10min, then resume — chat reconnects without reload
- [ ] Airplane mode on/off and cellular↔WiFi — reconnect indicator then restored send
- [ ] Drop mid-stream response — interrupted banner + conversation refetch
- [ ] Confirm nginx/`/ws/` proxy idle timeout stays above heartbeat (25s) or allows ping payloadsReviewed-on: #25

llm-fe/src/llm-fe/contexts/MessageContext.tsx

@@ -16,6 +16,9 @@ type IMessageContext = {
     conversationDetails: ConversationPrompt[];
     setConversationDetails: (conversationPrompts: ConversationPrompt[]) => void;
     isGeneratingMessage: boolean;
+    /** True when a stream was interrupted by a socket drop; cleared on refetch/retry. */
+    streamInterrupted: boolean;
+    clearStreamInterrupted: () => void;
 }
 
 const initialValues = {
@@ -23,59 +26,116 @@ const initialValues = {
     setStateMessage: () => {},
     conversationDetails: [],
     setConversationDetails: () => {},
-    isGeneratingMessage: false
+    isGeneratingMessage: false,
+    streamInterrupted: false,
+    clearStreamInterrupted: () => {},
 }
 
 const MessageContext = createContext<IMessageContext>(initialValues);
 
+const STREAM_INTERRUPTED_MESSAGE =
+    '_Response interrupted — connection lost. Re-send your message if needed._';
+
 const MessageProvider = ( {children}: MessageProviderProps) => {
-    const [subscribe, unsubscribe]= useContext(WebSocketContext)
+    const [subscribe, unsubscribe, , , , , reconnectGeneration] = useContext(WebSocketContext)
     const { account } = useContext(AccountContext)
     const {selectedConversation, setSelectedConversation} = useContext(ConversationContext);
 
     const [stateMessage, setStateMessage] = useState<string>('')
     const [conversationDetails, setConversationDetails] = useState<ConversationPrompt[]>([])
     const [isGeneratingMessage, setIsGeneratingMessage] = useState<boolean>(false)
+    const [streamInterrupted, setStreamInterrupted] = useState<boolean>(false)
 
     const messageRef = useRef('')
     const messageResponsePart = useRef(0);
     const conversationRef = useRef(conversationDetails)
     const selectedConversationRef = useRef<undefined | number>(undefined)
+    const isGeneratingRef = useRef(false)
+    const prevReconnectGenerationRef = useRef(0)
 
-    useEffect(() => {
-        async function GetConversationDetails(){
-            if(selectedConversation){
+    const clearStreamInterrupted = () => setStreamInterrupted(false)
 
-                try{
-                    //setPromptProcessing(true)
-                    selectedConversationRef.current = selectedConversation;
-                    const {data, }: AxiosResponse<ConversationPromptType[]> = await axiosInstance.get(`conversation_details?conversation_id=${selectedConversation}`)
+    async function GetConversationDetails(conversationId: number | undefined) {
+        if (!conversationId) {
+            setConversationDetails([])
+            return
+        }
 
-                    const tempConversations: ConversationPrompt[] = data.map((item) => new ConversationPrompt({
+        try {
+            selectedConversationRef.current = conversationId
+            const { data }: AxiosResponse<ConversationPromptType[]> =
+                await axiosInstance.get(
+                    `conversation_details?conversation_id=${conversationId}`,
+                )
+
+            const tempConversations: ConversationPrompt[] = data.map(
+                (item) =>
+                    new ConversationPrompt({
                         message: item.message,
                         user_created: item.user_created,
-                        created_timestamp: item.created_timestamp
-                    }))
-                    if(tempConversations.length === 1){
-                        // we need to add another card because this is the first message
-                        tempConversations.push(new ConversationPrompt({message: '', user_created:false}))
-                    }
-                    conversationRef.current = tempConversations
-                    setConversationDetails(tempConversations)
+                        created_timestamp: item.created_timestamp,
+                    }),
+            )
+            if (tempConversations.length === 1) {
+                // first message still awaiting assistant reply placeholder
+                tempConversations.push(
+                    new ConversationPrompt({ message: '', user_created: false }),
+                )
+            }
+            conversationRef.current = tempConversations
+            setConversationDetails(tempConversations)
+        } catch (err) {
+            console.log('Failed to refetch conversation after reconnect', err)
+        }
+    }
 
-                }finally{
-                    //setPromptProcessing(false)
+    useEffect(() => {
+        GetConversationDetails(selectedConversation)
+    }, [selectedConversation])
 
-                }
+    // Streaming recovery: on reconnect after a drop mid-stream, refetch conversation
+    // and mark the in-flight assistant turn interrupted (no resume protocol on backend).
+    useEffect(() => {
+        if (reconnectGeneration === prevReconnectGenerationRef.current) {
+            return
+        }
+        prevReconnectGenerationRef.current = reconnectGeneration
+        if (reconnectGeneration === 0) {
+            return
+        }
 
-            }else{
-                setConversationDetails([])
+        const wasGenerating = isGeneratingRef.current || messageResponsePart.current !== 0
+        if (!wasGenerating) {
+            return
+        }
 
+        const partial = messageRef.current
+        messageRef.current = ''
+        messageResponsePart.current = 0
+        isGeneratingRef.current = false
+        setIsGeneratingMessage(false)
+        setStateMessage('')
+        setStreamInterrupted(true)
+
+        const details = [...conversationRef.current]
+        if (details.length > 0) {
+            const last = details[details.length - 1]
+            if (last && !last.user_created && (!last.message || last.message.length === 0)) {
+                details[details.length - 1] = new ConversationPrompt({
+                    message: partial
+                        ? `${partial}\n\n${STREAM_INTERRUPTED_MESSAGE}`
+                        : STREAM_INTERRUPTED_MESSAGE,
+                    user_created: false,
+                })
+                conversationRef.current = details
+                setConversationDetails(details)
             }
         }
-        GetConversationDetails();
 
-    }, [selectedConversation])
+        if (selectedConversationRef.current) {
+            GetConversationDetails(selectedConversationRef.current)
+        }
+    }, [reconnectGeneration])
 
     useEffect(() => {
         /* register a consistent channel name for identifing this chat messages */
@@ -95,17 +155,24 @@ const MessageProvider = ( {children}: MessageProviderProps) => {
                 console.log([...conversationRef.current, new ConversationPrompt({message: `${messageRef.current}`, user_created:false})])
                 messageRef.current = ''
                 setStateMessage('')
+                isGeneratingRef.current = false
                 setIsGeneratingMessage(false)
+                setStreamInterrupted(false)
             }
             else if (message === 'START_OF_THE_STREAM_ENDER_GAME_42'){
                 conversationRef.current = conversationDetails
+                isGeneratingRef.current = true
                 setIsGeneratingMessage(true)
+                setStreamInterrupted(false)
                 messageResponsePart.current = 2
   
             }else if (message === 'CONVERSATION_ID'){
+                isGeneratingRef.current = true
                 setIsGeneratingMessage(true)
+                setStreamInterrupted(false)
                 messageResponsePart.current = 1
             }else{
+                isGeneratingRef.current = true
                 setIsGeneratingMessage(true)
                 if (messageResponsePart.current === 1){
                     // this has to do with the conversation id
@@ -130,7 +197,15 @@ const MessageProvider = ( {children}: MessageProviderProps) => {
 
 
     return(
-        <MessageContext.Provider value={{stateMessage, setStateMessage, conversationDetails, setConversationDetails, isGeneratingMessage}}>
+        <MessageContext.Provider value={{
+            stateMessage,
+            setStateMessage,
+            conversationDetails,
+            setConversationDetails,
+            isGeneratingMessage,
+            streamInterrupted,
+            clearStreamInterrupted,
+        }}>
             {children}
         </MessageContext.Provider>
     )

llm-fe/src/llm-fe/contexts/WebSocketContext.js

@@ -1,7 +1,18 @@
 import { AccountContext } from "./AccountContext";
+import { axiosInstance } from "../../axiosApi";
+import {
+  ConnectionStatus,
+  DEFAULT_RECONNECT_CONFIG,
+  buildHeartbeatPayload,
+  computeBackoffDelay,
+  ensureFreshAccessToken,
+  isSocketOpen,
+  subscribeLifecycleEvents,
+} from "./websocketReconnect";
 
 const {
   useEffect,
+  useCallback,
   createContext,
   useRef,
   useState,
@@ -12,102 +23,389 @@ const WebSocketContext = createContext();
 
 function WebSocketProvider({ children }) {
   const ws = useRef(null);
-  const [socket, setSocket] = useState(null);
-  const channels = useRef({}); // maps each channel to the callback
+  const channels = useRef({});
   const { account } = useContext(AccountContext);
-  const [, setCurrentChannel] = useState("");
-  const [isConnected, setIsConnected] = useState(false);
+  const accountRef = useRef(account);
+  const intentionalCloseRef = useRef(false);
+  const reconnectAttemptRef = useRef(0);
+  const reconnectTimerRef = useRef(null);
+  const heartbeatTimerRef = useRef(null);
+  const connectGenerationRef = useRef(0);
+  const connectingRef = useRef(false);
 
-  /* called from a component that registers a callback for a channel */
-  const subscribe = (channel, callback) => {
-    //console.log(`Subbing to ${channel}`)
-    setCurrentChannel(channel);
+  const [socket, setSocket] = useState(null);
+  const [connectionStatus, setConnectionStatus] = useState(
+    ConnectionStatus.DISCONNECTED,
+  );
+  /** Increments on each successful open after a drop; MessageContext uses for stream recovery. */
+  const [reconnectGeneration, setReconnectGeneration] = useState(0);
+  const hadOpenConnectionRef = useRef(false);
+
+  accountRef.current = account;
+
+  const isConnected = connectionStatus === ConnectionStatus.CONNECTED;
+
+  const clearReconnectTimer = useCallback(() => {
+    if (reconnectTimerRef.current != null) {
+      clearTimeout(reconnectTimerRef.current);
+      reconnectTimerRef.current = null;
+    }
+  }, []);
+
+  const clearHeartbeat = useCallback(() => {
+    if (heartbeatTimerRef.current != null) {
+      clearInterval(heartbeatTimerRef.current);
+      heartbeatTimerRef.current = null;
+    }
+  }, []);
+
+  const startHeartbeat = useCallback(() => {
+    clearHeartbeat();
+    heartbeatTimerRef.current = setInterval(() => {
+      const current = ws.current;
+      if (!isSocketOpen(current)) {
+        return;
+      }
+      try {
+        current.send(buildHeartbeatPayload(accountRef.current?.email));
+      } catch (err) {
+        console.log("WebSocket heartbeat failed", err);
+      }
+    }, DEFAULT_RECONNECT_CONFIG.heartbeatIntervalMs);
+  }, [clearHeartbeat]);
+
+  const subscribe = useCallback((channel, callback) => {
     channels.current[channel] = callback;
-  };
-  /* remove callback  */
-  const unsubscribe = (channel) => {
+  }, []);
+
+  const unsubscribe = useCallback((channel) => {
     delete channels.current[channel];
-  };
+  }, []);
+
+  const sendMessage = useCallback(
+    (message, conversation_id, file, fileType, modelName) => {
+      const current = ws.current;
+      if (!isSocketOpen(current)) {
+        console.log("Error sending message. WebSocket is not open");
+        return false;
+      }
 
-  const sendMessage = (message, conversation_id, file, fileType, modelName) => {
-    if (socket && socket.readyState === WebSocket.OPEN) {
       if (file) {
         const reader = new FileReader();
         reader.onload = () => {
           const base64File = reader.result?.toString().split(",")[1];
-          if (base64File) {
-            const data = {
-              message: message,
-              conversation_id: conversation_id,
-              email: account?.email,
-              file: base64File,
-              fileType: fileType,
-              modelName: modelName,
-            };
-            socket.send(JSON.stringify(data));
+          if (base64File && isSocketOpen(ws.current)) {
+            ws.current.send(
+              JSON.stringify({
+                message,
+                conversation_id,
+                email: accountRef.current?.email,
+                file: base64File,
+                fileType,
+                modelName,
+              }),
+            );
           }
         };
         reader.readAsDataURL(file);
       } else {
-        const data = {
-          message: message,
-          conversation_id: conversation_id,
-          email: account?.email,
-          file: null,
-          fileType: null,
-          modelName: modelName,
+        current.send(
+          JSON.stringify({
+            message,
+            conversation_id,
+            email: accountRef.current?.email,
+            file: null,
+            fileType: null,
+            modelName,
+          }),
+        );
+      }
+      return true;
+    },
+    [],
+  );
+
+  const scheduleReconnect = useCallback(
+    (connectFn) => {
+      if (intentionalCloseRef.current || !accountRef.current) {
+        return;
+      }
+      if (
+        reconnectAttemptRef.current >= DEFAULT_RECONNECT_CONFIG.maxAttempts
+      ) {
+        console.log("WebSocket reconnect attempts exhausted");
+        setConnectionStatus(ConnectionStatus.DISCONNECTED);
+        connectingRef.current = false;
+        return;
+      }
+
+      clearReconnectTimer();
+      setConnectionStatus(ConnectionStatus.RECONNECTING);
+      const attempt = reconnectAttemptRef.current;
+      const delay = computeBackoffDelay(attempt);
+      reconnectAttemptRef.current = attempt + 1;
+
+      reconnectTimerRef.current = setTimeout(() => {
+        reconnectTimerRef.current = null;
+        connectFn({ isRetry: true });
+      }, delay);
+    },
+    [clearReconnectTimer],
+  );
+
+  const connect = useCallback(
+    async ({ isRetry = false } = {}) => {
+      if (!accountRef.current) {
+        return;
+      }
+      if (intentionalCloseRef.current) {
+        return;
+      }
+      if (
+        connectingRef.current ||
+        isSocketOpen(ws.current) ||
+        ws.current?.readyState === WebSocket.CONNECTING
+      ) {
+        return;
+      }
+
+      connectingRef.current = true;
+      clearReconnectTimer();
+      setConnectionStatus(
+        isRetry || hadOpenConnectionRef.current
+          ? ConnectionStatus.RECONNECTING
+          : ConnectionStatus.CONNECTING,
+      );
+
+      // Re-auth before open so REST recovery (#22) and any future WS auth stay valid.
+      await ensureFreshAccessToken({
+        refreshRequest: async (refresh) => {
+          const response = await axiosInstance.post("/token/refresh/", {
+            refresh,
+          });
+          axiosInstance.defaults.headers["Authorization"] =
+            "JWT " + response.data.access;
+          return response.data;
+        },
+      });
+
+      if (!accountRef.current || intentionalCloseRef.current) {
+        connectingRef.current = false;
+        return;
+      }
+
+      const generation = ++connectGenerationRef.current;
+      const url = process.env.REACT_APP_BACKEND_WS_API_BASE_URL;
+      if (!url) {
+        console.log("REACT_APP_BACKEND_WS_API_BASE_URL is not set");
+        connectingRef.current = false;
+        setConnectionStatus(ConnectionStatus.DISCONNECTED);
+        return;
+      }
+
+      try {
+        if (ws.current) {
+          try {
+            ws.current.onopen = null;
+            ws.current.onclose = null;
+            ws.current.onerror = null;
+            ws.current.onmessage = null;
+            if (
+              ws.current.readyState === WebSocket.OPEN ||
+              ws.current.readyState === WebSocket.CONNECTING
+            ) {
+              ws.current.close();
+            }
+          } catch {
+            /* ignore */
+          }
+        }
+
+        const socketInstance = new WebSocket(url);
+        ws.current = socketInstance;
+
+        socketInstance.onopen = () => {
+          if (generation !== connectGenerationRef.current) {
+            return;
+          }
+          connectingRef.current = false;
+          reconnectAttemptRef.current = 0;
+          setSocket(socketInstance);
+          setConnectionStatus(ConnectionStatus.CONNECTED);
+          startHeartbeat();
+          if (hadOpenConnectionRef.current) {
+            setReconnectGeneration((g) => g + 1);
+          }
+          hadOpenConnectionRef.current = true;
         };
 
-        socket.send(JSON.stringify(data));
+        socketInstance.onmessage = (message) => {
+          if (generation !== connectGenerationRef.current) {
+            return;
+          }
+          const data = message.data;
+          if (data === "pong") {
+            return;
+          }
+          try {
+            const parsed = JSON.parse(data);
+            if (parsed && parsed.type === "pong") {
+              return;
+            }
+          } catch {
+            /* not JSON — treat as chat payload */
+          }
+
+          const entries = Object.entries(channels.current);
+          if (entries.length === 0) {
+            console.log("Error: no WebSocket channel subscribers");
+            return;
+          }
+          const [chatChannel] = entries[0];
+          if (channels.current[chatChannel]) {
+            channels.current[chatChannel](data);
+          } else {
+            console.log("Error");
+          }
+        };
+
+        socketInstance.onerror = () => {
+          // onclose follows; reconnect is scheduled there
+        };
+
+        socketInstance.onclose = () => {
+          if (generation !== connectGenerationRef.current) {
+            return;
+          }
+          connectingRef.current = false;
+          clearHeartbeat();
+          setSocket(null);
+          if (intentionalCloseRef.current) {
+            setConnectionStatus(ConnectionStatus.DISCONNECTED);
+            return;
+          }
+          scheduleReconnect(connect);
+        };
+      } catch (err) {
+        console.log("WebSocket connect failed", err);
+        connectingRef.current = false;
+        clearHeartbeat();
+        setSocket(null);
+        scheduleReconnect(connect);
       }
-      //socket.send(`${conversation_id} | ${message}`)
-    } else {
-      console.log("Error sending message. WebSocket is not open");
+    },
+    [clearHeartbeat, clearReconnectTimer, scheduleReconnect, startHeartbeat],
+  );
+
+  const reconnectNow = useCallback(() => {
+    if (!accountRef.current || intentionalCloseRef.current) {
+      return;
     }
-  };
+    if (isSocketOpen(ws.current)) {
+      return;
+    }
+    reconnectAttemptRef.current = 0;
+    clearReconnectTimer();
+    connect({ isRetry: hadOpenConnectionRef.current });
+  }, [clearReconnectTimer, connect]);
+
   useEffect(() => {
-    /* WS initialization and cleanup */
-    if (account) {
-      ws.current = new WebSocket(process.env.REACT_APP_BACKEND_WS_API_BASE_URL);
+    intentionalCloseRef.current = false;
+    reconnectAttemptRef.current = 0;
+    hadOpenConnectionRef.current = false;
+    setReconnectGeneration(0);
 
-      ws.current.onopen = () => {
-        setSocket(ws.current);
-        setIsConnected(true);
-      };
-      ws.current.onclose = () => {
-        console.log('websocket closed');
-        setIsConnected(false);
-      };
-      ws.current.onmessage = (message) => {
-        const data = message.data;
-        // lookup for an existing chat in which this message belongs
-        // if no chat is subscribed send message to generic channel
-        const chatChannel = Object.entries(channels.current)[0][0];
-
-        if (channels.current[chatChannel]) {
-          /* in chat component the subscribed channel is `MESSAGE_CREATE_${id}` */
-          channels.current[chatChannel](data);
-        } else {
-          /* in notifications wrapper the subscribed channel is `MESSAGE_CREATE` */
-          console.log("Error");
-          // channels.current[type]?.(data)
+    if (!account) {
+      intentionalCloseRef.current = true;
+      clearReconnectTimer();
+      clearHeartbeat();
+      connectGenerationRef.current += 1;
+      if (ws.current) {
+        try {
+          ws.current.close();
+        } catch {
+          /* ignore */
         }
-      };
-      return () => {
-        ws.current.close();
-      };
+        ws.current = null;
+      }
+      setSocket(null);
+      setConnectionStatus(ConnectionStatus.DISCONNECTED);
+      return undefined;
     }
-  }, [account]);
 
-  /* WS provider dom */
-  /* subscribe and unsubscribe are the only required prop for the context */
+    connect({ isRetry: false });
+
+    const unsubscribeLifecycle = subscribeLifecycleEvents({
+      onResume: () => {
+        if (!isSocketOpen(ws.current)) {
+          reconnectNow();
+        }
+      },
+      onOnline: () => {
+        if (!isSocketOpen(ws.current)) {
+          reconnectNow();
+        }
+      },
+      onOffline: () => {
+        clearReconnectTimer();
+        setConnectionStatus((status) =>
+          status === ConnectionStatus.CONNECTED
+            ? ConnectionStatus.DISCONNECTED
+            : status,
+        );
+      },
+    });
+
+    return () => {
+      intentionalCloseRef.current = true;
+      clearReconnectTimer();
+      clearHeartbeat();
+      unsubscribeLifecycle();
+      connectGenerationRef.current += 1;
+      if (ws.current) {
+        try {
+          ws.current.onopen = null;
+          ws.current.onclose = null;
+          ws.current.onerror = null;
+          ws.current.onmessage = null;
+          ws.current.close();
+        } catch {
+          /* ignore */
+        }
+        ws.current = null;
+      }
+      connectingRef.current = false;
+      setSocket(null);
+      setConnectionStatus(ConnectionStatus.DISCONNECTED);
+    };
+  }, [
+    account,
+    clearHeartbeat,
+    clearReconnectTimer,
+    connect,
+    reconnectNow,
+  ]);
+
   return (
     <WebSocketContext.Provider
-      value={[subscribe, unsubscribe, socket, sendMessage, isConnected]}
+      value={[
+        subscribe,
+        unsubscribe,
+        socket,
+        sendMessage,
+        isConnected,
+        connectionStatus,
+        reconnectGeneration,
+        reconnectNow,
+      ]}
     >
       {children}
     </WebSocketContext.Provider>
   );
 }
 
-export { WebSocketContext, WebSocketProvider };
+export {
+  WebSocketContext,
+  WebSocketProvider,
+  ConnectionStatus,
+};

llm-fe/src/llm-fe/contexts/WebSocketContext.test.js

@@ -0,0 +1,242 @@
+import React, { useContext } from 'react';
+import { act, render, screen, waitFor } from '@testing-library/react';
+import { AccountContext } from './AccountContext';
+import {
+  ConnectionStatus,
+  WebSocketContext,
+  WebSocketProvider,
+} from './WebSocketContext';
+
+jest.mock('../../axiosApi', () => ({
+  axiosInstance: {
+    post: jest.fn(),
+    defaults: { headers: {} },
+  },
+}));
+
+class MockWebSocket {
+  static OPEN = 1;
+  static CONNECTING = 0;
+  static CLOSING = 2;
+  static CLOSED = 3;
+  static instances = [];
+
+  constructor(url) {
+    this.url = url;
+    this.readyState = MockWebSocket.CONNECTING;
+    this.sent = [];
+    this.onopen = null;
+    this.onclose = null;
+    this.onerror = null;
+    this.onmessage = null;
+    MockWebSocket.instances.push(this);
+  }
+
+  send(data) {
+    this.sent.push(data);
+  }
+
+  close() {
+    this.readyState = MockWebSocket.CLOSED;
+    if (this.onclose) {
+      this.onclose({ code: 1000 });
+    }
+  }
+
+  /** Test helper: simulate successful open. */
+  open() {
+    this.readyState = MockWebSocket.OPEN;
+    if (this.onopen) {
+      this.onopen({});
+    }
+  }
+
+  static reset() {
+    MockWebSocket.instances = [];
+  }
+}
+
+const Probe = () => {
+  const [
+    subscribe,
+    ,
+    ,
+    sendMessage,
+    isConnected,
+    connectionStatus,
+    reconnectGeneration,
+  ] = useContext(WebSocketContext);
+
+  React.useEffect(() => {
+    subscribe('ACCOUNT_ID_test@example.com', () => {});
+  }, [subscribe]);
+
+  return (
+    <div>
+      <span data-testid="status">{connectionStatus}</span>
+      <span data-testid="connected">{isConnected ? 'yes' : 'no'}</span>
+      <span data-testid="reconnect-gen">{reconnectGeneration}</span>
+      <button
+        type="button"
+        onClick={() => sendMessage('hi', 1, null, null, 'FAST')}
+      >
+        send
+      </button>
+    </div>
+  );
+};
+
+const renderWithAccount = (account) =>
+  render(
+    <AccountContext.Provider value={{ account, setAccount: () => {} }}>
+      <WebSocketProvider>
+        <Probe />
+      </WebSocketProvider>
+    </AccountContext.Provider>,
+  );
+
+describe('WebSocketProvider', () => {
+  const originalWebSocket = global.WebSocket;
+  const originalEnv = process.env.REACT_APP_BACKEND_WS_API_BASE_URL;
+
+  beforeEach(() => {
+    MockWebSocket.reset();
+    global.WebSocket = MockWebSocket;
+    process.env.REACT_APP_BACKEND_WS_API_BASE_URL = 'ws://test/ws/chat_again/';
+    localStorage.clear();
+    jest.useFakeTimers();
+  });
+
+  afterEach(() => {
+    jest.useRealTimers();
+    global.WebSocket = originalWebSocket;
+    process.env.REACT_APP_BACKEND_WS_API_BASE_URL = originalEnv;
+    delete window.Capacitor;
+  });
+
+  it('stays disconnected when no account', async () => {
+    renderWithAccount(undefined);
+    expect(screen.getByTestId('status')).toHaveTextContent(
+      ConnectionStatus.DISCONNECTED,
+    );
+    expect(MockWebSocket.instances).toHaveLength(0);
+  });
+
+  it('connects when account is present', async () => {
+    renderWithAccount({ email: 'test@example.com' });
+
+    await waitFor(() => {
+      expect(MockWebSocket.instances.length).toBeGreaterThan(0);
+    });
+
+    await act(async () => {
+      MockWebSocket.instances[0].open();
+    });
+
+    expect(screen.getByTestId('status')).toHaveTextContent(
+      ConnectionStatus.CONNECTED,
+    );
+    expect(screen.getByTestId('connected')).toHaveTextContent('yes');
+  });
+
+  it('schedules reconnect with backoff after unexpected close', async () => {
+    renderWithAccount({ email: 'test@example.com' });
+
+    await waitFor(() => {
+      expect(MockWebSocket.instances.length).toBe(1);
+    });
+
+    await act(async () => {
+      MockWebSocket.instances[0].open();
+    });
+    expect(screen.getByTestId('status')).toHaveTextContent(
+      ConnectionStatus.CONNECTED,
+    );
+
+    await act(async () => {
+      MockWebSocket.instances[0].close();
+    });
+
+    expect(screen.getByTestId('status')).toHaveTextContent(
+      ConnectionStatus.RECONNECTING,
+    );
+
+    await act(async () => {
+      jest.advanceTimersByTime(35000);
+    });
+
+    await waitFor(() => {
+      expect(MockWebSocket.instances.length).toBeGreaterThan(1);
+    });
+  });
+
+  it('increments reconnectGeneration after drop and successful reopen', async () => {
+    renderWithAccount({ email: 'test@example.com' });
+
+    await waitFor(() => expect(MockWebSocket.instances.length).toBe(1));
+    await act(async () => {
+      MockWebSocket.instances[0].open();
+    });
+    expect(screen.getByTestId('reconnect-gen')).toHaveTextContent('0');
+
+    await act(async () => {
+      MockWebSocket.instances[0].close();
+    });
+    await act(async () => {
+      jest.advanceTimersByTime(35000);
+    });
+
+    await waitFor(() => expect(MockWebSocket.instances.length).toBeGreaterThan(1));
+    await act(async () => {
+      MockWebSocket.instances[MockWebSocket.instances.length - 1].open();
+    });
+
+    await waitFor(() => {
+      expect(screen.getByTestId('reconnect-gen')).toHaveTextContent('1');
+    });
+  });
+
+  it('reconnects on visibility resume when socket is down', async () => {
+    renderWithAccount({ email: 'test@example.com' });
+    await waitFor(() => expect(MockWebSocket.instances.length).toBe(1));
+    await act(async () => {
+      MockWebSocket.instances[0].open();
+    });
+
+    // Simulate dead socket without going through close handler scheduling
+    MockWebSocket.instances[0].readyState = MockWebSocket.CLOSED;
+    MockWebSocket.instances[0].onclose = null;
+
+    Object.defineProperty(document, 'visibilityState', {
+      configurable: true,
+      get: () => 'visible',
+    });
+
+    const before = MockWebSocket.instances.length;
+    await act(async () => {
+      document.dispatchEvent(new Event('visibilitychange'));
+    });
+
+    await waitFor(() => {
+      expect(MockWebSocket.instances.length).toBeGreaterThan(before);
+    });
+  });
+
+  it('sends chat payload when connected', async () => {
+    renderWithAccount({ email: 'test@example.com' });
+    await waitFor(() => expect(MockWebSocket.instances.length).toBe(1));
+    await act(async () => {
+      MockWebSocket.instances[0].open();
+    });
+
+    await act(async () => {
+      screen.getByRole('button', { name: 'send' }).click();
+    });
+
+    const sent = MockWebSocket.instances[0].sent;
+    expect(sent.length).toBeGreaterThan(0);
+    const payload = JSON.parse(sent[sent.length - 1]);
+    expect(payload.message).toBe('hi');
+    expect(payload.email).toBe('test@example.com');
+  });
+});

llm-fe/src/llm-fe/contexts/websocketReconnect.js

@@ -0,0 +1,243 @@
+/** Connection status values exposed by WebSocketProvider. */
+export const ConnectionStatus = Object.freeze({
+  CONNECTED: "connected",
+  CONNECTING: "connecting",
+  RECONNECTING: "reconnecting",
+  DISCONNECTED: "disconnected",
+});
+
+export const DEFAULT_RECONNECT_CONFIG = Object.freeze({
+  /** Base delay before first retry (ms). */
+  baseDelayMs: 1000,
+  /** Max delay between retries (ms). */
+  maxDelayMs: 30000,
+  /** Hard cap on reconnect attempts before staying disconnected. */
+  maxAttempts: 12,
+  /** Refresh access token if it expires within this window (seconds). */
+  tokenRefreshSkewSeconds: 60,
+  /** App-level keepalive interval (ms). Stay under common proxy idle timeouts (~60s). */
+  heartbeatIntervalMs: 25000,
+});
+
+/**
+ * Exponential backoff with full jitter.
+ * delay = random_between(0, min(maxDelay, base * 2^attempt))
+ *
+ * @param {number} attempt zero-based attempt index
+ * @param {{ baseDelayMs?: number, maxDelayMs?: number }} [config]
+ * @param {() => number} [random] injectable RNG for tests
+ * @returns {number} delay in milliseconds
+ */
+export function computeBackoffDelay(
+  attempt,
+  config = {},
+  random = Math.random,
+) {
+  const baseDelayMs = config.baseDelayMs ?? DEFAULT_RECONNECT_CONFIG.baseDelayMs;
+  const maxDelayMs = config.maxDelayMs ?? DEFAULT_RECONNECT_CONFIG.maxDelayMs;
+  const exp = Math.min(maxDelayMs, baseDelayMs * 2 ** Math.max(0, attempt));
+  return Math.floor(random() * exp);
+}
+
+/**
+ * @param {number|undefined|null} expSeconds JWT exp claim (seconds since epoch)
+ * @param {number} nowSeconds current unix time in seconds
+ * @param {number} [skewSeconds]
+ * @returns {boolean}
+ */
+export function shouldRefreshAccessToken(
+  expSeconds,
+  nowSeconds,
+  skewSeconds = DEFAULT_RECONNECT_CONFIG.tokenRefreshSkewSeconds,
+) {
+  if (expSeconds == null || Number.isNaN(expSeconds)) {
+    return false;
+  }
+  return expSeconds - nowSeconds <= skewSeconds;
+}
+
+/**
+ * Decode JWT payload without verifying signature (client-side expiry check only).
+ * @param {string|null|undefined} token
+ * @returns {{ exp?: number } | null}
+ */
+export function decodeJwtPayload(token) {
+  if (!token || typeof token !== "string") {
+    return null;
+  }
+  const parts = token.split(".");
+  if (parts.length < 2) {
+    return null;
+  }
+  try {
+    const normalized = parts[1].replace(/-/g, "+").replace(/_/g, "/");
+    const json = atob(normalized);
+    return JSON.parse(json);
+  } catch {
+    return null;
+  }
+}
+
+/**
+ * Refresh access token via REST if expired/near-expiry.
+ * Coordinates with JWT auth used by axios (#22); WS itself still uses account email.
+ *
+ * @param {{
+ *   getAccessToken?: () => string|null,
+ *   getRefreshToken?: () => string|null,
+ *   setTokens?: (access: string, refresh: string) => void,
+ *   refreshRequest?: (refresh: string) => Promise<{ access: string, refresh?: string }>,
+ *   nowSeconds?: () => number,
+ *   skewSeconds?: number,
+ * }} [deps]
+ * @returns {Promise<string|null>} access token to use, or null if unavailable
+ */
+export async function ensureFreshAccessToken(deps = {}) {
+  const getAccessToken =
+    deps.getAccessToken ?? (() => localStorage.getItem("access_token"));
+  const getRefreshToken =
+    deps.getRefreshToken ?? (() => localStorage.getItem("refresh_token"));
+  const setTokens =
+    deps.setTokens ??
+    ((access, refresh) => {
+      localStorage.setItem("access_token", access);
+      if (refresh) {
+        localStorage.setItem("refresh_token", refresh);
+      }
+    });
+  const nowSeconds = deps.nowSeconds ?? (() => Math.ceil(Date.now() / 1000));
+  const skewSeconds =
+    deps.skewSeconds ?? DEFAULT_RECONNECT_CONFIG.tokenRefreshSkewSeconds;
+
+  const access = getAccessToken();
+  const refresh = getRefreshToken();
+  if (!access) {
+    return null;
+  }
+
+  const payload = decodeJwtPayload(access);
+  if (
+    !shouldRefreshAccessToken(payload?.exp, nowSeconds(), skewSeconds) ||
+    !refresh
+  ) {
+    return access;
+  }
+
+  const refreshPayload = decodeJwtPayload(refresh);
+  if (
+    refreshPayload?.exp != null &&
+    refreshPayload.exp <= nowSeconds()
+  ) {
+    return null;
+  }
+
+  if (!deps.refreshRequest) {
+    return access;
+  }
+
+  try {
+    const data = await deps.refreshRequest(refresh);
+    if (data?.access) {
+      setTokens(data.access, data.refresh ?? refresh);
+      return data.access;
+    }
+  } catch (err) {
+    console.log("Token refresh before WS connect failed", err);
+  }
+  return access;
+}
+
+/**
+ * Subscribe to browser + Capacitor lifecycle events that imply the socket may be dead.
+ * Capacitor plugins are resolved via window.Capacitor.Plugins when present (#20/#21);
+ * web visibility/online listeners always register.
+ *
+ * @param {{
+ *   onResume: () => void,
+ *   onOnline: () => void,
+ *   onOffline?: () => void,
+ * }} handlers
+ * @returns {() => void} cleanup
+ */
+export function subscribeLifecycleEvents(handlers) {
+  const { onResume, onOnline, onOffline } = handlers;
+  const cleanups = [];
+
+  const onVisibility = () => {
+    if (document.visibilityState === "visible") {
+      onResume();
+    }
+  };
+  const onWindowOnline = () => onOnline();
+  const onWindowOffline = () => onOffline?.();
+
+  document.addEventListener("visibilitychange", onVisibility);
+  window.addEventListener("online", onWindowOnline);
+  window.addEventListener("offline", onWindowOffline);
+  cleanups.push(() => {
+    document.removeEventListener("visibilitychange", onVisibility);
+    window.removeEventListener("online", onWindowOnline);
+    window.removeEventListener("offline", onWindowOffline);
+  });
+
+  const Cap = typeof window !== "undefined" ? window.Capacitor : undefined;
+  const plugins = Cap?.Plugins;
+  if (plugins?.App?.addListener) {
+    const handle = plugins.App.addListener("appStateChange", (state) => {
+      if (state?.isActive) {
+        onResume();
+      }
+    });
+    cleanups.push(() => {
+      if (handle?.remove) {
+        handle.remove();
+      } else if (typeof handle?.then === "function") {
+        handle.then((h) => h?.remove?.());
+      }
+    });
+  }
+  if (plugins?.Network?.addListener) {
+    const handle = plugins.Network.addListener(
+      "networkStatusChange",
+      (status) => {
+        if (status?.connected) {
+          onOnline();
+        } else {
+          onOffline?.();
+        }
+      },
+    );
+    cleanups.push(() => {
+      if (handle?.remove) {
+        handle.remove();
+      } else if (typeof handle?.then === "function") {
+        handle.then((h) => h?.remove?.());
+      }
+    });
+  }
+
+  return () => {
+    cleanups.forEach((fn) => {
+      try {
+        fn();
+      } catch {
+        /* ignore plugin cleanup errors */
+      }
+    });
+  };
+}
+
+/**
+ * Heartbeat payload ignored by chat handlers that key on `message`.
+ * Keeps proxies from idling out `/ws/` connections.
+ */
+export function buildHeartbeatPayload(email) {
+  return JSON.stringify({
+    type: "ping",
+    email: email ?? null,
+  });
+}
+
+export function isSocketOpen(socket) {
+  return Boolean(socket && socket.readyState === WebSocket.OPEN);
+}

llm-fe/src/llm-fe/contexts/websocketReconnect.test.js

@@ -0,0 +1,191 @@
+import {
+  ConnectionStatus,
+  DEFAULT_RECONNECT_CONFIG,
+  buildHeartbeatPayload,
+  computeBackoffDelay,
+  decodeJwtPayload,
+  ensureFreshAccessToken,
+  shouldRefreshAccessToken,
+  subscribeLifecycleEvents,
+} from './websocketReconnect';
+
+describe('websocketReconnect helpers', () => {
+  describe('computeBackoffDelay', () => {
+    it('stays within exponential cap for attempt 0', () => {
+      const delay = computeBackoffDelay(0, { baseDelayMs: 1000, maxDelayMs: 30000 }, () => 1);
+      expect(delay).toBe(1000);
+    });
+
+    it('doubles upper bound each attempt until maxDelay', () => {
+      const delay = computeBackoffDelay(3, { baseDelayMs: 1000, maxDelayMs: 30000 }, () => 1);
+      expect(delay).toBe(8000);
+    });
+
+    it('never exceeds maxDelayMs', () => {
+      const delay = computeBackoffDelay(20, { baseDelayMs: 1000, maxDelayMs: 5000 }, () => 1);
+      expect(delay).toBe(5000);
+    });
+
+    it('applies jitter via random()', () => {
+      const delay = computeBackoffDelay(2, { baseDelayMs: 1000, maxDelayMs: 30000 }, () => 0.5);
+      expect(delay).toBe(2000);
+    });
+  });
+
+  describe('shouldRefreshAccessToken', () => {
+    it('returns true when token already expired', () => {
+      expect(shouldRefreshAccessToken(100, 200, 60)).toBe(true);
+    });
+
+    it('returns true within skew window', () => {
+      expect(shouldRefreshAccessToken(250, 200, 60)).toBe(true);
+    });
+
+    it('returns false when plenty of lifetime remains', () => {
+      expect(shouldRefreshAccessToken(500, 200, 60)).toBe(false);
+    });
+
+    it('returns false for missing exp', () => {
+      expect(shouldRefreshAccessToken(undefined, 200, 60)).toBe(false);
+    });
+  });
+
+  describe('decodeJwtPayload', () => {
+    it('decodes a minimal JWT payload', () => {
+      // {"exp":1999999999}
+      const token = `x.${btoa(JSON.stringify({ exp: 1999999999 }))}.y`;
+      expect(decodeJwtPayload(token)).toEqual({ exp: 1999999999 });
+    });
+
+    it('returns null for garbage', () => {
+      expect(decodeJwtPayload('not-a-jwt')).toBeNull();
+      expect(decodeJwtPayload(null)).toBeNull();
+    });
+  });
+
+  describe('ensureFreshAccessToken', () => {
+    it('returns existing access token when still fresh', async () => {
+      const refreshRequest = jest.fn();
+      const token = await ensureFreshAccessToken({
+        getAccessToken: () =>
+          `h.${btoa(JSON.stringify({ exp: 9999999999 }))}.s`,
+        getRefreshToken: () =>
+          `h.${btoa(JSON.stringify({ exp: 9999999999 }))}.s`,
+        refreshRequest,
+        nowSeconds: () => 1000,
+        skewSeconds: 60,
+      });
+      expect(token).toContain('h.');
+      expect(refreshRequest).not.toHaveBeenCalled();
+    });
+
+    it('refreshes when access token is near expiry', async () => {
+      const setTokens = jest.fn();
+      const refreshRequest = jest.fn().mockResolvedValue({
+        access: 'new-access',
+        refresh: 'new-refresh',
+      });
+      const token = await ensureFreshAccessToken({
+        getAccessToken: () =>
+          `h.${btoa(JSON.stringify({ exp: 1050 }))}.s`,
+        getRefreshToken: () =>
+          `h.${btoa(JSON.stringify({ exp: 9999999999 }))}.s`,
+        setTokens,
+        refreshRequest,
+        nowSeconds: () => 1000,
+        skewSeconds: 60,
+      });
+      expect(refreshRequest).toHaveBeenCalled();
+      expect(setTokens).toHaveBeenCalledWith('new-access', 'new-refresh');
+      expect(token).toBe('new-access');
+    });
+
+    it('returns null when refresh token is expired', async () => {
+      const refreshRequest = jest.fn();
+      const token = await ensureFreshAccessToken({
+        getAccessToken: () =>
+          `h.${btoa(JSON.stringify({ exp: 1000 }))}.s`,
+        getRefreshToken: () =>
+          `h.${btoa(JSON.stringify({ exp: 900 }))}.s`,
+        refreshRequest,
+        nowSeconds: () => 1000,
+        skewSeconds: 60,
+      });
+      expect(token).toBeNull();
+      expect(refreshRequest).not.toHaveBeenCalled();
+    });
+  });
+
+  describe('buildHeartbeatPayload', () => {
+    it('includes ping type and email', () => {
+      expect(JSON.parse(buildHeartbeatPayload('a@b.com'))).toEqual({
+        type: 'ping',
+        email: 'a@b.com',
+      });
+    });
+  });
+
+  describe('ConnectionStatus / defaults', () => {
+    it('exposes expected status strings', () => {
+      expect(ConnectionStatus.CONNECTED).toBe('connected');
+      expect(ConnectionStatus.RECONNECTING).toBe('reconnecting');
+      expect(DEFAULT_RECONNECT_CONFIG.maxAttempts).toBeGreaterThan(0);
+    });
+  });
+
+  describe('subscribeLifecycleEvents', () => {
+    it('fires onResume when document becomes visible', () => {
+      const onResume = jest.fn();
+      const onOnline = jest.fn();
+      const cleanup = subscribeLifecycleEvents({ onResume, onOnline });
+
+      Object.defineProperty(document, 'visibilityState', {
+        configurable: true,
+        get: () => 'visible',
+      });
+      document.dispatchEvent(new Event('visibilitychange'));
+      expect(onResume).toHaveBeenCalled();
+
+      cleanup();
+    });
+
+    it('fires onOnline for window online event', () => {
+      const onResume = jest.fn();
+      const onOnline = jest.fn();
+      const cleanup = subscribeLifecycleEvents({ onResume, onOnline });
+
+      window.dispatchEvent(new Event('online'));
+      expect(onOnline).toHaveBeenCalled();
+
+      cleanup();
+    });
+
+    it('hooks Capacitor App/Network plugins when present', () => {
+      const appListener = jest.fn(() => ({ remove: jest.fn() }));
+      const networkListener = jest.fn(() => ({ remove: jest.fn() }));
+      window.Capacitor = {
+        Plugins: {
+          App: { addListener: appListener },
+          Network: { addListener: networkListener },
+        },
+      };
+
+      const cleanup = subscribeLifecycleEvents({
+        onResume: jest.fn(),
+        onOnline: jest.fn(),
+      });
+
+      expect(appListener).toHaveBeenCalledWith(
+        'appStateChange',
+        expect.any(Function),
+      );
+      expect(networkListener).toHaveBeenCalledWith(
+        'networkStatusChange',
+        expect.any(Function),
+      );
+
+      cleanup();
+      delete window.Capacitor;
+    });
+  });
+});

llm-fe/src/llm-fe/pages/AsyncDashboard2/AsyncDashboard2.tsx

@@ -11,7 +11,7 @@ import {
 } from "../../data";
 import { ConversationContext } from "../../contexts/ConversationContext";
 import ConversationDetailCard from "../../components/ConversationDetailCard/ConversationDetailCard";
-import { WebSocketContext } from "../../contexts/WebSocketContext";
+import { WebSocketContext, ConnectionStatus } from "../../contexts/WebSocketContext";
 import { MessageContext } from "../../contexts/MessageContext";
 import ParticleBackground from "../../components/ParticleBackground/ParticleBackground";
 
@@ -290,7 +290,7 @@ const AlwaysScrollToBottom = (): JSX.Element => {
 };
 
 const AsyncDashboardInner = (): JSX.Element => {
-  const [, , , sendMessage, isConnected] =
+  const [, , , sendMessage, isConnected, connectionStatus] =
     useContext(WebSocketContext);
 
   const { conversations, selectedConversation, setSelectedConversation, deleteConversation } =
@@ -300,6 +300,7 @@ const AsyncDashboardInner = (): JSX.Element => {
     conversationDetails,
     setConversationDetails,
     stateMessage,
+    streamInterrupted,
   } = useContext(MessageContext);
 
   const conversationRef = useRef(conversationDetails);
@@ -307,6 +308,19 @@ const AsyncDashboardInner = (): JSX.Element => {
   const textareaRef = useRef<HTMLTextAreaElement>(null);
   const [isSidebarOpen, setIsSidebarOpen] = useState(false);
 
+  const connectionBanner = (() => {
+    if (connectionStatus === ConnectionStatus.CONNECTED || isConnected) {
+      return null;
+    }
+    if (connectionStatus === ConnectionStatus.CONNECTING) {
+      return { text: "Connecting to chat…", color: "#f0a500" };
+    }
+    if (connectionStatus === ConnectionStatus.RECONNECTING) {
+      return { text: "Reconnecting…", color: "#f0a500" };
+    }
+    return { text: "Connection to the LLM is not active", color: "#ff4444" };
+  })();
+
   const handlePromptSubmit = async (
     { prompt, file, fileType, modelName }: PromptValues,
     { resetForm }: any,
@@ -440,15 +454,26 @@ const AsyncDashboardInner = (): JSX.Element => {
           >
             {(formik) => (
               <Form>
-                {!isConnected && (
+                {connectionBanner && (
                   <div style={{
-                    color: '#ff4444',
+                    color: connectionBanner.color,
                     textAlign: 'center',
                     marginBottom: '0.5rem',
                     fontSize: '0.9rem',
                     fontWeight: 500
                   }}>
-                    Connection to the LLM is not active
+                    {connectionBanner.text}
+                  </div>
+                )}
+                {streamInterrupted && isConnected && (
+                  <div style={{
+                    color: '#f0a500',
+                    textAlign: 'center',
+                    marginBottom: '0.5rem',
+                    fontSize: '0.85rem',
+                    fontWeight: 500
+                  }}>
+                    Last response was interrupted. You can re-send your message.
                   </div>
                 )}
                 <StyledInputContainer>