From 4d266cf1a2f78a1ca53cb09c5414fd74f789925c Mon Sep 17 00:00:00 2001 From: Ehab Younes Date: Mon, 19 Jan 2026 19:34:50 +0300 Subject: [PATCH 1/6] Add state machine to ReconnectingWebSocket --- src/websocket/reconnectingWebSocket.ts | 95 ++++++++++++------- .../websocket/reconnectingWebSocket.test.ts | 48 +++++++++- 2 files changed, 106 insertions(+), 37 deletions(-) diff --git a/src/websocket/reconnectingWebSocket.ts b/src/websocket/reconnectingWebSocket.ts index 5956da87..2e0e3a51 100644 --- a/src/websocket/reconnectingWebSocket.ts +++ b/src/websocket/reconnectingWebSocket.ts @@ -17,6 +17,24 @@ import type { UnidirectionalStream, } from "./eventStreamConnection"; +/** + * Connection states for the ReconnectingWebSocket state machine. + */ +export enum ConnectionState { + /** Initial state, ready to connect */ + IDLE = "IDLE", + /** Actively running connect() - WS factory in progress */ + CONNECTING = "CONNECTING", + /** Socket is open and working */ + CONNECTED = "CONNECTED", + /** Waiting for backoff timer before attempting reconnection */ + AWAITING_RETRY = "AWAITING_RETRY", + /** Temporarily paused - user must call reconnect() to resume */ + DISCONNECTED = "DISCONNECTED", + /** Permanently closed - cannot be reused */ + DISPOSED = "DISPOSED", +} + export type SocketFactory = () => Promise>; export interface ReconnectingWebSocketOptions { @@ -46,10 +64,8 @@ export class ReconnectingWebSocket< #lastRoute = "unknown"; // Cached route for logging when socket is closed #backoffMs: number; #reconnectTimeoutId: NodeJS.Timeout | null = null; - #isDisconnected = false; // Temporary pause, can be resumed via reconnect() - #isDisposed = false; // Permanent disposal, cannot be resumed - #isConnecting = false; - #pendingReconnect = false; + #state: ConnectionState = ConnectionState.IDLE; + #pendingReconnect = false; // Queue reconnect during CONNECTING state #certRefreshAttempted = false; // Tracks if cert refresh was already attempted this connection cycle readonly #onDispose?: () => void; @@ -94,11 +110,10 @@ export class ReconnectingWebSocket< } /** - * Returns true if the socket is temporarily disconnected and not attempting to reconnect. - * Use reconnect() to resume. + * Returns the current connection state. */ - get isDisconnected(): boolean { - return this.#isDisconnected; + get state(): string { + return this.#state; } /** @@ -133,14 +148,14 @@ export class ReconnectingWebSocket< * Resumes the socket if previously disconnected via disconnect(). */ reconnect(): void { - if (this.#isDisconnected) { - this.#isDisconnected = false; - this.#backoffMs = this.#options.initialBackoffMs; - this.#certRefreshAttempted = false; // User-initiated reconnect, allow retry + if (this.#state === ConnectionState.DISPOSED) { + return; } - if (this.#isDisposed) { - return; + if (this.#state === ConnectionState.DISCONNECTED) { + this.#state = ConnectionState.IDLE; + this.#backoffMs = this.#options.initialBackoffMs; + this.#certRefreshAttempted = false; // User-initiated reconnect, allow retry } if (this.#reconnectTimeoutId !== null) { @@ -148,7 +163,7 @@ export class ReconnectingWebSocket< this.#reconnectTimeoutId = null; } - if (this.#isConnecting) { + if (this.#state === ConnectionState.CONNECTING) { this.#pendingReconnect = true; return; } @@ -161,16 +176,19 @@ export class ReconnectingWebSocket< * Temporarily disconnect the socket. Can be resumed via reconnect(). */ disconnect(code?: number, reason?: string): void { - if (this.#isDisposed || this.#isDisconnected) { + if ( + this.#state === ConnectionState.DISPOSED || + this.#state === ConnectionState.DISCONNECTED + ) { return; } - this.#isDisconnected = true; + this.#state = ConnectionState.DISCONNECTED; this.clearCurrentSocket(code, reason); } close(code?: number, reason?: string): void { - if (this.#isDisposed) { + if (this.#state === ConnectionState.DISPOSED) { return; } @@ -187,11 +205,16 @@ export class ReconnectingWebSocket< } private async connect(): Promise { - if (this.#isDisposed || this.#isDisconnected || this.#isConnecting) { + // Only allow connecting from IDLE, CONNECTED (reconnect), or AWAITING_RETRY states + if ( + this.#state === ConnectionState.DISPOSED || + this.#state === ConnectionState.DISCONNECTED || + this.#state === ConnectionState.CONNECTING + ) { return; } - this.#isConnecting = true; + this.#state = ConnectionState.CONNECTING; try { // Close any existing socket before creating a new one if (this.#currentSocket) { @@ -204,18 +227,20 @@ export class ReconnectingWebSocket< const socket = await this.#socketFactory(); - // Check if disconnected/disposed while waiting for factory - if (this.#isDisposed || this.#isDisconnected) { + // Check if state changed while waiting for factory (e.g., disconnect/dispose called) + if (this.#state !== ConnectionState.CONNECTING) { socket.close(WebSocketCloseCode.NORMAL, "Cancelled during connection"); return; } this.#currentSocket = socket; this.#lastRoute = this.#route; + this.#state = ConnectionState.CONNECTED; socket.addEventListener("open", (event) => { + // Reset backoff on successful connection this.#backoffMs = this.#options.initialBackoffMs; - this.#certRefreshAttempted = false; // Reset on successful connection + this.#certRefreshAttempted = false; this.executeHandlers("open", event); }); @@ -233,7 +258,10 @@ export class ReconnectingWebSocket< }); socket.addEventListener("close", (event) => { - if (this.#isDisposed || this.#isDisconnected) { + if ( + this.#state === ConnectionState.DISPOSED || + this.#state === ConnectionState.DISCONNECTED + ) { return; } @@ -256,8 +284,6 @@ export class ReconnectingWebSocket< } catch (error) { await this.handleConnectionError(error); } finally { - this.#isConnecting = false; - if (this.#pendingReconnect) { this.#pendingReconnect = false; this.reconnect(); @@ -267,13 +293,15 @@ export class ReconnectingWebSocket< private scheduleReconnect(): void { if ( - this.#isDisposed || - this.#isDisconnected || - this.#reconnectTimeoutId !== null + this.#state === ConnectionState.DISPOSED || + this.#state === ConnectionState.DISCONNECTED || + this.#state === ConnectionState.AWAITING_RETRY ) { return; } + this.#state = ConnectionState.AWAITING_RETRY; + const jitter = this.#backoffMs * this.#options.jitterFactor * (Math.random() * 2 - 1); const delayMs = Math.max(0, this.#backoffMs + jitter); @@ -354,7 +382,10 @@ export class ReconnectingWebSocket< * otherwise schedules a reconnect. */ private async handleConnectionError(error: unknown): Promise { - if (this.#isDisposed || this.#isDisconnected) { + if ( + this.#state === ConnectionState.DISPOSED || + this.#state === ConnectionState.DISCONNECTED + ) { return; } @@ -396,11 +427,11 @@ export class ReconnectingWebSocket< } private dispose(code?: number, reason?: string): void { - if (this.#isDisposed) { + if (this.#state === ConnectionState.DISPOSED) { return; } - this.#isDisposed = true; + this.#state = ConnectionState.DISPOSED; this.clearCurrentSocket(code, reason); for (const set of Object.values(this.#eventHandlers)) { diff --git a/test/unit/websocket/reconnectingWebSocket.test.ts b/test/unit/websocket/reconnectingWebSocket.test.ts index d81f4c1a..e06f1ec8 100644 --- a/test/unit/websocket/reconnectingWebSocket.test.ts +++ b/test/unit/websocket/reconnectingWebSocket.test.ts @@ -2,6 +2,7 @@ import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; import { WebSocketCloseCode, HttpStatusCode } from "@/websocket/codes"; import { + ConnectionState, ReconnectingWebSocket, type SocketFactory, } from "@/websocket/reconnectingWebSocket"; @@ -27,13 +28,17 @@ describe("ReconnectingWebSocket", () => { const { ws, sockets } = await createReconnectingWebSocket(); sockets[0].fireOpen(); + expect(ws.state).toBe(ConnectionState.CONNECTED); + sockets[0].fireClose({ code: WebSocketCloseCode.ABNORMAL, reason: "Network error", }); + expect(ws.state).toBe(ConnectionState.AWAITING_RETRY); await vi.advanceTimersByTimeAsync(300); expect(sockets).toHaveLength(2); + expect(ws.state).toBe(ConnectionState.CONNECTED); ws.close(); }); @@ -65,7 +70,10 @@ describe("ReconnectingWebSocket", () => { const { ws, sockets } = await createReconnectingWebSocket(); sockets[0].fireOpen(); + expect(ws.state).toBe(ConnectionState.CONNECTED); + sockets[0].fireClose({ code, reason: "Unrecoverable" }); + expect(ws.state).toBe(ConnectionState.DISCONNECTED); await vi.advanceTimersByTimeAsync(10000); expect(sockets).toHaveLength(1); @@ -98,7 +106,7 @@ describe("ReconnectingWebSocket", () => { ); // Should be disconnected after unrecoverable HTTP error - expect(ws.isDisconnected).toBe(true); + expect(ws.state).toBe(ConnectionState.DISCONNECTED); // Should not retry after unrecoverable HTTP error await vi.advanceTimersByTimeAsync(10000); @@ -121,6 +129,8 @@ describe("ReconnectingWebSocket", () => { sockets[0].fireError( new Error(`Unexpected server response: ${statusCode}`), ); + expect(ws.state).toBe(ConnectionState.DISCONNECTED); + sockets[0].fireClose({ code: WebSocketCloseCode.ABNORMAL, reason: "Connection failed", @@ -179,11 +189,13 @@ describe("ReconnectingWebSocket", () => { await createBlockingReconnectingWebSocket(); ws.reconnect(); + expect(ws.state).toBe(ConnectionState.CONNECTING); ws.reconnect(); // queued expect(sockets).toHaveLength(2); // This should cancel the queued request ws.disconnect(); + expect(ws.state).toBe(ConnectionState.DISCONNECTED); failConnection(new Error("No base URL")); await Promise.resolve(); @@ -200,10 +212,12 @@ describe("ReconnectingWebSocket", () => { // Start reconnect (will block on factory promise) ws.reconnect(); + expect(ws.state).toBe(ConnectionState.CONNECTING); expect(sockets).toHaveLength(2); // Disconnect while factory is still pending ws.disconnect(); + expect(ws.state).toBe(ConnectionState.DISCONNECTED); completeConnection(); await Promise.resolve(); @@ -274,6 +288,7 @@ describe("ReconnectingWebSocket", () => { it("preserves event handlers after suspend() and reconnect()", async () => { const { ws, sockets } = await createReconnectingWebSocket(); sockets[0].fireOpen(); + expect(ws.state).toBe(ConnectionState.CONNECTED); const handler = vi.fn(); ws.addEventListener("message", handler); @@ -282,12 +297,14 @@ describe("ReconnectingWebSocket", () => { // Suspend the socket ws.disconnect(); + expect(ws.state).toBe(ConnectionState.DISCONNECTED); // Reconnect (async operation) ws.reconnect(); await Promise.resolve(); // Wait for async connect() expect(sockets).toHaveLength(2); sockets[1].fireOpen(); + expect(ws.state).toBe(ConnectionState.CONNECTED); // Handler should still work after suspend/reconnect sockets[1].fireMessage({ test: 2 }); @@ -361,19 +378,26 @@ describe("ReconnectingWebSocket", () => { ); sockets[0].fireOpen(); + expect(ws.state).toBe(ConnectionState.CONNECTED); + sockets[0].fireClose({ code: WebSocketCloseCode.PROTOCOL_ERROR, reason: "Protocol error", }); // Should suspend, not dispose - allows recovery when credentials change + expect(ws.state).toBe(ConnectionState.DISCONNECTED); expect(disposeCount).toBe(0); // Should be able to reconnect after suspension ws.reconnect(); + await Promise.resolve(); expect(sockets).toHaveLength(2); + sockets[1].fireOpen(); + expect(ws.state).toBe(ConnectionState.CONNECTED); ws.close(); + expect(ws.state).toBe(ConnectionState.DISPOSED); }); it("does not call onDispose callback during reconnection", async () => { @@ -399,6 +423,7 @@ describe("ReconnectingWebSocket", () => { const { ws, sockets, setFactoryError } = await createReconnectingWebSocketWithErrorControl(); sockets[0].fireOpen(); + expect(ws.state).toBe(ConnectionState.CONNECTED); // Trigger reconnect that will fail with 403 setFactoryError( @@ -408,6 +433,7 @@ describe("ReconnectingWebSocket", () => { await Promise.resolve(); // Socket should be suspended - no automatic reconnection + expect(ws.state).toBe(ConnectionState.DISCONNECTED); await vi.advanceTimersByTimeAsync(10000); expect(sockets).toHaveLength(1); @@ -416,17 +442,23 @@ describe("ReconnectingWebSocket", () => { ws.reconnect(); await Promise.resolve(); expect(sockets).toHaveLength(2); + sockets[1].fireOpen(); + expect(ws.state).toBe(ConnectionState.CONNECTED); ws.close(); + expect(ws.state).toBe(ConnectionState.DISPOSED); }); it("reconnect() does nothing after close()", async () => { const { ws, sockets } = await createReconnectingWebSocket(); sockets[0].fireOpen(); + expect(ws.state).toBe(ConnectionState.CONNECTED); ws.close(); - ws.reconnect(); + expect(ws.state).toBe(ConnectionState.DISPOSED); + ws.reconnect(); + expect(ws.state).toBe(ConnectionState.DISPOSED); expect(sockets).toHaveLength(1); }); }); @@ -539,7 +571,9 @@ describe("ReconnectingWebSocket", () => { ); sockets[0].fireError(new Error("ssl alert certificate_expired")); - await vi.waitFor(() => expect(ws.isDisconnected).toBe(true)); + await vi.waitFor(() => + expect(ws.state).toBe(ConnectionState.DISCONNECTED), + ); expect(sockets).toHaveLength(1); ws.close(); @@ -556,7 +590,9 @@ describe("ReconnectingWebSocket", () => { await vi.waitFor(() => expect(sockets).toHaveLength(2)); sockets[1].fireError(new Error("ssl alert certificate_expired")); - await vi.waitFor(() => expect(ws.isDisconnected).toBe(true)); + await vi.waitFor(() => + expect(ws.state).toBe(ConnectionState.DISCONNECTED), + ); expect(refreshCount).toBe(1); ws.close(); @@ -583,7 +619,9 @@ describe("ReconnectingWebSocket", () => { ); sockets[0].fireError(new Error("ssl alert unknown_ca")); - await vi.waitFor(() => expect(ws.isDisconnected).toBe(true)); + await vi.waitFor(() => + expect(ws.state).toBe(ConnectionState.DISCONNECTED), + ); expect(refreshCallback).not.toHaveBeenCalled(); ws.close(); From 93355b9f5e24750de8eceffc42cae485a51eb2d7 Mon Sep 17 00:00:00 2001 From: Ehab Younes Date: Mon, 26 Jan 2026 13:57:13 +0300 Subject: [PATCH 2/6] Handle review comments --- src/websocket/reconnectingWebSocket.ts | 25 ++++++++++++++++++++----- 1 file changed, 20 insertions(+), 5 deletions(-) diff --git a/src/websocket/reconnectingWebSocket.ts b/src/websocket/reconnectingWebSocket.ts index 2e0e3a51..67fe902a 100644 --- a/src/websocket/reconnectingWebSocket.ts +++ b/src/websocket/reconnectingWebSocket.ts @@ -205,11 +205,10 @@ export class ReconnectingWebSocket< } private async connect(): Promise { - // Only allow connecting from IDLE, CONNECTED (reconnect), or AWAITING_RETRY states if ( - this.#state === ConnectionState.DISPOSED || - this.#state === ConnectionState.DISCONNECTED || - this.#state === ConnectionState.CONNECTING + this.#state !== ConnectionState.IDLE && + this.#state !== ConnectionState.CONNECTED && + this.#state !== ConnectionState.AWAITING_RETRY ) { return; } @@ -235,9 +234,13 @@ export class ReconnectingWebSocket< this.#currentSocket = socket; this.#lastRoute = this.#route; - this.#state = ConnectionState.CONNECTED; socket.addEventListener("open", (event) => { + if (this.#currentSocket !== socket) { + return; + } + + this.#state = ConnectionState.CONNECTED; // Reset backoff on successful connection this.#backoffMs = this.#options.initialBackoffMs; this.#certRefreshAttempted = false; @@ -245,10 +248,18 @@ export class ReconnectingWebSocket< }); socket.addEventListener("message", (event) => { + if (this.#currentSocket !== socket) { + return; + } + this.executeHandlers("message", event); }); socket.addEventListener("error", (event) => { + if (this.#currentSocket !== socket) { + return; + } + this.executeHandlers("error", event); // Errors during initial connection are caught by the factory (waitForOpen). // This handler is for errors AFTER successful connection. @@ -258,6 +269,10 @@ export class ReconnectingWebSocket< }); socket.addEventListener("close", (event) => { + if (this.#currentSocket !== socket) { + return; + } + if ( this.#state === ConnectionState.DISPOSED || this.#state === ConnectionState.DISCONNECTED From 2e27a6bc1832f209945531eafec733878ff16474 Mon Sep 17 00:00:00 2001 From: Ehab Younes Date: Mon, 26 Jan 2026 15:00:51 +0300 Subject: [PATCH 3/6] Refactor WebSocket state machine to use action-based dispatch Replace direct state assignments with a pure reducer pattern: - Add StateAction type and reduceState() pure function - Add #dispatch() method that validates transitions - Remove #pendingReconnect flag - reconnect() now restarts immediately --- src/websocket/reconnectingWebSocket.ts | 136 +++++++++++++----- test/unit/api/coderApi.test.ts | 3 + .../websocket/reconnectingWebSocket.test.ts | 20 +-- 3 files changed, 111 insertions(+), 48 deletions(-) diff --git a/src/websocket/reconnectingWebSocket.ts b/src/websocket/reconnectingWebSocket.ts index 67fe902a..1f2fb520 100644 --- a/src/websocket/reconnectingWebSocket.ts +++ b/src/websocket/reconnectingWebSocket.ts @@ -35,6 +35,73 @@ export enum ConnectionState { DISPOSED = "DISPOSED", } +/** + * Actions that trigger state transitions. + */ +type StateAction = + | { readonly type: "CONNECT" } + | { readonly type: "OPEN" } + | { readonly type: "SCHEDULE_RETRY" } + | { readonly type: "DISCONNECT" } + | { readonly type: "DISPOSE" }; + +/** + * Pure reducer function for state transitions. + */ +function reduceState( + state: ConnectionState, + action: StateAction, +): ConnectionState { + switch (action.type) { + case "CONNECT": + switch (state) { + case ConnectionState.IDLE: + case ConnectionState.CONNECTED: + case ConnectionState.AWAITING_RETRY: + case ConnectionState.DISCONNECTED: + return ConnectionState.CONNECTING; + default: + return state; + } + + case "OPEN": + switch (state) { + case ConnectionState.CONNECTING: + return ConnectionState.CONNECTED; + default: + return state; + } + + case "SCHEDULE_RETRY": + switch (state) { + case ConnectionState.CONNECTING: + case ConnectionState.CONNECTED: + return ConnectionState.AWAITING_RETRY; + default: + return state; + } + + case "DISCONNECT": + switch (state) { + case ConnectionState.IDLE: + case ConnectionState.CONNECTING: + case ConnectionState.CONNECTED: + case ConnectionState.AWAITING_RETRY: + return ConnectionState.DISCONNECTED; + default: + return state; + } + + case "DISPOSE": + switch (state) { + case ConnectionState.DISPOSED: + return state; + default: + return ConnectionState.DISPOSED; + } + } +} + export type SocketFactory = () => Promise>; export interface ReconnectingWebSocketOptions { @@ -65,10 +132,28 @@ export class ReconnectingWebSocket< #backoffMs: number; #reconnectTimeoutId: NodeJS.Timeout | null = null; #state: ConnectionState = ConnectionState.IDLE; - #pendingReconnect = false; // Queue reconnect during CONNECTING state #certRefreshAttempted = false; // Tracks if cert refresh was already attempted this connection cycle readonly #onDispose?: () => void; + /** + * Dispatch an action to transition state. Returns true if transition is allowed. + */ + #dispatch(action: StateAction): boolean { + const newState = reduceState(this.#state, action); + if (newState === this.#state) { + // Allow CONNECT from CONNECTING as a "restart" operation + if ( + action.type === "CONNECT" && + this.#state === ConnectionState.CONNECTING + ) { + return true; + } + return false; + } + this.#state = newState; + return true; + } + private constructor( socketFactory: SocketFactory, logger: Logger, @@ -153,7 +238,6 @@ export class ReconnectingWebSocket< } if (this.#state === ConnectionState.DISCONNECTED) { - this.#state = ConnectionState.IDLE; this.#backoffMs = this.#options.initialBackoffMs; this.#certRefreshAttempted = false; // User-initiated reconnect, allow retry } @@ -163,11 +247,6 @@ export class ReconnectingWebSocket< this.#reconnectTimeoutId = null; } - if (this.#state === ConnectionState.CONNECTING) { - this.#pendingReconnect = true; - return; - } - // connect() handles all errors internally void this.connect(); } @@ -176,14 +255,9 @@ export class ReconnectingWebSocket< * Temporarily disconnect the socket. Can be resumed via reconnect(). */ disconnect(code?: number, reason?: string): void { - if ( - this.#state === ConnectionState.DISPOSED || - this.#state === ConnectionState.DISCONNECTED - ) { + if (!this.#dispatch({ type: "DISCONNECT" })) { return; } - - this.#state = ConnectionState.DISCONNECTED; this.clearCurrentSocket(code, reason); } @@ -205,15 +279,9 @@ export class ReconnectingWebSocket< } private async connect(): Promise { - if ( - this.#state !== ConnectionState.IDLE && - this.#state !== ConnectionState.CONNECTED && - this.#state !== ConnectionState.AWAITING_RETRY - ) { + if (!this.#dispatch({ type: "CONNECT" })) { return; } - - this.#state = ConnectionState.CONNECTING; try { // Close any existing socket before creating a new one if (this.#currentSocket) { @@ -240,7 +308,9 @@ export class ReconnectingWebSocket< return; } - this.#state = ConnectionState.CONNECTED; + if (!this.#dispatch({ type: "OPEN" })) { + return; + } // Reset backoff on successful connection this.#backoffMs = this.#options.initialBackoffMs; this.#certRefreshAttempted = false; @@ -298,25 +368,14 @@ export class ReconnectingWebSocket< }); } catch (error) { await this.handleConnectionError(error); - } finally { - if (this.#pendingReconnect) { - this.#pendingReconnect = false; - this.reconnect(); - } } } private scheduleReconnect(): void { - if ( - this.#state === ConnectionState.DISPOSED || - this.#state === ConnectionState.DISCONNECTED || - this.#state === ConnectionState.AWAITING_RETRY - ) { + if (!this.#dispatch({ type: "SCHEDULE_RETRY" })) { return; } - this.#state = ConnectionState.AWAITING_RETRY; - const jitter = this.#backoffMs * this.#options.jitterFactor * (Math.random() * 2 - 1); const delayMs = Math.max(0, this.#backoffMs + jitter); @@ -401,6 +460,10 @@ export class ReconnectingWebSocket< this.#state === ConnectionState.DISPOSED || this.#state === ConnectionState.DISCONNECTED ) { + this.#logger.debug( + `Ignoring connection error in ${this.#state} state for ${this.#route}`, + error, + ); return; } @@ -442,11 +505,9 @@ export class ReconnectingWebSocket< } private dispose(code?: number, reason?: string): void { - if (this.#state === ConnectionState.DISPOSED) { + if (!this.#dispatch({ type: "DISPOSE" })) { return; } - - this.#state = ConnectionState.DISPOSED; this.clearCurrentSocket(code, reason); for (const set of Object.values(this.#eventHandlers)) { @@ -457,9 +518,6 @@ export class ReconnectingWebSocket< } private clearCurrentSocket(code?: number, reason?: string): void { - // Clear pending reconnect to prevent resume - this.#pendingReconnect = false; - if (this.#reconnectTimeoutId !== null) { clearTimeout(this.#reconnectTimeoutId); this.#reconnectTimeoutId = null; diff --git a/test/unit/api/coderApi.test.ts b/test/unit/api/coderApi.test.ts index 6f8dfd45..c97e3f8a 100644 --- a/test/unit/api/coderApi.test.ts +++ b/test/unit/api/coderApi.test.ts @@ -801,6 +801,7 @@ describe("CoderApi", () => { const sockets = setupAutoOpeningWebSocket(); api = createApi(CODER_URL, AXIOS_TOKEN); await api.watchAgentMetadata(AGENT_ID); + await tick(); // Wait for open event to fire (socket becomes CONNECTED) mockConfig.set("coder.insecure", true); await tick(); @@ -820,6 +821,7 @@ describe("CoderApi", () => { const sockets = setupAutoOpeningWebSocket(); api = createApi(CODER_URL, AXIOS_TOKEN); await api.watchAgentMetadata(AGENT_ID); + await tick(); // Wait for open event to fire (socket becomes CONNECTED) mockConfig.set(key, value); await tick(); @@ -833,6 +835,7 @@ describe("CoderApi", () => { const sockets = setupAutoOpeningWebSocket(); api = createApi(CODER_URL, AXIOS_TOKEN); await api.watchAgentMetadata(AGENT_ID); + await tick(); // Wait for open event to fire (socket becomes CONNECTED) api.dispose(); mockConfig.set("coder.insecure", true); diff --git a/test/unit/websocket/reconnectingWebSocket.test.ts b/test/unit/websocket/reconnectingWebSocket.test.ts index e06f1ec8..c449623c 100644 --- a/test/unit/websocket/reconnectingWebSocket.test.ts +++ b/test/unit/websocket/reconnectingWebSocket.test.ts @@ -38,6 +38,7 @@ describe("ReconnectingWebSocket", () => { await vi.advanceTimersByTimeAsync(300); expect(sockets).toHaveLength(2); + sockets[1].fireOpen(); expect(ws.state).toBe(ConnectionState.CONNECTED); ws.close(); @@ -164,7 +165,7 @@ describe("ReconnectingWebSocket", () => { ws.close(); }); - it("queues reconnect() calls made during connection", async () => { + it("reconnect() during CONNECTING immediately restarts connection", async () => { const { ws, sockets, completeConnection } = await createBlockingReconnectingWebSocket(); @@ -172,33 +173,33 @@ describe("ReconnectingWebSocket", () => { ws.reconnect(); expect(sockets).toHaveLength(2); // Call reconnect again while first reconnect is in progress + // This immediately restarts (creates a new socket) ws.reconnect(); - // Still only 2 sockets (queued reconnect hasn't started) - expect(sockets).toHaveLength(2); + expect(sockets).toHaveLength(3); + // Complete the third socket's connection completeConnection(); await Promise.resolve(); - // Now queued reconnect should have executed, creating third socket - expect(sockets).toHaveLength(3); + expect(ws.state).toBe(ConnectionState.CONNECTED); ws.close(); }); - it("suspend() cancels pending reconnect queued during connection", async () => { + it("disconnect() cancels in-progress reconnect and prevents new connections", async () => { const { ws, sockets, failConnection } = await createBlockingReconnectingWebSocket(); ws.reconnect(); expect(ws.state).toBe(ConnectionState.CONNECTING); - ws.reconnect(); // queued expect(sockets).toHaveLength(2); - // This should cancel the queued request + // Disconnect while reconnect is in progress ws.disconnect(); expect(ws.state).toBe(ConnectionState.DISCONNECTED); failConnection(new Error("No base URL")); await Promise.resolve(); + // No new socket should be created after disconnect expect(sockets).toHaveLength(2); await vi.advanceTimersByTimeAsync(10000); expect(sockets).toHaveLength(2); @@ -791,7 +792,8 @@ async function createBlockingReconnectingWebSocket(): Promise<{ completeConnection: () => { const socket = sockets.at(-1)!; pendingResolve?.(socket); - socket.fireOpen(); + // Fire open after microtask so event listener is attached + queueMicrotask(() => socket.fireOpen()); }, failConnection: (error: Error) => pendingReject?.(error), }; From 69df29136d74cf34f15f5456e02bd4ddf3426167 Mon Sep 17 00:00:00 2001 From: Ehab Younes Date: Mon, 26 Jan 2026 13:38:48 +0300 Subject: [PATCH 4/6] Only trigger reconnection if the WS is not connected --- src/api/coderApi.ts | 19 ++--- src/websocket/reconnectingWebSocket.ts | 2 +- test/unit/api/coderApi.test.ts | 70 ++++++++++++++----- .../websocket/reconnectingWebSocket.test.ts | 49 +++++++++++++ 4 files changed, 113 insertions(+), 27 deletions(-) diff --git a/src/api/coderApi.ts b/src/api/coderApi.ts index 1e5a5e85..166b2012 100644 --- a/src/api/coderApi.ts +++ b/src/api/coderApi.ts @@ -46,6 +46,7 @@ import { type OneWayWebSocketInit, } from "../websocket/oneWayWebSocket"; import { + ConnectionState, ReconnectingWebSocket, type SocketFactory, } from "../websocket/reconnectingWebSocket"; @@ -164,7 +165,8 @@ export class CoderApi extends Api implements vscode.Disposable { /** * Watch for configuration changes that affect WebSocket connections. - * When any watched setting changes, all active WebSockets are reconnected. + * Only reconnects DISCONNECTED sockets since they require an explicit reconnect() call. + * Other states will pick up settings naturally. */ private watchConfigChanges(): vscode.Disposable { const settings = webSocketConfigSettings.map((setting) => ({ @@ -172,13 +174,14 @@ export class CoderApi extends Api implements vscode.Disposable { getValue: () => vscode.workspace.getConfiguration().get(setting), })); return watchConfigurationChanges(settings, () => { - if (this.reconnectingSockets.size > 0) { - this.output.info( - `Configuration changed, reconnecting ${this.reconnectingSockets.size} WebSocket(s)`, - ); - for (const socket of this.reconnectingSockets) { - socket.reconnect(); - } + const socketsToReconnect = [...this.reconnectingSockets].filter( + (socket) => socket.state === ConnectionState.DISCONNECTED, + ); + this.output.debug( + `Configuration changed, ${socketsToReconnect.length}/${this.reconnectingSockets.size} socket(s) in DISCONNECTED state`, + ); + for (const socket of socketsToReconnect) { + socket.reconnect(); } }); } diff --git a/src/websocket/reconnectingWebSocket.ts b/src/websocket/reconnectingWebSocket.ts index 1f2fb520..ec2887d6 100644 --- a/src/websocket/reconnectingWebSocket.ts +++ b/src/websocket/reconnectingWebSocket.ts @@ -197,7 +197,7 @@ export class ReconnectingWebSocket< /** * Returns the current connection state. */ - get state(): string { + get state(): ConnectionState { return this.#state; } diff --git a/test/unit/api/coderApi.test.ts b/test/unit/api/coderApi.test.ts index c97e3f8a..f4377aae 100644 --- a/test/unit/api/coderApi.test.ts +++ b/test/unit/api/coderApi.test.ts @@ -591,9 +591,11 @@ describe("CoderApi", () => { const setupAutoOpeningWebSocket = () => { const sockets: Array> = []; + const handlers: Record void> = {}; vi.mocked(Ws).mockImplementation(function (url: string | URL) { const mockWs = createMockWebSocket(String(url), { - on: vi.fn((event, handler) => { + on: vi.fn((event: string, handler: (...args: unknown[]) => void) => { + handlers[event] = handler; if (event === "open") { setImmediate(() => handler()); } @@ -603,12 +605,12 @@ describe("CoderApi", () => { sockets.push(mockWs); return mockWs as Ws; }); - return sockets; + return { sockets, handlers }; }; describe("Reconnection on Host/Token Changes", () => { it("triggers reconnection when session token changes", async () => { - const sockets = setupAutoOpeningWebSocket(); + const { sockets } = setupAutoOpeningWebSocket(); api = createApi(CODER_URL, AXIOS_TOKEN); await api.watchAgentMetadata(AGENT_ID); @@ -623,7 +625,7 @@ describe("CoderApi", () => { }); it("triggers reconnection when host changes", async () => { - const sockets = setupAutoOpeningWebSocket(); + const { sockets } = setupAutoOpeningWebSocket(); api = createApi(CODER_URL, AXIOS_TOKEN); const wsWrap = await api.watchAgentMetadata(AGENT_ID); expect(wsWrap.url).toContain(CODER_URL.replace("http", "ws")); @@ -642,7 +644,7 @@ describe("CoderApi", () => { }); it("does not reconnect when token or host are unchanged", async () => { - const sockets = setupAutoOpeningWebSocket(); + const { sockets } = setupAutoOpeningWebSocket(); api = createApi(CODER_URL, AXIOS_TOKEN); await api.watchAgentMetadata(AGENT_ID); @@ -655,7 +657,7 @@ describe("CoderApi", () => { }); it("suspends sockets when host is set to empty string (logout)", async () => { - const sockets = setupAutoOpeningWebSocket(); + const { sockets } = setupAutoOpeningWebSocket(); api = createApi(CODER_URL, AXIOS_TOKEN); await api.watchAgentMetadata(AGENT_ID); @@ -668,7 +670,7 @@ describe("CoderApi", () => { }); it("does not reconnect when setting token after clearing host", async () => { - const sockets = setupAutoOpeningWebSocket(); + const { sockets } = setupAutoOpeningWebSocket(); api = createApi(CODER_URL, AXIOS_TOKEN); await api.watchAgentMetadata(AGENT_ID); @@ -682,7 +684,7 @@ describe("CoderApi", () => { }); it("setCredentials sets both host and token together", async () => { - const sockets = setupAutoOpeningWebSocket(); + const { sockets } = setupAutoOpeningWebSocket(); api = createApi(CODER_URL, AXIOS_TOKEN); await api.watchAgentMetadata(AGENT_ID); @@ -695,7 +697,7 @@ describe("CoderApi", () => { }); it("setCredentials suspends when host is cleared", async () => { - const sockets = setupAutoOpeningWebSocket(); + const { sockets } = setupAutoOpeningWebSocket(); api = createApi(CODER_URL, AXIOS_TOKEN); await api.watchAgentMetadata(AGENT_ID); @@ -796,9 +798,9 @@ describe("CoderApi", () => { describe("Configuration Change Reconnection", () => { const tick = () => new Promise((resolve) => setImmediate(resolve)); - it("reconnects sockets when watched config value changes", async () => { + it("does not reconnect connected sockets when config value changes", async () => { mockConfig.set("coder.insecure", false); - const sockets = setupAutoOpeningWebSocket(); + const { sockets } = setupAutoOpeningWebSocket(); api = createApi(CODER_URL, AXIOS_TOKEN); await api.watchAgentMetadata(AGENT_ID); await tick(); // Wait for open event to fire (socket becomes CONNECTED) @@ -806,11 +808,8 @@ describe("CoderApi", () => { mockConfig.set("coder.insecure", true); await tick(); - expect(sockets[0].close).toHaveBeenCalledWith( - 1000, - "Replacing connection", - ); - expect(sockets).toHaveLength(2); + expect(sockets[0].close).not.toHaveBeenCalled(); + expect(sockets).toHaveLength(1); }); it.each([ @@ -818,7 +817,7 @@ describe("CoderApi", () => { ["unrelated setting", "unrelated.setting", "new-value"], ])("does not reconnect for %s", async (_desc, key, value) => { mockConfig.set("coder.insecure", false); - const sockets = setupAutoOpeningWebSocket(); + const { sockets } = setupAutoOpeningWebSocket(); api = createApi(CODER_URL, AXIOS_TOKEN); await api.watchAgentMetadata(AGENT_ID); await tick(); // Wait for open event to fire (socket becomes CONNECTED) @@ -832,7 +831,7 @@ describe("CoderApi", () => { it("stops watching after dispose", async () => { mockConfig.set("coder.insecure", false); - const sockets = setupAutoOpeningWebSocket(); + const { sockets } = setupAutoOpeningWebSocket(); api = createApi(CODER_URL, AXIOS_TOKEN); await api.watchAgentMetadata(AGENT_ID); await tick(); // Wait for open event to fire (socket becomes CONNECTED) @@ -843,6 +842,41 @@ describe("CoderApi", () => { expect(sockets).toHaveLength(1); }); + + it("does not reconnect sockets in AWAITING_RETRY state when config changes", async () => { + mockConfig.set("coder.insecure", false); + const { sockets, handlers } = setupAutoOpeningWebSocket(); + api = createApi(CODER_URL, AXIOS_TOKEN); + await api.watchAgentMetadata(AGENT_ID); + + // Trigger close with abnormal code to put socket in AWAITING_RETRY + handlers["close"]?.({ code: 1006, reason: "Abnormal closure" }); + await tick(); + + mockConfig.set("coder.insecure", true); + await tick(); + + // AWAITING_RETRY will naturally retry, so no config-triggered reconnect needed + expect(sockets).toHaveLength(1); + }); + + it("reconnects sockets in DISCONNECTED state when config changes", async () => { + mockConfig.set("coder.insecure", false); + const { sockets, handlers } = setupAutoOpeningWebSocket(); + api = createApi(CODER_URL, AXIOS_TOKEN); + await api.watchAgentMetadata(AGENT_ID); + await tick(); + + // Trigger close with unrecoverable code to put socket in DISCONNECTED + handlers["close"]?.({ code: 1002, reason: "Protocol error" }); + await tick(); + + mockConfig.set("coder.insecure", true); + await tick(); + + // Only DISCONNECTED sockets get reconnected by config changes + expect(sockets).toHaveLength(2); + }); }); }); diff --git a/test/unit/websocket/reconnectingWebSocket.test.ts b/test/unit/websocket/reconnectingWebSocket.test.ts index c449623c..8e028e7f 100644 --- a/test/unit/websocket/reconnectingWebSocket.test.ts +++ b/test/unit/websocket/reconnectingWebSocket.test.ts @@ -510,6 +510,55 @@ describe("ReconnectingWebSocket", () => { }); describe("Error Handling", () => { + it("error event when CONNECTED schedules retry", async () => { + const { ws, sockets } = await createReconnectingWebSocket(); + + sockets[0].fireOpen(); + expect(ws.state).toBe(ConnectionState.CONNECTED); + + sockets[0].fireError(new Error("Connection lost")); + expect(ws.state).toBe(ConnectionState.AWAITING_RETRY); + + await vi.advanceTimersByTimeAsync(300); + expect(sockets).toHaveLength(2); + + ws.close(); + }); + + it("error event when DISCONNECTED is ignored", async () => { + const { ws, sockets } = await createReconnectingWebSocket(); + + sockets[0].fireOpen(); + expect(ws.state).toBe(ConnectionState.CONNECTED); + + ws.disconnect(); + expect(ws.state).toBe(ConnectionState.DISCONNECTED); + + // Error after disconnect should be ignored + sockets[0].fireError(new Error("Connection lost")); + expect(ws.state).toBe(ConnectionState.DISCONNECTED); + + // No reconnection should be scheduled + await vi.advanceTimersByTimeAsync(10000); + expect(sockets).toHaveLength(1); + + ws.close(); + }); + + it("multiple errors while AWAITING_RETRY only creates one reconnection", async () => { + const { ws, sockets } = await createReconnectingWebSocket(); + sockets[0].fireOpen(); + + sockets[0].fireError(new Error("First error")); + sockets[0].fireError(new Error("Second error")); + sockets[0].fireError(new Error("Third error")); + + await vi.advanceTimersByTimeAsync(300); + expect(sockets).toHaveLength(2); + + ws.close(); + }); + it("schedules retry when socket factory throws error", async () => { const sockets: MockSocket[] = []; let shouldFail = false; From 6b4fe99d2fc15c1a85da64117f58a75eb492cfdc Mon Sep 17 00:00:00 2001 From: Ehab Younes Date: Mon, 26 Jan 2026 21:54:34 +0300 Subject: [PATCH 5/6] Add more logging and some cleanup --- src/api/coderApi.ts | 13 ++++++++----- src/inbox.ts | 5 ----- src/websocket/reconnectingWebSocket.ts | 23 +++++++++-------------- src/workspace/workspacesProvider.ts | 4 ---- 4 files changed, 17 insertions(+), 28 deletions(-) diff --git a/src/api/coderApi.ts b/src/api/coderApi.ts index 166b2012..a04bfeeb 100644 --- a/src/api/coderApi.ts +++ b/src/api/coderApi.ts @@ -177,11 +177,14 @@ export class CoderApi extends Api implements vscode.Disposable { const socketsToReconnect = [...this.reconnectingSockets].filter( (socket) => socket.state === ConnectionState.DISCONNECTED, ); - this.output.debug( - `Configuration changed, ${socketsToReconnect.length}/${this.reconnectingSockets.size} socket(s) in DISCONNECTED state`, - ); - for (const socket of socketsToReconnect) { - socket.reconnect(); + if (socketsToReconnect.length) { + this.output.debug( + `Configuration changed, ${socketsToReconnect.length}/${this.reconnectingSockets.size} socket(s) in DISCONNECTED state`, + ); + for (const socket of socketsToReconnect) { + this.output.debug(`Reconnecting WebSocket: ${socket.url}`); + socket.reconnect(); + } } }); } diff --git a/src/inbox.ts b/src/inbox.ts index 59b9ae0b..ed5956c8 100644 --- a/src/inbox.ts +++ b/src/inbox.ts @@ -50,11 +50,6 @@ export class Inbox implements vscode.Disposable { logger.info("Listening to Coder Inbox"); }); - socket.addEventListener("error", () => { - // Errors are already logged internally - inbox.dispose(); - }); - socket.addEventListener("message", (data) => { if (data.parseError) { logger.error("Failed to parse inbox message", data.parseError); diff --git a/src/websocket/reconnectingWebSocket.ts b/src/websocket/reconnectingWebSocket.ts index ec2887d6..bd33a247 100644 --- a/src/websocket/reconnectingWebSocket.ts +++ b/src/websocket/reconnectingWebSocket.ts @@ -93,12 +93,7 @@ function reduceState( } case "DISPOSE": - switch (state) { - case ConnectionState.DISPOSED: - return state; - default: - return ConnectionState.DISPOSED; - } + return ConnectionState.DISPOSED; } } @@ -172,7 +167,7 @@ export class ReconnectingWebSocket< this.#onDispose = onDispose; } - static async create( + public static async create( socketFactory: SocketFactory, logger: Logger, options: ReconnectingWebSocketOptions, @@ -190,14 +185,14 @@ export class ReconnectingWebSocket< return instance; } - get url(): string { + public get url(): string { return this.#currentSocket?.url ?? ""; } /** * Returns the current connection state. */ - get state(): ConnectionState { + public get state(): ConnectionState { return this.#state; } @@ -214,14 +209,14 @@ export class ReconnectingWebSocket< return url.pathname + url.search; } - addEventListener( + public addEventListener( event: TEvent, callback: EventHandler, ): void { this.#eventHandlers[event].add(callback); } - removeEventListener( + public removeEventListener( event: TEvent, callback: EventHandler, ): void { @@ -232,7 +227,7 @@ export class ReconnectingWebSocket< * Force an immediate reconnection attempt. * Resumes the socket if previously disconnected via disconnect(). */ - reconnect(): void { + public reconnect(): void { if (this.#state === ConnectionState.DISPOSED) { return; } @@ -254,14 +249,14 @@ export class ReconnectingWebSocket< /** * Temporarily disconnect the socket. Can be resumed via reconnect(). */ - disconnect(code?: number, reason?: string): void { + public disconnect(code?: number, reason?: string): void { if (!this.#dispatch({ type: "DISCONNECT" })) { return; } this.clearCurrentSocket(code, reason); } - close(code?: number, reason?: string): void { + public close(code?: number, reason?: string): void { if (this.#state === ConnectionState.DISPOSED) { return; } diff --git a/src/workspace/workspacesProvider.ts b/src/workspace/workspacesProvider.ts index ff462e75..d03851a7 100644 --- a/src/workspace/workspacesProvider.ts +++ b/src/workspace/workspacesProvider.ts @@ -101,10 +101,6 @@ export class WorkspaceProvider * logged in or the query fails. */ private async fetch(): Promise { - this.logger.debug( - `Fetching workspaces: ${this.getWorkspacesQuery || "no filter"}...`, - ); - // If there is no URL configured, assume we are logged out. const url = this.client.getAxiosInstance().defaults.baseURL; if (!url) { From 59d414763a47c8f39d7f609c5a2faeca4294c073 Mon Sep 17 00:00:00 2001 From: Ehab Younes Date: Tue, 27 Jan 2026 00:27:25 +0300 Subject: [PATCH 6/6] Add changelog entry --- CHANGELOG.md | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 40ff8a40..aeaaf8b9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,11 @@ ## Unreleased +### Changed + +- WebSocket connections are now more robust and reconnect less frequently, only when truly + necessary, reducing unnecessary disconnections and improving stability. + ## [v1.12.1](https://github.com/coder/vscode-coder/releases/tag/v1.12.1) 2026-01-23 ### Fixed