fix(site): add WebSocket reconnection with exponential backoff to chat stream (#22662)

## 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.

site/src/pages/AgentsPage/AgentDetail/ChatContext.test.tsx

@@ -26,41 +26,67 @@ type MessageListener = (
 	payload: OneWayMessageEvent<TypesGen.ServerSentEvent>,
 ) => 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<MessageListener>();
 	const errorListeners = new Set<ErrorListener>();
+	const openListeners = new Set<OpenListener>();
+	const closeListeners = new Set<CloseListener>();
 
 	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) => (
+			<QueryClientProvider client={queryClient}>{children}</QueryClientProvider>
+		);
+
+		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) => (
+			<QueryClientProvider client={queryClient}>{children}</QueryClientProvider>
+		);
+
+		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();

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<string> = new Set<TypesGen.ChatStatus>([
 	"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<typeof setTimeout> | null = null;
+		let activeSocket: ReturnType<typeof watchChat> | null = null;
+
 		const handleMessage = (
 			payload: OneWayMessageEvent<TypesGen.ServerSentEvent>,
 		) => {
@@ -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;
 		};