mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
refactor(site): extract shared WebSocket reconnect utility (#22809)
This commit is contained in:
Danielle Maywood
@@ -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<string> = new Set<TypesGen.ChatStatus>([
|
||||
"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<typeof setTimeout> | null = null;
|
||||
let activeSocket: ReturnType<typeof watchChat> | null = null;
|
||||
|
||||
const handleMessage = (
|
||||
payload: OneWayMessageEvent<TypesGen.ServerSentEvent>,
|
||||
@@ -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;
|
||||
};
|
||||
|
||||
@@ -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<typeof setTimeout> | null = null;
|
||||
let activeSocket: ReturnType<typeof watchChats> | 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<TypesGen.Chat[]>(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<TypesGen.Chat[]>(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<TypesGen.ChatWithMessages | undefined>(
|
||||
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<TypesGen.ChatWithMessages | undefined>(
|
||||
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(() => {
|
||||
|
||||
@@ -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<string, Array<(...args: unknown[]) => 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<typeof createMockSocket>[] = [];
|
||||
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);
|
||||
});
|
||||
});
|
||||
@@ -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<ServerSentEvent>`).
|
||||
*/
|
||||
interface ReconnectingWebSocketOptions<TSocket extends Closable> {
|
||||
/**
|
||||
* 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<TSocket extends Closable>(
|
||||
options: ReconnectingWebSocketOptions<TSocket>,
|
||||
): () => 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<typeof setTimeout> | 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();
|
||||
}
|
||||
};
|
||||
}
|
||||
Reference in New Issue
Block a user