Add reconnecting websocket class and use it in CoderApi

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,