Add automatic reconnection for WebSocket connections (#654)

Implement ReconnectingWebSocket class to automatically recover from
network failures when communicating with Coder deployments. Uses
exponential backoff with jitter and distinguishes between recoverable
errors (network issues) and unrecoverable errors (auth failures).

- Auto-reconnect on abnormal closures and network failures
- Stop on unrecoverable errors (HTTP 403/410/426, WS 1002/1003)
- Reconnect when session token or host changes
- Event handlers persist across reconnections

Closes #595

Author: EhabY

CHANGELOG.md

@@ -8,6 +8,11 @@
   Other windows now wait and display real-time progress instead of attempting concurrent downloads,
   preventing corruption and failures.
 
+### Changed
+
+- WebSocket connections now automatically reconnect on network failures, improving reliability when
+  communicating with Coder deployments.
+
 ## [v1.11.4](https://github.com/coder/vscode-coder/releases/tag/v1.11.4) 2025-11-20
 
 ### Fixed

src/api/agentMetadataHelper.ts

@@ -53,7 +53,11 @@ export async function createAgentMetadataWatcher(
 				event.parsedMessage.data,
 			);
 
-			// Overwrite metadata if it changed.
+			if (watcher.error !== undefined) {
+				watcher.error = undefined;
+				onChange.fire(null);
+			}
+
 			if (JSON.stringify(watcher.metadata) !== JSON.stringify(metadata)) {
 				watcher.metadata = metadata;
 				onChange.fire(null);

src/api/coderApi.ts

@@ -14,7 +14,7 @@ import {
 	type WorkspaceAgentLog,
 } from "coder/site/src/api/typesGenerated";
 import * as vscode from "vscode";
-import { type ClientOptions, type CloseEvent, type ErrorEvent } from "ws";
+import { type ClientOptions } from "ws";
 
 import { CertificateError } from "../error";
 import { getHeaderCommand, getHeaders } from "../headers";
@@ -31,11 +31,20 @@ import {
 	HttpClientLogLevel,
 } from "../logging/types";
 import { sizeOf } from "../logging/utils";
-import { type UnidirectionalStream } from "../websocket/eventStreamConnection";
+import { HttpStatusCode } from "../websocket/codes";
+import {
+	type UnidirectionalStream,
+	type CloseEvent,
+	type ErrorEvent,
+} from "../websocket/eventStreamConnection";
 import {
 	OneWayWebSocket,
 	type OneWayWebSocketInit,
 } from "../websocket/oneWayWebSocket";
+import {
+	ReconnectingWebSocket,
+	type SocketFactory,
+} from "../websocket/reconnectingWebSocket";
 import { SseConnection } from "../websocket/sseConnection";
 
 import { createHttpAgent } from "./utils";
@@ -47,6 +56,10 @@ const coderSessionTokenHeader = "Coder-Session-Token";
  * and WebSocket methods for real-time functionality.
  */
 export class CoderApi extends Api {
+	private readonly reconnectingSockets = new Set<
+		ReconnectingWebSocket<unknown>
+	>();
+
 	private constructor(private readonly output: Logger) {
 		super();
 	}
@@ -66,10 +79,34 @@ export class CoderApi extends Api {
 			client.setSessionToken(token);
 		}
 
-		setupInterceptors(client, baseUrl, output);
+		setupInterceptors(client, output);
 		return client;
 	}
 
+	setSessionToken = (token: string): void => {
+		const defaultHeaders = this.getAxiosInstance().defaults.headers.common;
+		const currentToken = defaultHeaders[coderSessionTokenHeader];
+		defaultHeaders[coderSessionTokenHeader] = token;
+
+		if (currentToken !== token) {
+			for (const socket of this.reconnectingSockets) {
+				socket.reconnect();
+			}
+		}
+	};
+
+	setHost = (host: string | undefined): void => {
+		const defaults = this.getAxiosInstance().defaults;
+		const currentHost = defaults.baseURL;
+		defaults.baseURL = host;
+
+		if (currentHost !== host) {
+			for (const socket of this.reconnectingSockets) {
+				socket.reconnect();
+			}
+		}
+	};
+
 	watchInboxNotifications = async (
 		watchTemplates: string[],
 		watchTargets: string[],
@@ -83,6 +120,7 @@ export class CoderApi extends Api {
 				targets: watchTargets.join(","),
 			},
 			options,
+			enableRetry: true,
 		});
 	};
 
@@ -91,6 +129,7 @@ export class CoderApi extends Api {
 			apiRoute: `/api/v2/workspaces/${workspace.id}/watch-ws`,
 			fallbackApiRoute: `/api/v2/workspaces/${workspace.id}/watch`,
 			options,
+			enableRetry: true,
 		});
 	};
 
@@ -102,6 +141,7 @@ export class CoderApi extends Api {
 			apiRoute: `/api/v2/workspaceagents/${agentId}/watch-metadata-ws`,
 			fallbackApiRoute: `/api/v2/workspaceagents/${agentId}/watch-metadata`,
 			options,
+			enableRetry: true,
 		});
 	};
 
