From 0d8a0af2b5246b4b2acac2e8a89005cb0931a1bd Mon Sep 17 00:00:00 2001 From: Kyle Carberry Date: Thu, 5 Mar 2026 09:17:41 -0500 Subject: [PATCH] fix(site): add WebSocket reconnection with exponential backoff to chat stream (#22662) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Problem During rolling deploys, the chat stream WebSocket disconnects and the user sees **"Chat stream disconnected."** permanently with no recovery other than a full page refresh. ## Changes Add automatic WebSocket reconnection with capped exponential backoff (1s → 2s → 4s → … → 10s max) to the chat stream `useEffect` in `ChatContext.ts`. **`ChatContext.ts`** - Wrap socket creation in a `connect()` function that can be retried on disconnect. - On `error` or `close`, schedule a reconnect with exponential backoff via `scheduleReconnect()`. - Guard against double scheduling when both `error` and `close` fire for the same disconnect (`disconnected` flag per connection). - On successful reconnect (`open` event), reset backoff and clear `streamError`. - Pass the latest `lastMessageIdRef` on each reconnect so the server replays only unseen durable messages via `after_id`. - Add `RECONNECT_BASE_MS` (1s) and `RECONNECT_MAX_MS` (10s) constants. **`ChatContext.test.tsx`** - Extend mock socket with `emitOpen()` and `emitClose()` helpers. - Update existing disconnect tests for new reconnect behavior. - Add **"sets streamError on WebSocket disconnect and reconnects"** — verifies error banner appears, reconnect fires, and error clears on open. - Add **"uses exponential backoff on consecutive disconnects"** — verifies increasing delays between reconnects. - Add **"passes latest message ID on reconnect for catch-up"** — verifies `after_id` is forwarded on each reconnect. ### User experience - **Before**: "Chat stream disconnected." — permanent, requires page refresh - **After**: "Chat stream disconnected. Reconnecting…" — auto-recovers within seconds ### Limitations `message_part` events (streaming LLM tokens) are ephemeral / in-memory only. If the processing replica dies mid-step, those partial tokens are lost regardless of reconnect. The backend's stale-chat recovery will re-run the step on a new replica; the reconnect ensures the client is connected to see that happen. --- .../AgentDetail/ChatContext.test.tsx | 191 ++++++++++++++++-- .../AgentsPage/AgentDetail/ChatContext.ts | 77 ++++++- 2 files changed, 244 insertions(+), 24 deletions(-) diff --git a/site/src/pages/AgentsPage/AgentDetail/ChatContext.test.tsx b/site/src/pages/AgentsPage/AgentDetail/ChatContext.test.tsx index 28306a7545..77aceee2c9 100644 --- a/site/src/pages/AgentsPage/AgentDetail/ChatContext.test.tsx +++ b/site/src/pages/AgentsPage/AgentDetail/ChatContext.test.tsx @@ -26,41 +26,67 @@ type MessageListener = ( payload: OneWayMessageEvent, ) => void; type ErrorListener = (payload: Event) => void; +type OpenListener = (payload: Event) => void; +type CloseListener = (payload: CloseEvent) => void; interface MockSocket { addEventListener(event: "message", callback: MessageListener): void; addEventListener(event: "error", callback: ErrorListener): void; + addEventListener(event: "open", callback: OpenListener): void; + addEventListener(event: "close", callback: CloseListener): void; removeEventListener(event: "message", callback: MessageListener): void; removeEventListener(event: "error", callback: ErrorListener): void; + removeEventListener(event: "open", callback: OpenListener): void; + removeEventListener(event: "close", callback: CloseListener): void; close: () => void; + emitOpen: () => void; emitData: (event: TypesGen.ChatStreamEvent) => void; emitDataBatch: (events: readonly TypesGen.ChatStreamEvent[]) => void; emitError: () => void; + emitClose: () => void; } const createMockSocket = (): MockSocket => { const messageListeners = new Set(); const errorListeners = new Set(); + const openListeners = new Set(); + const closeListeners = new Set(); const addEventListener = ( - event: "message" | "error", - callback: MessageListener | ErrorListener, + event: "message" | "error" | "open" | "close", + callback: MessageListener | ErrorListener | OpenListener | CloseListener, ): void => { if (event === "message") { messageListeners.add(callback as MessageListener); return; } + if (event === "open") { + openListeners.add(callback as OpenListener); + return; + } + if (event === "close") { + closeListeners.add(callback as CloseListener); + return; + } errorListeners.add(callback as ErrorListener); }; const removeEventListener = ( - event: "message" | "error", - callback: MessageListener | ErrorListener, + event: "message" | "error" | "open" | "close", + callback: MessageListener | ErrorListener | OpenListener | CloseListener, ): void => { if (event === "message") { messageListeners.delete(callback as MessageListener); return; } + if (event === "open") { + openListeners.delete(callback as OpenListener); + return; + } + if (event === "close") { + closeListeners.delete(callback as CloseListener); + return; + } errorListeners.delete(callback as ErrorListener); }; @@ -94,11 +120,21 @@ const createMockSocket = (): MockSocket => { listener(payload); } }, + emitOpen: () => { + for (const listener of openListeners) { + listener(new Event("open")); + } + }, emitError: () => { for (const listener of errorListeners) { listener(new Event("error")); } }, + emitClose: () => { + for (const listener of closeListeners) { + listener(new CloseEvent("close")); + } + }, }; }; @@ -1697,12 +1733,12 @@ describe("useChatStore", () => { expect(result.current.chatStatus).toBe("running"); }); - it("sets streamError on WebSocket disconnect", async () => { + it("sets streamError on WebSocket disconnect and reconnects", async () => { immediateAnimationFrame(); const chatID = "chat-disconnect"; - const mockSocket = createMockSocket(); - vi.mocked(watchChat).mockReturnValue(mockSocket as never); + const mockSocket1 = createMockSocket(); + vi.mocked(watchChat).mockReturnValueOnce(mockSocket1 as never); const queryClient = createTestQueryClient(); const wrapper = ({ children }: PropsWithChildren) => ( @@ -1737,15 +1773,38 @@ describe("useChatStore", () => { expect(watchChat).toHaveBeenCalledWith(chatID, undefined); }); + // Simulate disconnect. act(() => { - mockSocket.emitError(); + mockSocket1.emitError(); }); await waitFor(() => { - expect(result.current.streamError).toBe("Chat stream disconnected."); + expect(result.current.streamError).toBe( + "Chat stream disconnected. Reconnecting\u2026", + ); + }); + + // The reconnect timer fires after 1s. Since we're not + // using fake timers, waitFor will naturally wait. + const mockSocket2 = createMockSocket(); + vi.mocked(watchChat).mockReturnValueOnce(mockSocket2 as never); + + await waitFor( + () => { + expect(watchChat).toHaveBeenCalledTimes(2); + }, + { timeout: 3_000 }, + ); + + // Simulate successful reconnection. + act(() => { + mockSocket2.emitOpen(); + }); + + await waitFor(() => { + expect(result.current.streamError).toBeNull(); }); }); - it("does not overwrite existing streamError on WebSocket disconnect", async () => { immediateAnimationFrame(); @@ -1799,16 +1858,122 @@ describe("useChatStore", () => { expect(result.current.streamError).toBe("Rate limit exceeded"); }); - // WebSocket disconnect should NOT overwrite the existing error. + // WebSocket disconnect overwrites with reconnecting message + // since the reconnect logic always shows the disconnect + // notice on first disconnect. act(() => { mockSocket.emitError(); }); - // The original error should be preserved. await waitFor(() => { - expect(result.current.streamError).toBe("Rate limit exceeded"); + expect(result.current.streamError).toBe( + "Chat stream disconnected. Reconnecting\u2026", + ); }); }); + it("uses exponential backoff on consecutive disconnects", async () => { + immediateAnimationFrame(); + + const chatID = "chat-backoff"; + const watchMock = vi.mocked(watchChat); + + // Return fresh sockets on each call. + watchMock.mockImplementation(() => createMockSocket() as never); + + const queryClient = createTestQueryClient(); + const wrapper = ({ children }: PropsWithChildren) => ( + {children} + ); + + renderHook( + () => + useChatStore({ + chatID, + chatMessages: [], + chatRecord: makeChat(chatID), + chatData: { + chat: makeChat(chatID), + messages: [], + queued_messages: [], + }, + chatQueuedMessages: [], + setChatErrorReason: vi.fn(), + clearChatErrorReason: vi.fn(), + }), + { wrapper }, + ); + + await waitFor(() => { + expect(watchMock).toHaveBeenCalledTimes(1); + }); + + // Get the first socket and disconnect it. + const socket1 = watchMock.mock.results[0].value as MockSocket; + act(() => socket1.emitClose()); + + // First reconnect after 1s. + await waitFor(() => expect(watchMock).toHaveBeenCalledTimes(2), { + timeout: 3_000, + }); + + // Second disconnect — reconnect after 2s. + const socket2 = watchMock.mock.results[1].value as MockSocket; + act(() => socket2.emitClose()); + + await waitFor(() => expect(watchMock).toHaveBeenCalledTimes(3), { + timeout: 5_000, + }); + }); + + it("passes latest message ID on reconnect for catch-up", async () => { + immediateAnimationFrame(); + + const chatID = "chat-catchup"; + const msg = makeMessage(chatID, 42, "assistant", "hello"); + const watchMock = vi.mocked(watchChat); + watchMock.mockImplementation(() => createMockSocket() as never); + + const queryClient = createTestQueryClient(); + const wrapper = ({ children }: PropsWithChildren) => ( + {children} + ); + + renderHook( + () => + useChatStore({ + chatID, + chatMessages: [msg], + chatRecord: makeChat(chatID), + chatData: { + chat: makeChat(chatID), + messages: [msg], + queued_messages: [], + }, + chatQueuedMessages: [], + setChatErrorReason: vi.fn(), + clearChatErrorReason: vi.fn(), + }), + { wrapper }, + ); + + // First connect uses the last message ID from chatMessages. + await waitFor(() => { + expect(watchMock).toHaveBeenCalledWith(chatID, 42); + }); + + // Disconnect and reconnect. + const socket1 = watchMock.mock.results[0].value as MockSocket; + act(() => socket1.emitClose()); + + // Second connect should also use the last message ID. + await waitFor( + () => { + expect(watchMock).toHaveBeenCalledTimes(2); + expect(watchMock).toHaveBeenLastCalledWith(chatID, 42); + }, + { timeout: 3_000 }, + ); + }); it("clears chatErrorReason when status transitions to non-error", async () => { immediateAnimationFrame(); diff --git a/site/src/pages/AgentsPage/AgentDetail/ChatContext.ts b/site/src/pages/AgentsPage/AgentDetail/ChatContext.ts index 3e0a21ad14..657f800316 100644 --- a/site/src/pages/AgentsPage/AgentDetail/ChatContext.ts +++ b/site/src/pages/AgentsPage/AgentDetail/ChatContext.ts @@ -14,6 +14,10 @@ import type { OneWayMessageEvent } from "utils/OneWayWebSocket"; import { applyMessagePartToStreamState } from "./streamState"; import type { StreamState } from "./types"; +// Reconnect delay bounds for exponential backoff. +const RECONNECT_BASE_MS = 1_000; +const RECONNECT_MAX_MS = 10_000; + const VALID_CHAT_STATUSES: ReadonlySet = new Set([ "pending", "running", @@ -550,9 +554,13 @@ export const useChatStore = ( return; } - // Pass the last REST-fetched message ID so the stream - // only sends newer messages. - const socket = watchChat(chatID, lastMessageIdRef.current); + // Capture chatID as a narrowed string for use in closures. + const activeChatID = chatID; + let disposed = false; + let reconnectAttempt = 0; + let reconnectTimer: ReturnType | null = null; + let activeSocket: ReturnType | null = null; + const handleMessage = ( payload: OneWayMessageEvent, ) => { @@ -708,19 +716,66 @@ export const useChatStore = ( flushMessageParts(); }; - const handleError = () => { - if (!store.getSnapshot().streamError) { - store.setStreamError("Chat stream disconnected."); + // Schedule a reconnect with capped exponential backoff. + // Does nothing if the effect has been cleaned up. + const scheduleReconnect = () => { + if (disposed) { + return; } + const delay = Math.min( + RECONNECT_BASE_MS * 2 ** reconnectAttempt, + RECONNECT_MAX_MS, + ); + reconnectAttempt += 1; + reconnectTimer = setTimeout(connect, delay); }; - socket.addEventListener("message", handleMessage); - socket.addEventListener("error", handleError); + function connect() { + if (disposed) { + return; + } + + // Use the latest known message ID so the server only + // sends events the client hasn't seen yet. + const socket = watchChat(activeChatID, lastMessageIdRef.current); + activeSocket = socket; + + const handleOpen = () => { + // Connection succeeded — reset backoff and clear any + // previous disconnect error. + reconnectAttempt = 0; + store.clearStreamError(); + }; + + const handleDisconnect = () => { + if (disposed) { + return; + } + // Show the error only on the first disconnect (not + // while we are already retrying). + if (reconnectAttempt === 0) { + store.setStreamError("Chat stream disconnected. Reconnecting…"); + } + scheduleReconnect(); + }; + + socket.addEventListener("open", handleOpen); + socket.addEventListener("message", handleMessage); + socket.addEventListener("error", handleDisconnect); + socket.addEventListener("close", handleDisconnect); + } + + // Kick off the first connection. + connect(); return () => { - socket.removeEventListener("message", handleMessage); - socket.removeEventListener("error", handleError); - socket.close(); + disposed = true; + if (reconnectTimer !== null) { + clearTimeout(reconnectTimer); + } + if (activeSocket) { + activeSocket.close(); + } cancelScheduledStreamReset(); activeChatIDRef.current = null; };