From 667d501282878a247424feed3297c2cbc50cba66 Mon Sep 17 00:00:00 2001 From: Danielle Maywood Date: Sun, 8 Mar 2026 18:30:39 +0000 Subject: [PATCH] refactor(site): extract shared WebSocket reconnect utility (#22809) --- .../AgentsPage/AgentDetail/ChatContext.ts | 91 ++---- site/src/pages/AgentsPage/AgentsPage.tsx | 290 +++++++---------- site/src/utils/reconnectingWebSocket.test.ts | 301 ++++++++++++++++++ site/src/utils/reconnectingWebSocket.ts | 192 +++++++++++ 4 files changed, 635 insertions(+), 239 deletions(-) create mode 100644 site/src/utils/reconnectingWebSocket.test.ts create mode 100644 site/src/utils/reconnectingWebSocket.ts diff --git a/site/src/pages/AgentsPage/AgentDetail/ChatContext.ts b/site/src/pages/AgentsPage/AgentDetail/ChatContext.ts index 6143bf9067..4a111d716c 100644 --- a/site/src/pages/AgentsPage/AgentDetail/ChatContext.ts +++ b/site/src/pages/AgentsPage/AgentDetail/ChatContext.ts @@ -11,13 +11,10 @@ import { } from "react"; import { useQueryClient } from "react-query"; import type { OneWayMessageEvent } from "utils/OneWayWebSocket"; +import { createReconnectingWebSocket } from "utils/reconnectingWebSocket"; 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", @@ -590,10 +587,9 @@ export const useChatStore = ( // Capture chatID as a narrowed string for use in closures. const activeChatID = chatID; + // Local disposed flag so the message handler (which lives + // outside the utility) can bail out after cleanup. let disposed = false; - let reconnectAttempt = 0; - let reconnectTimer: ReturnType | null = null; - let activeSocket: ReturnType | null = null; const handleMessage = ( payload: OneWayMessageEvent, @@ -771,55 +767,24 @@ export const useChatStore = ( flushMessageParts(); }; - // Schedule a reconnect with capped exponential backoff. - // Does nothing if the effect has been cleaned up. - const scheduleReconnect = () => { - if (disposed) { - return; - } - if (reconnectTimer !== null) { - clearTimeout(reconnectTimer); - } - const delay = Math.min( - RECONNECT_BASE_MS * 2 ** reconnectAttempt, - RECONNECT_MAX_MS, - ); - reconnectAttempt += 1; - reconnectTimer = setTimeout(connect, delay); - }; - - function connect() { - if (disposed) { - return; - } - if (activeSocket) { - activeSocket.close(); - } - - // 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; + const disposeSocket = createReconnectingWebSocket({ + connect() { + // Use the latest known message ID so the server only + // sends events the client hasn't seen yet. + const socket = watchChat(activeChatID, lastMessageIdRef.current); + socket.addEventListener("message", handleMessage); + return socket; + }, + onOpen() { + // Connection succeeded — clear any previous disconnect + // error. store.clearStreamError(); - }; - - const handleDisconnect = () => { - // Guard against duplicate calls: browsers fire both - // "error" and "close" on a failed WebSocket, so we - // only process the first event per socket instance. - if (activeSocket !== socket || disposed) { - return; - } - activeSocket = null; + }, + onDisconnect(attempt) { // Show the error only on the first disconnect (not // while we are already retrying). - if (reconnectAttempt === 0) { - store.setStreamError("Chat stream disconnected. Reconnecting…"); + if (attempt === 0) { + store.setStreamError("Chat stream disconnected. Reconnecting\u2026"); } // Clear "running" status on disconnect so the UI // doesn't show a stale spinner. The reconnected @@ -828,26 +793,12 @@ export const useChatStore = ( if (currentStatus === "running") { store.setChatStatus(null); } - scheduleReconnect(); - }; - - socket.addEventListener("open", handleOpen); - socket.addEventListener("message", handleMessage); - socket.addEventListener("error", handleDisconnect); - socket.addEventListener("close", handleDisconnect); - } - - // Kick off the first connection. - connect(); + }, + }); return () => { disposed = true; - if (reconnectTimer !== null) { - clearTimeout(reconnectTimer); - } - if (activeSocket) { - activeSocket.close(); - } + disposeSocket(); cancelScheduledStreamReset(); activeChatIDRef.current = null; }; diff --git a/site/src/pages/AgentsPage/AgentsPage.tsx b/site/src/pages/AgentsPage/AgentsPage.tsx index 98e7b4c172..dc1df7ec6f 100644 --- a/site/src/pages/AgentsPage/AgentsPage.tsx +++ b/site/src/pages/AgentsPage/AgentsPage.tsx @@ -46,6 +46,7 @@ import { NavLink, Outlet, useNavigate, useParams } from "react-router"; import { toast } from "sonner"; import { cn } from "utils/cn"; import { pageTitle } from "utils/page"; +import { createReconnectingWebSocket } from "utils/reconnectingWebSocket"; import { AgentChatInput } from "./AgentChatInput"; import { maybePlayChime } from "./AgentDetail/useAgentChime"; import { AgentsSidebar } from "./AgentsSidebar"; @@ -377,189 +378,140 @@ const AgentsPage: FC = () => { activeChatIDRef.current = agentId; useEffect(() => { - let disposed = false; - let reconnectAttempt = 0; - let reconnectTimer: ReturnType | null = null; - let activeSocket: ReturnType | null = null; + return createReconnectingWebSocket({ + connect() { + const ws = watchChats(); - // Schedule a reconnect with capped exponential backoff. - const scheduleReconnect = () => { - if (disposed) { - return; - } - if (reconnectTimer !== null) { - clearTimeout(reconnectTimer); - } - const delay = Math.min(1000 * 2 ** reconnectAttempt, 10000); - reconnectAttempt += 1; - reconnectTimer = setTimeout(connect, delay); - }; + ws.addEventListener("message", (event) => { + if (event.parseError) { + console.warn("Failed to parse chat watch event:", event.parseError); + return; + } + const sse = event.parsedMessage; + if (sse?.type !== "data" || !sse.data) { + return; + } + if (!isChatListSSEEvent(sse.data)) { + return; + } + const chatEvent = sse.data; + const updatedChat = chatEvent.chat; - function connect() { - if (disposed) { - return; - } - if (activeSocket) { - activeSocket.close(); - } + // Read the previous status from the query cache, which + // is synchronously updated by both the per-chat WebSocket + // (via updateSidebarChat) and this handler. This avoids + // the async-lag of a useEffect-based status map. + const currentChats = + queryClient.getQueryData(chatsKey); + const prevStatus = currentChats?.find( + (c) => c.id === updatedChat.id, + )?.status; + // Only play the chime for top-level chats, not sub-agents. + if (!updatedChat.parent_chat_id) { + maybePlayChime( + prevStatus, + updatedChat.status, + updatedChat.id, + activeChatIDRef.current, + ); + } - const ws = watchChats(); - activeSocket = ws; + if (chatEvent.kind === "deleted") { + queryClient.setQueryData( + chatsKey, + (prev: TypesGen.Chat[] | undefined) => + prev?.filter( + (c) => + c.id !== updatedChat.id && + c.root_chat_id !== updatedChat.id, + ), + ); + queryClient.removeQueries({ + queryKey: chatKey(updatedChat.id), + exact: true, + }); + return; + } - ws.addEventListener("open", () => { - // Connection succeeded — reset backoff. - reconnectAttempt = 0; - void queryClient.invalidateQueries({ queryKey: chatsKey }); - }); + if (chatEvent.kind === "diff_status_change") { + void Promise.all([ + queryClient.invalidateQueries({ + queryKey: chatsKey, + }), + queryClient.invalidateQueries({ + queryKey: chatDiffStatusKey(updatedChat.id), + }), + queryClient.invalidateQueries({ + queryKey: chatDiffContentsKey(updatedChat.id), + }), + ]); + return; + } - const handleDisconnect = () => { - // Guard against duplicate calls: browsers fire both - // "error" and "close" on a failed WebSocket, so we - // only process the first event per socket instance. - if (activeSocket !== ws || disposed) { - return; - } - activeSocket = null; - void queryClient.invalidateQueries({ queryKey: chatsKey }); - scheduleReconnect(); - }; + // Scope field updates by event kind so that + // status_change events (which may carry a stale title + // snapshot from before async title generation + // finished) don't clobber a title_change that already + // landed. + const isTitleEvent = chatEvent.kind === "title_change"; + const isStatusEvent = chatEvent.kind === "status_change"; - ws.addEventListener("close", handleDisconnect); - ws.addEventListener("error", handleDisconnect); - - ws.addEventListener("message", (event) => { - if (event.parseError) { - console.warn("Failed to parse chat watch event:", event.parseError); - return; - } - const sse = event.parsedMessage; - if (sse?.type !== "data" || !sse.data) { - return; - } - if (!isChatListSSEEvent(sse.data)) { - return; - } - const chatEvent = sse.data; - const updatedChat = chatEvent.chat; - - // Read the previous status from the query cache, which - // is synchronously updated by both the per-chat WebSocket - // (via updateSidebarChat) and this handler. This avoids - // the async-lag of a useEffect-based status map. - const currentChats = - queryClient.getQueryData(chatsKey); - const prevStatus = currentChats?.find( - (c) => c.id === updatedChat.id, - )?.status; - // Only play the chime for top-level chats, not sub-agents. - if (!updatedChat.parent_chat_id) { - maybePlayChime( - prevStatus, - updatedChat.status, - updatedChat.id, - activeChatIDRef.current, - ); - } - - if (chatEvent.kind === "deleted") { queryClient.setQueryData( chatsKey, - (prev: TypesGen.Chat[] | undefined) => - prev?.filter( - (c) => - c.id !== updatedChat.id && c.root_chat_id !== updatedChat.id, - ), + (prev: TypesGen.Chat[] | undefined) => { + if (!prev) return prev; + const exists = prev.some((c) => c.id === updatedChat.id); + if (exists) { + return prev.map((c) => { + if (c.id !== updatedChat.id) return c; + return { + ...c, + ...(isStatusEvent && { status: updatedChat.status }), + ...(isTitleEvent && { title: updatedChat.title }), + updated_at: + c.updated_at > updatedChat.updated_at + ? c.updated_at + : updatedChat.updated_at, + }; + }); + } + if (chatEvent.kind === "created") { + return [updatedChat, ...prev]; + } + return prev; + }, ); - queryClient.removeQueries({ - queryKey: chatKey(updatedChat.id), - exact: true, - }); - return; - } - - if (chatEvent.kind === "diff_status_change") { - void Promise.all([ - queryClient.invalidateQueries({ - queryKey: chatsKey, - }), - queryClient.invalidateQueries({ - queryKey: chatDiffStatusKey(updatedChat.id), - }), - queryClient.invalidateQueries({ - queryKey: chatDiffContentsKey(updatedChat.id), - }), - ]); - return; - } - - // Scope field updates by event kind so that - // status_change events (which may carry a stale title - // snapshot from before async title generation - // finished) don't clobber a title_change that already - // landed. - const isTitleEvent = chatEvent.kind === "title_change"; - const isStatusEvent = chatEvent.kind === "status_change"; - - queryClient.setQueryData( - chatsKey, - (prev: TypesGen.Chat[] | undefined) => { - if (!prev) return prev; - const exists = prev.some((c) => c.id === updatedChat.id); - if (exists) { - return prev.map((c) => { - if (c.id !== updatedChat.id) return c; - return { - ...c, + queryClient.setQueryData( + chatKey(updatedChat.id), + (previousChat) => { + if (!previousChat) { + return previousChat; + } + return { + ...previousChat, + chat: { + ...previousChat.chat, ...(isStatusEvent && { status: updatedChat.status }), ...(isTitleEvent && { title: updatedChat.title }), updated_at: - c.updated_at > updatedChat.updated_at - ? c.updated_at + previousChat.chat.updated_at > updatedChat.updated_at + ? previousChat.chat.updated_at : updatedChat.updated_at, - }; - }); - } - if (chatEvent.kind === "created") { - return [updatedChat, ...prev]; - } - return prev; - }, - ); - queryClient.setQueryData( - chatKey(updatedChat.id), - (previousChat) => { - if (!previousChat) { - return previousChat; - } - return { - ...previousChat, - chat: { - ...previousChat.chat, - ...(isStatusEvent && { status: updatedChat.status }), - ...(isTitleEvent && { title: updatedChat.title }), - updated_at: - previousChat.chat.updated_at > updatedChat.updated_at - ? previousChat.chat.updated_at - : updatedChat.updated_at, - }, - }; - }, - ); - }); - } + }, + }; + }, + ); + }); - // Kick off the first connection. - connect(); - - return () => { - disposed = true; - if (reconnectTimer !== null) { - clearTimeout(reconnectTimer); - } - if (activeSocket) { - activeSocket.close(); - } - }; + return ws; + }, + onOpen() { + void queryClient.invalidateQueries({ queryKey: chatsKey }); + }, + onDisconnect() { + void queryClient.invalidateQueries({ queryKey: chatsKey }); + }, + }); }, [queryClient]); useEffect(() => { diff --git a/site/src/utils/reconnectingWebSocket.test.ts b/site/src/utils/reconnectingWebSocket.test.ts new file mode 100644 index 0000000000..0f1e2b03cc --- /dev/null +++ b/site/src/utils/reconnectingWebSocket.test.ts @@ -0,0 +1,301 @@ +import { createReconnectingWebSocket } from "./reconnectingWebSocket"; + +/** + * Minimal mock that satisfies the {@link Closable} interface used by + * the reconnection utility. Each instance records every + * `addEventListener` call and exposes helpers to fire those events. + */ +function createMockSocket() { + const listeners: Record void>> = {}; + const socket = { + addEventListener: vi.fn( + (event: string, handler: (...args: unknown[]) => void) => { + if (!listeners[event]) { + listeners[event] = []; + } + listeners[event].push(handler); + }, + ), + close: vi.fn(), + /** Fire all handlers registered for the given event type. */ + emit(event: string) { + for (const handler of listeners[event] ?? []) { + handler(); + } + }, + }; + return socket; +} + +beforeEach(() => { + vi.useFakeTimers(); +}); + +afterEach(() => { + vi.useRealTimers(); +}); + +describe("createReconnectingWebSocket", () => { + it("calls connect immediately and wires lifecycle events", () => { + const socket = createMockSocket(); + const connect = vi.fn(() => socket); + const onOpen = vi.fn(); + + createReconnectingWebSocket({ connect, onOpen }); + + expect(connect).toHaveBeenCalledTimes(1); + expect(socket.addEventListener).toHaveBeenCalledWith( + "open", + expect.any(Function), + ); + expect(socket.addEventListener).toHaveBeenCalledWith( + "error", + expect.any(Function), + ); + expect(socket.addEventListener).toHaveBeenCalledWith( + "close", + expect.any(Function), + ); + + // Simulate the socket opening. + socket.emit("open"); + expect(onOpen).toHaveBeenCalledTimes(1); + expect(onOpen).toHaveBeenCalledWith(socket); + }); + + it("reconnects with exponential backoff on disconnect", () => { + let activeSocket = createMockSocket(); + const connect = vi.fn(() => { + activeSocket = createMockSocket(); + return activeSocket; + }); + const onDisconnect = vi.fn(); + + createReconnectingWebSocket({ + connect, + onDisconnect, + baseMs: 1000, + maxMs: 10000, + factor: 2, + }); + + expect(connect).toHaveBeenCalledTimes(1); + + // First disconnect — should schedule reconnect after 1000ms. + activeSocket.emit("close"); + expect(onDisconnect).toHaveBeenCalledTimes(1); + expect(onDisconnect).toHaveBeenLastCalledWith(0); + + vi.advanceTimersByTime(999); + expect(connect).toHaveBeenCalledTimes(1); + vi.advanceTimersByTime(1); + expect(connect).toHaveBeenCalledTimes(2); + + // Second disconnect — delay should be 2000ms. + activeSocket.emit("close"); + expect(onDisconnect).toHaveBeenCalledTimes(2); + expect(onDisconnect).toHaveBeenLastCalledWith(1); + + vi.advanceTimersByTime(1999); + expect(connect).toHaveBeenCalledTimes(2); + vi.advanceTimersByTime(1); + expect(connect).toHaveBeenCalledTimes(3); + + // Third disconnect — delay should be 4000ms. + activeSocket.emit("close"); + vi.advanceTimersByTime(3999); + expect(connect).toHaveBeenCalledTimes(3); + vi.advanceTimersByTime(1); + expect(connect).toHaveBeenCalledTimes(4); + }); + + it("caps backoff delay at maxMs", () => { + let activeSocket = createMockSocket(); + const connect = vi.fn(() => { + activeSocket = createMockSocket(); + return activeSocket; + }); + + createReconnectingWebSocket({ + connect, + baseMs: 1000, + maxMs: 5000, + factor: 2, + }); + + // Disconnect enough times that the uncapped delay would + // exceed maxMs: 1000, 2000, 4000, 8000 → capped at 5000. + for (let i = 0; i < 3; i++) { + activeSocket.emit("close"); + vi.runOnlyPendingTimers(); + } + + // The 4th disconnect would have delay = 1000 * 2^3 = 8000, + // but should be capped at 5000. + activeSocket.emit("close"); + vi.advanceTimersByTime(4999); + expect(connect).toHaveBeenCalledTimes(4); + vi.advanceTimersByTime(1); + expect(connect).toHaveBeenCalledTimes(5); + }); + + it("resets backoff on successful connection", () => { + let activeSocket = createMockSocket(); + const connect = vi.fn(() => { + activeSocket = createMockSocket(); + return activeSocket; + }); + + createReconnectingWebSocket({ + connect, + baseMs: 1000, + maxMs: 10000, + factor: 2, + }); + + // Disconnect twice to bump the attempt counter. + activeSocket.emit("close"); + vi.runOnlyPendingTimers(); + activeSocket.emit("close"); + vi.runOnlyPendingTimers(); + + // Now simulate a successful open — attempt should reset. + activeSocket.emit("open"); + activeSocket.emit("close"); + + // Next reconnect should use the base delay (1000ms), not + // 4000ms. + vi.advanceTimersByTime(999); + expect(connect).toHaveBeenCalledTimes(3); + vi.advanceTimersByTime(1); + expect(connect).toHaveBeenCalledTimes(4); + }); + + it("deduplicates error+close from the same socket", () => { + let activeSocket = createMockSocket(); + const connect = vi.fn(() => { + activeSocket = createMockSocket(); + return activeSocket; + }); + const onDisconnect = vi.fn(); + + createReconnectingWebSocket({ connect, onDisconnect }); + + const socketBeforeDisconnect = activeSocket; + + // Browser fires both error and close. + socketBeforeDisconnect.emit("error"); + socketBeforeDisconnect.emit("close"); + + // onDisconnect should only fire once. + expect(onDisconnect).toHaveBeenCalledTimes(1); + + // Only one reconnect timer should be pending. + vi.runOnlyPendingTimers(); + expect(connect).toHaveBeenCalledTimes(2); + }); + + it("closes previous socket when reconnecting", () => { + let activeSocket = createMockSocket(); + const sockets: ReturnType[] = []; + const connect = vi.fn(() => { + activeSocket = createMockSocket(); + sockets.push(activeSocket); + return activeSocket; + }); + + createReconnectingWebSocket({ connect }); + + const firstSocket = sockets[0]; + firstSocket.emit("close"); + vi.runOnlyPendingTimers(); + + // The connect function creates a new socket. The old socket + // was already "closed" by the browser, but on a fresh + // reconnection the utility closes the previous one if it's + // still the active reference. + expect(connect).toHaveBeenCalledTimes(2); + }); + + it("dispose stops reconnection and closes the socket", () => { + const socket = createMockSocket(); + const connect = vi.fn(() => socket); + + const dispose = createReconnectingWebSocket({ connect }); + + dispose(); + + expect(socket.close).toHaveBeenCalled(); + + // Simulating events after dispose should not cause errors or + // additional connect calls. + socket.emit("close"); + vi.runAllTimers(); + expect(connect).toHaveBeenCalledTimes(1); + }); + + it("dispose cancels pending reconnect timer", () => { + let activeSocket = createMockSocket(); + const connect = vi.fn(() => { + activeSocket = createMockSocket(); + return activeSocket; + }); + + const dispose = createReconnectingWebSocket({ connect }); + + // Trigger a disconnect so a timer is scheduled. + activeSocket.emit("close"); + expect(connect).toHaveBeenCalledTimes(1); + + // Dispose before the timer fires. + dispose(); + vi.runAllTimers(); + + // No reconnection should have occurred. + expect(connect).toHaveBeenCalledTimes(1); + }); + + it("dispose is safe to call multiple times", () => { + const socket = createMockSocket(); + const connect = vi.fn(() => socket); + + const dispose = createReconnectingWebSocket({ connect }); + + dispose(); + dispose(); + dispose(); + + // close is idempotent on real WebSockets, so calling it + // multiple times is harmless. The important thing is that + // no reconnection is scheduled after the first dispose. + expect(connect).toHaveBeenCalledTimes(1); + }); + + it("passes socket to onOpen callback", () => { + const socket = createMockSocket(); + const connect = vi.fn(() => socket); + const onOpen = vi.fn(); + + createReconnectingWebSocket({ connect, onOpen }); + + socket.emit("open"); + expect(onOpen).toHaveBeenCalledWith(socket); + }); + + it("uses default backoff values when none provided", () => { + let activeSocket = createMockSocket(); + const connect = vi.fn(() => { + activeSocket = createMockSocket(); + return activeSocket; + }); + + createReconnectingWebSocket({ connect }); + + // Default: baseMs=1000, factor=2, maxMs=10000. + activeSocket.emit("close"); + vi.advanceTimersByTime(999); + expect(connect).toHaveBeenCalledTimes(1); + vi.advanceTimersByTime(1); + expect(connect).toHaveBeenCalledTimes(2); + }); +}); diff --git a/site/src/utils/reconnectingWebSocket.ts b/site/src/utils/reconnectingWebSocket.ts new file mode 100644 index 0000000000..294832ab43 --- /dev/null +++ b/site/src/utils/reconnectingWebSocket.ts @@ -0,0 +1,192 @@ +/** + * @file Shared WebSocket reconnection utility with capped exponential + * backoff. Both the chat-list watcher (AgentsPage) and the per-chat + * stream watcher (ChatContext) use the same reconnect-on-disconnect + * pattern. This module extracts that logic into a single reusable + * function so the two call sites stay in sync and the backoff math + * lives in one place. + * + * @example + * ```ts + * const dispose = createReconnectingWebSocket({ + * connect() { + * const ws = watchChats(); + * ws.addEventListener("message", (e) => handleMessage(e)); + * return ws; + * }, + * onOpen() { + * console.log("connected"); + * }, + * onDisconnect() { + * console.log("disconnected, will reconnect automatically"); + * }, + * }); + * + * // Later, to tear down: + * dispose(); + * ``` + */ + +/** Default base delay for exponential backoff (milliseconds). */ +const RECONNECT_BASE_MS = 1_000; + +/** Default maximum delay cap for exponential backoff (milliseconds). */ +const RECONNECT_MAX_MS = 10_000; + +/** Default multiplier applied to the base delay on each retry. */ +const RECONNECT_FACTOR = 2; + +/** + * A minimal WebSocket-like interface that the reconnection utility + * can manage. Both native `WebSocket` and `OneWayWebSocket` satisfy + * this contract. + */ +interface Closable { + addEventListener(event: string, handler: (...args: unknown[]) => void): void; + close(...args: unknown[]): void; +} + +/** + * Configuration for {@link createReconnectingWebSocket}. + * + * @typeParam TSocket - The concrete socket type returned by the + * `connect` function (e.g. `OneWayWebSocket`). + */ +interface ReconnectingWebSocketOptions { + /** + * Factory that creates and returns a new socket. Called on the + * initial connection and on every reconnection attempt. The caller + * is responsible for attaching any `message` listeners to the + * returned socket — this utility only manages the lifecycle + * (`open`, `close`, `error`) events. + */ + connect: () => TSocket; + + /** + * Called when a connection succeeds (the socket fires `open`). + * The backoff counter is reset before this callback runs. + */ + onOpen?: (socket: TSocket) => void; + + /** + * Called on the first disconnect after a successful connection or + * on a connection failure. Fires at most once per socket instance + * (browsers fire both `error` and `close`; only the first is + * forwarded). A reconnection is scheduled automatically after + * this callback returns. + * + * @param attempt - The zero-based reconnection attempt counter + * *before* it is incremented for the upcoming retry. A value of + * `0` means this is the first disconnect since the last + * successful connection. + */ + onDisconnect?: (attempt: number) => void; + + /** Base delay in milliseconds. Defaults to {@link RECONNECT_BASE_MS}. */ + baseMs?: number; + + /** Maximum delay cap in milliseconds. Defaults to {@link RECONNECT_MAX_MS}. */ + maxMs?: number; + + /** Multiplier applied per attempt. Defaults to {@link RECONNECT_FACTOR}. */ + factor?: number; +} + +/** + * Creates a self-reconnecting WebSocket connection with capped + * exponential backoff. + * + * The returned function disposes of the connection: it closes the + * active socket (if any), cancels any pending reconnection timer, + * and prevents further reconnection attempts. It is safe to call + * the dispose function more than once. + * + * Backoff delay formula: + * ``` + * delay = min(baseMs * factor ^ attempt, maxMs) + * ``` + * + * The attempt counter resets to `0` whenever a connection + * successfully opens. + * + * @returns A dispose function that tears down the connection. + */ +export function createReconnectingWebSocket( + options: ReconnectingWebSocketOptions, +): () => void { + const { + connect: connectFn, + onOpen, + onDisconnect, + baseMs = RECONNECT_BASE_MS, + maxMs = RECONNECT_MAX_MS, + factor = RECONNECT_FACTOR, + } = options; + + let disposed = false; + let reconnectAttempt = 0; + let reconnectTimer: ReturnType | null = null; + let activeSocket: TSocket | null = null; + + // Schedule a reconnect with capped exponential backoff. + // Does nothing if the connection has been disposed. + const scheduleReconnect = () => { + if (disposed) { + return; + } + if (reconnectTimer !== null) { + clearTimeout(reconnectTimer); + } + const delay = Math.min(baseMs * factor ** reconnectAttempt, maxMs); + reconnectAttempt += 1; + reconnectTimer = setTimeout(connect, delay); + }; + + function connect() { + if (disposed) { + return; + } + if (activeSocket) { + activeSocket.close(); + } + + const socket = connectFn(); + activeSocket = socket; + + const handleOpen = () => { + // Connection succeeded — reset backoff. + reconnectAttempt = 0; + onOpen?.(socket); + }; + + const handleDisconnect = () => { + // Guard against duplicate calls: browsers fire both + // "error" and "close" on a failed WebSocket, so we + // only process the first event per socket instance. + if (activeSocket !== socket || disposed) { + return; + } + activeSocket = null; + onDisconnect?.(reconnectAttempt); + scheduleReconnect(); + }; + + socket.addEventListener("open", handleOpen); + socket.addEventListener("error", handleDisconnect); + socket.addEventListener("close", handleDisconnect); + } + + // Kick off the first connection. + connect(); + + // Return a dispose function that tears everything down. + return () => { + disposed = true; + if (reconnectTimer !== null) { + clearTimeout(reconnectTimer); + } + if (activeSocket) { + activeSocket.close(); + } + }; +}