authorgravatar for julius0216@outlook.comJulius Marminge <julius0216@outlook.com> 2026-04-09 17:23:37-07:00
committergravatar for noreply@github.comGitHub <noreply@github.com> 2026-04-09 17:23:37-07:00
log528bb2a1fa318962437281e68723bf0bb5a589d7
treea9f7cbfa2212d493f7bc24ba03fa1b2edff91c64
parent869789b4edd7ff71ed717feb1dc1935aef36e5ad
signature Signed by PGP key B5690EEEBB952194

[codex] Harden WebSocket reconnect recovery (#1864)

Co-authored-by: codex <codex@users.noreply.github.com>

8 files changed, 237 insertions(+), 83 deletions(-)

apps/web/src/components/WebSocketConnectionSurface.logic.test.ts+31-1
......@@ -1,7 +1,7 @@
11import { describe, expect, it } from "vitest";
22
33import type { WsConnectionStatus } from "../rpc/wsConnectionState";
4import { shouldAutoReconnect } from "./WebSocketConnectionSurface";
4import { shouldAutoReconnect, shouldRestartStalledReconnect } from "./WebSocketConnectionSurface";
55
66function makeStatus(overrides: Partial<WsConnectionStatus> = {}): WsConnectionStatus {
77 return {
......@@ -80,4 +80,34 @@ describe("WebSocketConnectionSurface.logic", () => {
8080 ),
8181 ).toBe(true);
8282 });
83
84 it("restarts a stalled reconnect window after the scheduled retry time passes", () => {
85 expect(
86 shouldRestartStalledReconnect(
87 makeStatus({
88 hasConnected: true,
89 nextRetryAt: "2026-04-03T20:00:01.000Z",
90 online: true,
91 phase: "disconnected",
92 reconnectAttemptCount: 3,
93 reconnectPhase: "waiting",
94 }),
95 "2026-04-03T20:00:01.000Z",
96 ),
97 ).toBe(true);
98
99 expect(
100 shouldRestartStalledReconnect(
101 makeStatus({
102 hasConnected: true,
103 nextRetryAt: "2026-04-03T20:00:01.000Z",
104 online: true,
105 phase: "disconnected",
106 reconnectAttemptCount: 3,
107 reconnectPhase: "attempting",
108 }),
109 "2026-04-03T20:00:01.000Z",
110 ),
111 ).toBe(false);
112 });
83113});
apps/web/src/components/WebSocketConnectionSurface.tsx+18-2
......@@ -2,7 +2,6 @@ import { type ReactNode, useEffect, useEffectEvent, useRef, useState } from "rea
22
33import { type SlowRpcAckRequest, useSlowRpcAckRequests } from "../rpc/requestLatencyState";
44import {
5 exhaustWsReconnectIfStillWaiting,
65 getWsConnectionStatus,
76 getWsConnectionUiState,
87 setBrowserOnlineStatus,
......@@ -105,6 +104,18 @@ export function shouldAutoReconnect(
105104 );
106105}
107106
107export function shouldRestartStalledReconnect(
108 status: WsConnectionStatus,
109 expectedNextRetryAt: string,
110): boolean {
111 return (
112 status.reconnectPhase === "waiting" &&
113 status.nextRetryAt === expectedNextRetryAt &&
114 status.online &&
115 status.hasConnected
116 );
117}
118
108119export function WebSocketConnectionCoordinator() {
109120 const status = useWsConnectionStatus();
110121 const [nowMs, setNowMs] = useState(() => Date.now());
......@@ -205,7 +216,12 @@ export function WebSocketConnectionCoordinator() {
205216 const nextRetryAt = status.nextRetryAt;
206217 const timeoutMs = Math.max(0, new Date(nextRetryAt).getTime() - Date.now()) + 1_500;
207218 const timeoutId = window.setTimeout(() => {
208 exhaustWsReconnectIfStillWaiting(nextRetryAt);
219 const currentStatus = getWsConnectionStatus();
220 if (!shouldRestartStalledReconnect(currentStatus, nextRetryAt)) {
221 return;
222 }
223
224 runReconnect(false);
209225 }, timeoutMs);
210226
211227 return () => {
apps/web/src/environments/runtime/connection.test.ts+59
......@@ -100,6 +100,7 @@ function createTestClient(options?: {
100100 return {
101101 client,
102102 getSnapshot,
103 replayEvents,
103104 emitWelcome: (environmentId: EnvironmentId) => {
104105 for (const listener of lifecycleListeners) {
105106 listener({
......@@ -227,6 +228,64 @@ describe("createEnvironmentConnection", () => {
227228 await connection.dispose();
228229 });
229230
231 it("retries replay recovery after transport disconnects during resubscribe", async () => {
232 const environmentId = EnvironmentId.makeUnsafe("env-1");
233 let replayAttempts = 0;
234 const applyEventBatch = vi.fn();
235 const { client, replayEvents, triggerDomainResubscribe } = createTestClient({
236 replayEvents: async () => {
237 replayAttempts += 1;
238 if (replayAttempts === 1) {
239 throw new Error("SocketCloseError: 1006");
240 }
241
242 return [
243 {
244 sequence: 2,
245 type: "thread.created",
246 payload: {},
247 },
248 ];
249 },
250 });
251
252 const connection = createEnvironmentConnection({
253 kind: "saved",
254 knownEnvironment: {
255 id: "env-1",
256 label: "Remote env",
257 source: "manual",
258 target: {
259 httpBaseUrl: "http://example.test",
260 wsBaseUrl: "ws://example.test",
261 },
262 environmentId,
263 },
264 client,
265 applyEventBatch,
266 syncSnapshot: vi.fn(),
267 applyTerminalEvent: vi.fn(),
268 });
269
270 await Promise.resolve();
271 await Promise.resolve();
272
273 triggerDomainResubscribe();
274
275 await vi.waitFor(() => {
276 expect(replayEvents).toHaveBeenCalledTimes(2);
277 expect(applyEventBatch).toHaveBeenCalledWith(
278 [
279 expect.objectContaining({
280 sequence: 2,
281 }),
282 ],
283 environmentId,
284 );
285 });
286
287 await connection.dispose();
288 });
230289 it("swallows replay recovery failures triggered by resubscribe", async () => {
231290 const environmentId = EnvironmentId.makeUnsafe("env-1");
232291 const snapshotError = new Error("snapshot failed");
apps/web/src/environments/runtime/connection.ts+36-2
......@@ -16,10 +16,13 @@ import {
1616 createOrchestrationRecoveryCoordinator,
1717 type ReplayRetryTracker,
1818} from "../../orchestrationRecovery";
19import { isTransportConnectionErrorMessage } from "~/rpc/transportError";
1920import type { WsRpcClient } from "~/rpc/wsRpcClient";
2021
2122const REPLAY_RECOVERY_RETRY_DELAY_MS = 100;
2223const MAX_NO_PROGRESS_REPLAY_RETRIES = 3;
24const RECOVERY_TRANSPORT_RETRY_DELAY_MS = 250;
25const MAX_RECOVERY_TRANSPORT_RETRIES = 20;
2326
2427export interface EnvironmentConnection {
2528 readonly kind: "primary" | "saved";
......@@ -128,6 +131,30 @@ export function createEnvironmentConnection(
128131 queueMicrotask(flushPendingDomainEvents);
129132 };
130133
134 const retryTransportRecoveryOperation = async <T>(operation: () => Promise<T>): Promise<T> => {
135 for (let attempt = 0; ; attempt += 1) {
136 try {
137 return await operation();
138 } catch (error) {
139 const message = error instanceof Error ? error.message : String(error);
140 if (
141 disposed ||
142 !isTransportConnectionErrorMessage(message) ||
143 attempt >= MAX_RECOVERY_TRANSPORT_RETRIES - 1
144 ) {
145 throw error;
146 }
147
148 await new Promise<void>((resolve) => {
149 setTimeout(resolve, RECOVERY_TRANSPORT_RETRY_DELAY_MS);
150 });
151
152 if (disposed) {
153 throw error;
154 }
155 }
156 }
157 };
131158 const scheduleReplayRecovery = (reason: "sequence-gap" | "resubscribe") => {
132159 void runReplayRecovery(reason).catch(() => undefined);
133160 };
......@@ -139,7 +166,9 @@ export function createEnvironmentConnection(
139166
140167 const fromSequenceExclusive = recovery.getState().latestSequence;
141168 try {
142 const events = await input.client.orchestration.replayEvents({ fromSequenceExclusive });
169 const events = await retryTransportRecoveryOperation(() =>
170 input.client.orchestration.replayEvents({ fromSequenceExclusive }),
171 );
143172 if (!disposed) {
144173 const nextEvents = recovery.markEventBatchApplied(events);
145174 if (nextEvents.length > 0) {
......@@ -149,6 +178,9 @@ export function createEnvironmentConnection(
149178 } catch {
150179 replayRetryTracker = null;
151180 recovery.failReplayRecovery();
181 if (disposed) {
182 return;
183 }
152184 await snapshotBootstrap.ensureSnapshotRecovery("replay-failed");
153185 return;
154186 }
......@@ -198,7 +230,9 @@ export function createEnvironmentConnection(
198230 }
199231
200232 try {
201 const snapshot = await input.client.orchestration.getSnapshot();
233 const snapshot = await retryTransportRecoveryOperation(() =>
234 input.client.orchestration.getSnapshot(),
235 );
202236 if (!disposed) {
203237 input.syncSnapshot(snapshot, environmentId);
204238 if (recovery.completeSnapshotRecovery(snapshot.snapshotSequence)) {
apps/web/src/rpc/protocol.ts+26-4
......@@ -62,14 +62,36 @@ function defaultLifecycleHandlers(): Required<WsProtocolLifecycleHandlers> {
6262 };
6363}
6464
65function composeLifecycleHandlers(
66 handlers?: WsProtocolLifecycleHandlers,
67): Required<WsProtocolLifecycleHandlers> {
68 const defaults = defaultLifecycleHandlers();
69
70 return {
71 onAttempt: (socketUrl) => {
72 defaults.onAttempt(socketUrl);
73 handlers?.onAttempt?.(socketUrl);
74 },
75 onOpen: () => {
76 defaults.onOpen();
77 handlers?.onOpen?.();
78 },
79 onError: (message) => {
80 defaults.onError(message);
81 handlers?.onError?.(message);
82 },
83 onClose: (details) => {
84 defaults.onClose(details);
85 handlers?.onClose?.(details);
86 },
87 };
88}
89
6590export function createWsRpcProtocolLayer(
6691 url: WsRpcProtocolSocketUrlProvider,
6792 handlers?: WsProtocolLifecycleHandlers,
6893) {
69 const lifecycle = {
70 ...defaultLifecycleHandlers(),
71 ...handlers,
72 };
94 const lifecycle = composeLifecycleHandlers(handlers);
7395 const resolvedUrl =
7496 typeof url === "function"
7597 ? Effect.promise(() => url()).pipe(
apps/web/src/rpc/wsConnectionState.test.ts-37
......@@ -1,7 +1,6 @@
11import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
22
33import {
4 exhaustWsReconnectIfStillWaiting,
54 getWsConnectionStatus,
65 getWsReconnectDelayMsForRetry,
76 getWsConnectionUiState,
......@@ -73,40 +72,4 @@ describe("wsConnectionState", () => {
7372 reconnectPhase: "exhausted",
7473 });
7574 });
76
77 it("can exhaust a stalled final retry window when no new attempt starts", () => {
78 recordWsConnectionAttempt("ws://localhost:3020/ws");
79 recordWsConnectionOpened();
80
81 for (let attempt = 0; attempt < WS_RECONNECT_MAX_ATTEMPTS - 1; attempt += 1) {
82 recordWsConnectionAttempt("ws://localhost:3020/ws");
83 recordWsConnectionErrored("Unable to connect to the T3 server WebSocket.");
84 }
85
86 const finalRetryDelayMs = getWsReconnectDelayMsForRetry(WS_RECONNECT_MAX_ATTEMPTS - 2);
87 if (finalRetryDelayMs === null) {
88 throw new Error("Expected a final retry delay.");
89 }
90
91 const statusBeforeExhaust = getWsConnectionStatus();
92 expect(statusBeforeExhaust).toMatchObject({
93 nextRetryAt: new Date(Date.now() + finalRetryDelayMs).toISOString(),
94 reconnectAttemptCount: 7,
95 reconnectPhase: "waiting",
96 });
97
98 const nextRetryAt = statusBeforeExhaust.nextRetryAt;
99 if (!nextRetryAt) {
100 throw new Error("Expected a scheduled retry.");
101 }
102
103 vi.setSystemTime(new Date(Date.now() + finalRetryDelayMs + 1_000));
104 exhaustWsReconnectIfStillWaiting(nextRetryAt);
105
106 expect(getWsConnectionStatus()).toMatchObject({
107 nextRetryAt: null,
108 reconnectAttemptCount: WS_RECONNECT_MAX_ATTEMPTS,
109 reconnectPhase: "exhausted",
110 });
111 });
11275});
apps/web/src/rpc/wsConnectionState.ts-20
......@@ -149,26 +149,6 @@ export function resetWsReconnectBackoff(): WsConnectionStatus {
149149 }));
150150}
151151
152export function exhaustWsReconnectIfStillWaiting(expectedNextRetryAt: string): WsConnectionStatus {
153 return updateWsConnectionStatus((current) => {
154 if (
155 current.reconnectPhase !== "waiting" ||
156 current.nextRetryAt !== expectedNextRetryAt ||
157 !current.online ||
158 !current.hasConnected
159 ) {
160 return current;
161 }
162
163 return {
164 ...current,
165 nextRetryAt: null,
166 reconnectAttemptCount: current.reconnectMaxAttempts,
167 reconnectPhase: "exhausted",
168 };
169 });
170}
171
172152export function resetWsConnectionStateForTests(): void {
173153 appAtomRegistry.set(wsConnectionStatusAtom, INITIAL_WS_CONNECTION_STATUS);
174154}
apps/web/src/rpc/wsTransport.test.ts+67-17
......@@ -83,6 +83,7 @@ class MockWebSocket {
8383
8484const originalWebSocket = globalThis.WebSocket;
8585const originalFetch = globalThis.fetch;
86const transports: WsTransport[] = [];
8687
8788function getSocket(): MockWebSocket {
8889 const socket = sockets.at(-1);
......@@ -107,9 +108,16 @@ async function waitFor(assertion: () => void, timeoutMs = 1_000): Promise<void>
107108 }
108109}
109110
111function createTransport(...args: ConstructorParameters<typeof WsTransport>): WsTransport {
112 const transport = new WsTransport(...args);
113 transports.push(transport);
114 return transport;
115}
116
110117beforeEach(() => {
111118 vi.useRealTimers();
112119 sockets.length = 0;
120 transports.length = 0;
113121 resetRequestLatencyStateForTests();
114122 resetWsConnectionStateForTests();
115123
......@@ -134,6 +142,8 @@ beforeEach(() => {
134142});
135143
136144afterEach(async () => {
145 await Promise.allSettled(transports.map((transport) => transport.dispose()));
146 transports.length = 0;
137147 globalThis.WebSocket = originalWebSocket;
138148 globalThis.fetch = originalFetch;
139149 resetRequestLatencyStateForTests();
......@@ -144,7 +154,7 @@ afterEach(async () => {
144154
145155describe("WsTransport", () => {
146156 it("normalizes root websocket urls to /ws and preserves query params", async () => {
147 const transport = new WsTransport("ws://localhost:3020/?token=secret-token");
157 const transport = createTransport("ws://localhost:3020/?token=secret-token");
148158
149159 await waitFor(() => {
150160 expect(sockets).toHaveLength(1);
......@@ -155,7 +165,7 @@ describe("WsTransport", () => {
155165 });
156166
157167 it("uses an explicit secure websocket base url", async () => {
158 const transport = new WsTransport("wss://app.example.com");
168 const transport = createTransport("wss://app.example.com");
159169
160170 await waitFor(() => {
161171 expect(sockets).toHaveLength(1);
......@@ -166,7 +176,7 @@ describe("WsTransport", () => {
166176 });
167177
168178 it("uses an explicit insecure websocket base url for remote backends", async () => {
169 const transport = new WsTransport("ws://192.168.1.44:3773");
179 const transport = createTransport("ws://192.168.1.44:3773");
170180
171181 await waitFor(() => {
172182 expect(sockets).toHaveLength(1);
......@@ -177,7 +187,7 @@ describe("WsTransport", () => {
177187 });
178188
179189 it("supports async websocket url providers", async () => {
180 const transport = new WsTransport(async () => "wss://remote.example.com/?wsToken=dynamic");
190 const transport = createTransport(async () => "wss://remote.example.com/?wsToken=dynamic");
181191
182192 await waitFor(() => {
183193 expect(sockets).toHaveLength(1);
......@@ -188,7 +198,7 @@ describe("WsTransport", () => {
188198 });
189199
190200 it("tracks initial connection failures for the app error state", async () => {
191 const transport = new WsTransport("ws://localhost:3020");
201 const transport = createTransport("ws://localhost:3020");
192202
193203 await waitFor(() => {
194204 expect(sockets).toHaveLength(1);
......@@ -219,7 +229,7 @@ describe("WsTransport", () => {
219229 });
220230
221231 it("surfaces reconnecting state after a live socket disconnects", async () => {
222 const transport = new WsTransport("ws://localhost:3020");
232 const transport = createTransport("ws://localhost:3020");
223233
224234 await waitFor(() => {
225235 expect(sockets).toHaveLength(1);
......@@ -248,8 +258,48 @@ describe("WsTransport", () => {
248258 await transport.dispose();
249259 });
250260
261 it("composes custom lifecycle handlers with default websocket state tracking", async () => {
262 const onOpen = vi.fn();
263 const onClose = vi.fn();
264 const transport = createTransport("ws://localhost:3020", {
265 onOpen,
266 onClose,
267 });
268
269 await waitFor(() => {
270 expect(sockets).toHaveLength(1);
271 });
272
273 const socket = getSocket();
274 socket.open();
275
276 await waitFor(() => {
277 expect(onOpen).toHaveBeenCalledOnce();
278 expect(getWsConnectionStatus()).toMatchObject({
279 hasConnected: true,
280 phase: "connected",
281 });
282 });
283
284 socket.close(1012, "service restart");
285
286 await waitFor(() => {
287 expect(onClose).toHaveBeenCalledWith({
288 code: 1012,
289 reason: "service restart",
290 });
291 expect(getWsConnectionStatus()).toMatchObject({
292 attemptCount: 2,
293 closeReason: "service restart",
294 phase: "connecting",
295 });
296 }, 2_000);
297
298 await transport.dispose();
299 });
300
251301 it("reconnects the websocket session without disposing the transport", async () => {
252 const transport = new WsTransport("ws://localhost:3020");
302 const transport = createTransport("ws://localhost:3020");
253303
254304 await waitFor(() => {
255305 expect(sockets).toHaveLength(1);
......@@ -314,7 +364,7 @@ describe("WsTransport", () => {
314364 it("marks unary requests as slow until the first server ack arrives", async () => {
315365 const slowAckThresholdMs = 25;
316366 setSlowRpcAckThresholdMsForTests(slowAckThresholdMs);
317 const transport = new WsTransport("ws://localhost:3020");
367 const transport = createTransport("ws://localhost:3020");
318368
319369 const requestPromise = transport.request((client) =>
320370 client[WS_METHODS.serverUpsertKeybinding]({
......@@ -368,7 +418,7 @@ describe("WsTransport", () => {
368418 }, 5_000);
369419
370420 it("sends unary RPC requests and resolves successful exits", async () => {
371 const transport = new WsTransport("ws://localhost:3020");
421 const transport = createTransport("ws://localhost:3020");
372422
373423 const requestPromise = transport.request((client) =>
374424 client[WS_METHODS.serverUpsertKeybinding]({
......@@ -426,7 +476,7 @@ describe("WsTransport", () => {
426476 });
427477
428478 it("delivers stream chunks to subscribers", async () => {
429 const transport = new WsTransport("ws://localhost:3020");
479 const transport = createTransport("ws://localhost:3020");
430480 const listener = vi.fn();
431481
432482 const unsubscribe = transport.subscribe(
......@@ -481,7 +531,7 @@ describe("WsTransport", () => {
481531 });
482532
483533 it("re-subscribes stream listeners after the stream exits", async () => {
484 const transport = new WsTransport("ws://localhost:3020");
534 const transport = createTransport("ws://localhost:3020");
485535 const listener = vi.fn();
486536 const onResubscribe = vi.fn();
487537
......@@ -590,7 +640,7 @@ describe("WsTransport", () => {
590640 });
591641
592642 it("does not fire onResubscribe when the first stream attempt exits before any value", async () => {
593 const transport = new WsTransport("ws://localhost:3020");
643 const transport = createTransport("ws://localhost:3020");
594644 const listener = vi.fn();
595645 const onResubscribe = vi.fn();
596646
......@@ -636,7 +686,7 @@ describe("WsTransport", () => {
636686 });
637687
638688 it("does not retry stream subscriptions after application-level failures", async () => {
639 const transport = new WsTransport("ws://localhost:3020");
689 const transport = createTransport("ws://localhost:3020");
640690 const warnSpy = vi.spyOn(console, "warn").mockImplementation(() => undefined);
641691 let attempts = 0;
642692
......@@ -675,7 +725,7 @@ describe("WsTransport", () => {
675725 });
676726
677727 it("keeps retrying stream subscriptions after transport failures", async () => {
678 const transport = new WsTransport("ws://localhost:3020");
728 const transport = createTransport("ws://localhost:3020");
679729 const warnSpy = vi.spyOn(console, "warn").mockImplementation(() => undefined);
680730 let attempts = 0;
681731
......@@ -708,7 +758,7 @@ describe("WsTransport", () => {
708758 });
709759
710760 it("logs a transport disconnect once even when multiple subscriptions fail together", async () => {
711 const transport = new WsTransport("ws://localhost:3020");
761 const transport = createTransport("ws://localhost:3020");
712762 const warnSpy = vi.spyOn(console, "warn").mockImplementation(() => undefined);
713763
714764 const unsubscribeA = transport.subscribe(
......@@ -741,7 +791,7 @@ describe("WsTransport", () => {
741791 });
742792
743793 it("streams finite request events without re-subscribing", async () => {
744 const transport = new WsTransport("ws://localhost:3020");
794 const transport = createTransport("ws://localhost:3020");
745795 const listener = vi.fn();
746796
747797 await waitFor(() => {
......@@ -857,7 +907,7 @@ describe("WsTransport", () => {
857907 exportIntervalMs: 10,
858908 });
859909
860 const transport = new WsTransport("ws://localhost:3020");
910 const transport = createTransport("ws://localhost:3020");
861911 const requestPromise = transport.request((client) => client[WS_METHODS.serverGetSettings]({}));
862912
863913 await waitFor(() => {