From 22f7ea6aaa181e71cbe823c07a3e5746878b431c Mon Sep 17 00:00:00 2001 From: Ryan Westfall Date: Sun, 26 Jul 2026 07:01:53 -0500 Subject: [PATCH] Harden chat WebSocket for mobile lifecycle (#23) Reconnect with backoff/jitter, lifecycle + network resume hooks, JWT refresh before reconnect, heartbeat, and visible connection/stream-interrupted UI so backgrounding or network switches recover without a reload. --- llm-fe/src/llm-fe/contexts/MessageContext.tsx | 137 ++++-- .../src/llm-fe/contexts/WebSocketContext.js | 436 +++++++++++++++--- .../llm-fe/contexts/WebSocketContext.test.js | 242 ++++++++++ .../src/llm-fe/contexts/websocketReconnect.js | 243 ++++++++++ .../contexts/websocketReconnect.test.js | 191 ++++++++ .../pages/AsyncDashboard2/AsyncDashboard2.tsx | 35 +- 6 files changed, 1179 insertions(+), 105 deletions(-) create mode 100644 llm-fe/src/llm-fe/contexts/WebSocketContext.test.js create mode 100644 llm-fe/src/llm-fe/contexts/websocketReconnect.js create mode 100644 llm-fe/src/llm-fe/contexts/websocketReconnect.test.js diff --git a/llm-fe/src/llm-fe/contexts/MessageContext.tsx b/llm-fe/src/llm-fe/contexts/MessageContext.tsx index 38b7463..7ae4a44 100644 --- a/llm-fe/src/llm-fe/contexts/MessageContext.tsx +++ b/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(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('') const [conversationDetails, setConversationDetails] = useState([]) const [isGeneratingMessage, setIsGeneratingMessage] = useState(false) + const [streamInterrupted, setStreamInterrupted] = useState(false) const messageRef = useRef('') const messageResponsePart = useRef(0); const conversationRef = useRef(conversationDetails) const selectedConversationRef = useRef(undefined) + const isGeneratingRef = useRef(false) + const prevReconnectGenerationRef = useRef(0) - useEffect(() => { - async function GetConversationDetails(){ - if(selectedConversation){ - - try{ - //setPromptProcessing(true) - selectedConversationRef.current = selectedConversation; - const {data, }: AxiosResponse = await axiosInstance.get(`conversation_details?conversation_id=${selectedConversation}`) - - const tempConversations: ConversationPrompt[] = data.map((item) => new ConversationPrompt({ + const clearStreamInterrupted = () => setStreamInterrupted(false) + + async function GetConversationDetails(conversationId: number | undefined) { + if (!conversationId) { + setConversationDetails([]) + return + } + + try { + selectedConversationRef.current = conversationId + const { data }: AxiosResponse = + 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) - - }finally{ - //setPromptProcessing(false) - - } - - }else{ - setConversationDetails([]) + 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) + } + } + 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 + } + + 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( - + {children} ) diff --git a/llm-fe/src/llm-fe/contexts/WebSocketContext.js b/llm-fe/src/llm-fe/contexts/WebSocketContext.js index c5f8a0b..6eeb84c 100644 --- a/llm-fe/src/llm-fe/contexts/WebSocketContext.js +++ b/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 ( {children} ); } -export { WebSocketContext, WebSocketProvider }; +export { + WebSocketContext, + WebSocketProvider, + ConnectionStatus, +}; diff --git a/llm-fe/src/llm-fe/contexts/WebSocketContext.test.js b/llm-fe/src/llm-fe/contexts/WebSocketContext.test.js new file mode 100644 index 0000000..fc7f3d2 --- /dev/null +++ b/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 ( +
+ {connectionStatus} + {isConnected ? 'yes' : 'no'} + {reconnectGeneration} + +
+ ); +}; + +const renderWithAccount = (account) => + render( + {} }}> + + + + , + ); + +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'); + }); +}); diff --git a/llm-fe/src/llm-fe/contexts/websocketReconnect.js b/llm-fe/src/llm-fe/contexts/websocketReconnect.js new file mode 100644 index 0000000..045b572 --- /dev/null +++ b/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} 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); +} diff --git a/llm-fe/src/llm-fe/contexts/websocketReconnect.test.js b/llm-fe/src/llm-fe/contexts/websocketReconnect.test.js new file mode 100644 index 0000000..2a6105e --- /dev/null +++ b/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; + }); + }); +}); diff --git a/llm-fe/src/llm-fe/pages/AsyncDashboard2/AsyncDashboard2.tsx b/llm-fe/src/llm-fe/pages/AsyncDashboard2/AsyncDashboard2.tsx index 5e61d97..b40b414 100644 --- a/llm-fe/src/llm-fe/pages/AsyncDashboard2/AsyncDashboard2.tsx +++ b/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(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) => (
- {!isConnected && ( + {connectionBanner && (
- Connection to the LLM is not active + {connectionBanner.text} +
+ )} + {streamInterrupted && isConnected && ( +
+ Last response was interrupted. You can re-send your message.
)}