Add reconnecting websocket class and use it in CoderApi

Author: EhabY

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

@@ -36,6 +36,10 @@ 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 +51,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();
 	}
@@ -70,6 +78,30 @@ export class CoderApi extends Api {
 		return client;
 	}
 
+	setSessionToken = (token: string): void => {
+		const currentToken =
+			this.getAxiosInstance().defaults.headers.common[coderSessionTokenHeader];
+		this.getAxiosInstance().defaults.headers.common[coderSessionTokenHeader] =
+			token;
+
+		if (currentToken !== token) {
+			for (const socket of this.reconnectingSockets) {
+				socket.reconnect();
+			}
+		}
+	};
+
+	setHost = (host: string | undefined): void => {
+		const currentHost = this.getAxiosInstance().defaults.baseURL;
+		this.getAxiosInstance().defaults.baseURL = host;
+
+		if (currentHost !== host) {
+			for (const socket of this.reconnectingSockets) {
+				socket.reconnect();
+			}
+		}
+	};
+
 	watchInboxNotifications = async (
 		watchTemplates: string[],
 		watchTargets: string[],
@@ -83,6 +115,7 @@ export class CoderApi extends Api {
 				targets: watchTargets.join(","),
 			},
 			options,
+			enableRetry: true,
 		});
 	};
 
@@ -91,6 +124,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 +136,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 +183,73 @@ 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,
+			);
+
+			this.reconnectingSockets.add(
+				reconnectingSocket as ReconnectingWebSocket<unknown>,
+			);
 
-		this.attachStreamLogger(webSocket);
-		return webSocket;
+			return reconnectingSocket;
+		} else {
+			return socketFactory();
+		}
 	}
 
 	private attachStreamLogger<TData>(
@@ -230,13 +285,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

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,