@@ -148,53 +188,78 @@ export class CoderApi extends Api {
 	}
 
 	private async createWebSocket<TData = unknown>(
-		configs: Omit<OneWayWebSocketInit, "location">,
-	) {
-		const baseUrlRaw = this.getAxiosInstance().defaults.baseURL;
-		if (!baseUrlRaw) {
-			throw new Error("No base URL set on REST client");
-		}
+		configs: Omit<OneWayWebSocketInit, "location"> & { enableRetry?: boolean },
+	): Promise<UnidirectionalStream<TData>> {
+		const { enableRetry, ...socketConfigs } = configs;
+
+		const socketFactory: SocketFactory<TData> = async () => {
+			const baseUrlRaw = this.getAxiosInstance().defaults.baseURL;
+			if (!baseUrlRaw) {
+				throw new Error("No base URL set on REST client");
+			}
+
+			const baseUrl = new URL(baseUrlRaw);
+			const token = this.getAxiosInstance().defaults.headers.common[
+				coderSessionTokenHeader
+			] as string | undefined;
+
+			const headersFromCommand = await getHeaders(
+				baseUrlRaw,
+				getHeaderCommand(vscode.workspace.getConfiguration()),
+				this.output,
+			);
 
-		const baseUrl = new URL(baseUrlRaw);
-		const token = this.getAxiosInstance().defaults.headers.common[
-			coderSessionTokenHeader
-		] as string | undefined;
+			const httpAgent = await createHttpAgent(
+				vscode.workspace.getConfiguration(),
+			);
 
-		const headersFromCommand = await getHeaders(
-			baseUrlRaw,
-			getHeaderCommand(vscode.workspace.getConfiguration()),
-			this.output,
-		);
+			/**
+			 * Similar to the REST client, we want to prioritize headers in this order (highest to lowest):
+			 * 1. Headers from the header command
+			 * 2. Any headers passed directly to this function
+			 * 3. Coder session token from the Api client (if set)
+			 */
+			const headers = {
+				...(token ? { [coderSessionTokenHeader]: token } : {}),
+				...configs.options?.headers,
+				...headersFromCommand,
+			};
 
-		const httpAgent = await createHttpAgent(
-			vscode.workspace.getConfiguration(),
-		);
+			const webSocket = new OneWayWebSocket<TData>({
+				location: baseUrl,
+				...socketConfigs,
+				options: {
+					...configs.options,
+					agent: httpAgent,
+					followRedirects: true,
+					headers,
+				},
+			});
 
-		/**
-		 * Similar to the REST client, we want to prioritize headers in this order (highest to lowest):
-		 * 1. Headers from the header command
-		 * 2. Any headers passed directly to this function
-		 * 3. Coder session token from the Api client (if set)
-		 */
-		const headers = {
-			...(token ? { [coderSessionTokenHeader]: token } : {}),
-			...configs.options?.headers,
-			...headersFromCommand,
+			this.attachStreamLogger(webSocket);
+			return webSocket;
 		};
 
-		const webSocket = new OneWayWebSocket<TData>({
-			location: baseUrl,
-			...configs,
-			options: {
-				...configs.options,
-				agent: httpAgent,
-				followRedirects: true,
-				headers,
-			},
-		});
+		if (enableRetry) {
+			const reconnectingSocket = await ReconnectingWebSocket.create<TData>(
+				socketFactory,
+				this.output,
+				configs.apiRoute,
+				undefined,
+				() =>
+					this.reconnectingSockets.delete(
+						reconnectingSocket as ReconnectingWebSocket<unknown>,
+					),
+			);
+
+			this.reconnectingSockets.add(
+				reconnectingSocket as ReconnectingWebSocket<unknown>,
+			);
 
-		this.attachStreamLogger(webSocket);
-		return webSocket;
+			return reconnectingSocket;
+		} else {
+			return socketFactory();
+		}
 	}
 
 	private attachStreamLogger<TData>(
@@ -230,13 +295,15 @@ export class CoderApi extends Api {
 		fallbackApiRoute: string;
 		searchParams?: Record<string, string> | URLSearchParams;
 		options?: ClientOptions;
+		enableRetry?: boolean;
 	}): Promise<UnidirectionalStream<TData>> {
-		let webSocket: OneWayWebSocket<TData>;
+		let webSocket: UnidirectionalStream<TData>;
 		try {
 			webSocket = await this.createWebSocket<TData>({
 				apiRoute: configs.apiRoute,
 				searchParams: configs.searchParams,
 				options: configs.options,
+				enableRetry: configs.enableRetry,
 			});
 		} catch {
 			// Failed to create WebSocket, use SSE fallback
@@ -274,8 +341,8 @@ export class CoderApi extends Api {
 			const handleError = (event: ErrorEvent) => {
 				cleanup();
 				const is404 =
-					event.message?.includes("404") ||
-					event.error?.message?.includes("404");
+					event.message?.includes(String(HttpStatusCode.NOT_FOUND)) ||
+					event.error?.message?.includes(String(HttpStatusCode.NOT_FOUND));
 
 				if (is404 && onNotFound) {
 					connection.close();
@@ -323,14 +390,11 @@ export class CoderApi extends Api {
 /**
  * Set up logging and request interceptors for the CoderApi instance.
  */
-function setupInterceptors(
-	client: CoderApi,
-	baseUrl: string,
-	output: Logger,
-): void {
+function setupInterceptors(client: CoderApi, output: Logger): void {
 	addLoggingInterceptors(client.getAxiosInstance(), output);
 
 	client.getAxiosInstance().interceptors.request.use(async (config) => {
+		const baseUrl = client.getAxiosInstance().defaults.baseURL;
 		const headers = await getHeaders(
 			baseUrl,
 			getHeaderCommand(vscode.workspace.getConfiguration()),
@@ -356,7 +420,12 @@ function setupInterceptors(
 	client.getAxiosInstance().interceptors.response.use(
 		(r) => r,
 		async (err) => {
-			throw await CertificateError.maybeWrap(err, baseUrl, output);
+			const baseUrl = client.getAxiosInstance().defaults.baseURL;
+			if (baseUrl) {
+				throw await CertificateError.maybeWrap(err, baseUrl, output);
+			} else {
+				throw err;
+			}
 		},
 	);
 }

src/api/workspace.ts

@@ -11,7 +11,7 @@ import * as vscode from "vscode";
 import { type FeatureSet } from "../featureSet";
 import { getGlobalFlags } from "../globalFlags";
 import { escapeCommandArg } from "../util";
-import { type OneWayWebSocket } from "../websocket/oneWayWebSocket";
+import { type UnidirectionalStream } from "../websocket/eventStreamConnection";
 
 import { errToStr, createWorkspaceIdentifier } from "./api-helper";
 import { type CoderApi } from "./coderApi";
@@ -93,7 +93,7 @@ export async function streamBuildLogs(
 	client: CoderApi,
 	writeEmitter: vscode.EventEmitter<string>,
 	workspace: Workspace,
-): Promise<OneWayWebSocket<ProvisionerJobLog>> {
+): Promise<UnidirectionalStream<ProvisionerJobLog>> {
 	const socket = await client.watchBuildLogsByBuildId(
 		workspace.latest_build.id,
 		[],
@@ -131,7 +131,7 @@ export async function streamAgentLogs(
 	client: CoderApi,
 	writeEmitter: vscode.EventEmitter<string>,
 	agent: WorkspaceAgent,
-): Promise<OneWayWebSocket<WorkspaceAgentLog[]>> {
+): Promise<UnidirectionalStream<WorkspaceAgentLog[]>> {
 	const socket = await client.watchWorkspaceAgentLogs(agent.id, []);
 
 	socket.addEventListener("message", (data) => {

src/inbox.ts

@@ -7,7 +7,7 @@ import type {
 
 import type { CoderApi } from "./api/coderApi";
 import type { Logger } from "./logging/logger";
-import type { OneWayWebSocket } from "./websocket/oneWayWebSocket";
+import type { UnidirectionalStream } from "./websocket/eventStreamConnection";
 
 // These are the template IDs of our notifications.
 // Maybe in the future we should avoid hardcoding
@@ -16,7 +16,9 @@ const TEMPLATE_WORKSPACE_OUT_OF_MEMORY = "a9d027b4-ac49-4fb1-9f6d-45af15f64e7a";
 const TEMPLATE_WORKSPACE_OUT_OF_DISK = "f047f6a3-5713-40f7-85aa-0394cce9fa3a";
 
 export class Inbox implements vscode.Disposable {
-	private socket: OneWayWebSocket<GetInboxNotificationResponse> | undefined;
+	private socket:
+		| UnidirectionalStream<GetInboxNotificationResponse>
+		| undefined;
 	private disposed = false;
 
 	private constructor(private readonly logger: Logger) {}

src/remote/workspaceStateMachine.ts

@@ -21,7 +21,7 @@ import type { CoderApi } from "../api/coderApi";
 import type { PathResolver } from "../core/pathResolver";
 import type { FeatureSet } from "../featureSet";
 import type { Logger } from "../logging/logger";
-import type { OneWayWebSocket } from "../websocket/oneWayWebSocket";
+import type { UnidirectionalStream } from "../websocket/eventStreamConnection";
 
 /**
  * Manages workspace and agent state transitions until ready for SSH connection.
@@ -32,9 +32,10 @@ export class WorkspaceStateMachine implements vscode.Disposable {
 
 	private agent: { id: string; name: string } | undefined;
 
-	private buildLogSocket: OneWayWebSocket<ProvisionerJobLog> | null = null;
+	private buildLogSocket: UnidirectionalStream<ProvisionerJobLog> | null = null;
 
-	private agentLogSocket: OneWayWebSocket<WorkspaceAgentLog[]> | null = null;
+	private agentLogSocket: UnidirectionalStream<WorkspaceAgentLog[]> | null =
+		null;
 
 	constructor(
 		private readonly parts: AuthorityParts,