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

Author: EhabY

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<TData> = () => Promise<UnidirectionalStream<TData>>;
 
 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<TData>,
 		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<void> {
-		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;

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);

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,41 +165,41 @@ 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();
 
 			// Start first reconnect (will block on factory promise)
 			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),
 	};