Add WebSocket suspension

Author: EhabY

src/api/coderApi.ts

@@ -31,7 +31,7 @@ import {
 	HttpClientLogLevel,
 } from "../logging/types";
 import { sizeOf } from "../logging/utils";
-import { HttpStatusCode } from "../websocket/codes";
+import { HttpStatusCode, WebSocketCloseCode } from "../websocket/codes";
 import {
 	type UnidirectionalStream,
 	type CloseEvent,
@@ -55,7 +55,7 @@ const coderSessionTokenHeader = "Coder-Session-Token";
  * Unified API class that includes both REST API methods from the base Api class
  * and WebSocket methods for real-time functionality.
  */
-export class CoderApi extends Api {
+export class CoderApi extends Api implements vscode.Disposable {
 	private readonly reconnectingSockets = new Set<
 		ReconnectingWebSocket<unknown>
 	>();
@@ -102,11 +102,27 @@ export class CoderApi extends Api {
 
 		if (currentHost !== host) {
 			for (const socket of this.reconnectingSockets) {
-				socket.reconnect();
+				if (host) {
+					socket.reconnect();
+				} else {
+					// No host means logout - suspend sockets (can resume when host is set again)
+					socket.suspend(WebSocketCloseCode.NORMAL, "Host cleared");
+				}
 			}
 		}
 	};
 
+	/**
+	 * Permanently dispose all WebSocket connections.
+	 * This clears handlers and prevents reconnection.
+	 */
+	dispose(): void {
+		for (const socket of this.reconnectingSockets) {
+			socket.close();
+		}
+		this.reconnectingSockets.clear();
+	}
+
 	watchInboxNotifications = async (
 		watchTemplates: string[],
 		watchTargets: string[],
@@ -125,7 +141,7 @@ export class CoderApi extends Api {
 	};
 
 	watchWorkspace = async (workspace: Workspace, options?: ClientOptions) => {
-		return this.createWebSocketWithFallback<ServerSentEvent>({
+		return this.createWebSocketWithFallback({
 			apiRoute: `/api/v2/workspaces/${workspace.id}/watch-ws`,
 			fallbackApiRoute: `/api/v2/workspaces/${workspace.id}/watch`,
 			options,
@@ -137,7 +153,7 @@ export class CoderApi extends Api {
 		agentId: WorkspaceAgent["id"],
 		options?: ClientOptions,
 	) => {
-		return this.createWebSocketWithFallback<ServerSentEvent>({
+		return this.createWebSocketWithFallback({
 			apiRoute: `/api/v2/workspaceagents/${agentId}/watch-metadata-ws`,
 			fallbackApiRoute: `/api/v2/workspaceagents/${agentId}/watch-metadata`,
 			options,
@@ -290,43 +306,43 @@ export class CoderApi extends Api {
 	 *
 	 * Note: The fallback on SSE ignores all passed client options except the headers.
 	 */
-	private async createWebSocketWithFallback<TData = unknown>(configs: {
+	private async createWebSocketWithFallback(configs: {
 		apiRoute: string;
 		fallbackApiRoute: string;
 		searchParams?: Record<string, string> | URLSearchParams;
 		options?: ClientOptions;
 		enableRetry?: boolean;
-	}): Promise<UnidirectionalStream<TData>> {
-		let webSocket: UnidirectionalStream<TData>;
+	}): Promise<UnidirectionalStream<ServerSentEvent>> {
+		let webSocket: UnidirectionalStream<ServerSentEvent>;
 		try {
-			webSocket = await this.createWebSocket<TData>({
+			webSocket = await this.createWebSocket<ServerSentEvent>({
 				apiRoute: configs.apiRoute,
 				searchParams: configs.searchParams,
 				options: configs.options,
 				enableRetry: configs.enableRetry,
 			});
 		} catch {
 			// Failed to create WebSocket, use SSE fallback
-			return this.createSseFallback<TData>(
+			return this.createSseFallback(
 				configs.fallbackApiRoute,
 				configs.searchParams,
 				configs.options?.headers,
 			);
 		}
 
 		return this.waitForConnection(webSocket, () =>
-			this.createSseFallback<TData>(
+			this.createSseFallback(
 				configs.fallbackApiRoute,
 				configs.searchParams,
 				configs.options?.headers,
 			),
 		);
 	}
 
-	private waitForConnection<TData>(
-		connection: UnidirectionalStream<TData>,
-		onNotFound?: () => Promise<UnidirectionalStream<TData>>,
-	): Promise<UnidirectionalStream<TData>> {
+	private waitForConnection(
+		connection: UnidirectionalStream<ServerSentEvent>,
+		onNotFound?: () => Promise<UnidirectionalStream<ServerSentEvent>>,
+	): Promise<UnidirectionalStream<ServerSentEvent>> {
 		return new Promise((resolve, reject) => {
 			const cleanup = () => {
 				connection.removeEventListener("open", handleOpen);
@@ -345,7 +361,12 @@ export class CoderApi extends Api {
 					event.error?.message?.includes(String(HttpStatusCode.NOT_FOUND));
 
 				if (is404 && onNotFound) {
-					connection.close();
+					if (connection instanceof ReconnectingWebSocket) {
+						// We can attempt this again if we change the host
+						connection.suspend();
+					} else {
+						connection.close();
+					}
 					onNotFound().then(resolve).catch(reject);
 				} else {
 					reject(event.error || new Error(event.message));
@@ -360,11 +381,11 @@ export class CoderApi extends Api {
 	/**
 	 * Create SSE fallback connection
 	 */
-	private async createSseFallback<TData = unknown>(
+	private async createSseFallback(
 		apiRoute: string,
 		searchParams?: Record<string, string> | URLSearchParams,
 		optionsHeaders?: Record<string, string>,
-	): Promise<UnidirectionalStream<TData>> {
+	): Promise<UnidirectionalStream<ServerSentEvent>> {
 		this.output.warn(`WebSocket failed, using SSE fallback: ${apiRoute}`);
 
 		const baseUrlRaw = this.getAxiosInstance().defaults.baseURL;

src/extension.ts

@@ -93,6 +93,7 @@ export async function activate(ctx: vscode.ExtensionContext): Promise<void> {
 		await secretsManager.getSessionToken(deployment?.label ?? ""),
 		output,
 	);
+	ctx.subscriptions.push(client);
 	attachOAuthInterceptors(client, output, oauthSessionManager);
 
 	const myWorkspacesProvider = new WorkspaceProvider(

src/remote/remote.ts

@@ -175,6 +175,7 @@ export class Remote {
 			// disallow logging out/in altogether, but for now just use a separate
 			// client to remain unaffected by whatever the plugin is doing.
 			const workspaceClient = CoderApi.create(baseUrlRaw, token, this.logger);
+			disposables.push(workspaceClient);
 			attachOAuthInterceptors(workspaceClient, this.logger, remoteOAuthManager);
 			// Store for use in commands.
 			this.commands.workspaceRestClient = workspaceClient;

src/websocket/codes.ts

@@ -44,6 +44,7 @@ export const UNRECOVERABLE_WS_CLOSE_CODES = new Set<number>([
  */
 export const UNRECOVERABLE_HTTP_CODES = new Set<number>([
 	HttpStatusCode.FORBIDDEN,
+	HttpStatusCode.NOT_FOUND,
 	HttpStatusCode.GONE,
 	HttpStatusCode.UPGRADE_REQUIRED,
 ]);

src/websocket/reconnectingWebSocket.ts

@@ -41,7 +41,8 @@ export class ReconnectingWebSocket<TData = unknown>
 	#currentSocket: UnidirectionalStream<TData> | null = null;
 	#backoffMs: number;
 	#reconnectTimeoutId: NodeJS.Timeout | null = null;
-	#isDisposed = false;
+	#isSuspended = false; // Temporary pause, can be resumed via reconnect()
+	#isDisposed = false; // Permanent disposal, cannot be resumed
 	#isConnecting = false;
 	#pendingReconnect = false;
 	readonly #onDispose?: () => void;
@@ -102,6 +103,11 @@ export class ReconnectingWebSocket<TData = unknown>
 	}
 
 	reconnect(): void {
+		if (this.#isSuspended) {
+			this.#isSuspended = false;
+			this.#backoffMs = this.#options.initialBackoffMs;
+		}
+
 		if (this.#isDisposed) {
 			return;
 		}
@@ -121,6 +127,18 @@ export class ReconnectingWebSocket<TData = unknown>
 		this.connect().catch((error) => this.handleConnectionError(error));
 	}
 
+	/**
+	 * Temporarily suspend the socket. Can be resumed via reconnect().
+	 */
+	suspend(code?: number, reason?: string): void {
+		if (this.#isDisposed || this.#isSuspended) {
+			return;
+		}
+
+		this.#isSuspended = true;
+		this.clearCurrentSocket(code, reason);
+	}
+
 	close(code?: number, reason?: string): void {
 		if (this.#isDisposed) {
 			return;
@@ -139,7 +157,7 @@ export class ReconnectingWebSocket<TData = unknown>
 	}
 
 	private async connect(): Promise<void> {
-		if (this.#isDisposed || this.#isConnecting) {
+		if (this.#isDisposed || this.#isSuspended || this.#isConnecting) {
 			return;
 		}
 
@@ -171,7 +189,7 @@ export class ReconnectingWebSocket<TData = unknown>
 			});
 
 			socket.addEventListener("close", (event) => {
-				if (this.#isDisposed) {
+				if (this.#isDisposed || this.#isSuspended) {
 					return;
 				}
 
@@ -181,7 +199,8 @@ export class ReconnectingWebSocket<TData = unknown>
 					this.#logger.error(
 						`WebSocket connection closed with unrecoverable error code ${event.code}`,
 					);
-					this.dispose();
+					// Suspend instead of dispose - allows recovery when credentials change
+					this.suspend();
 					return;
 				}
 
@@ -204,7 +223,11 @@ export class ReconnectingWebSocket<TData = unknown>
 	}
 
 	private scheduleReconnect(): void {
-		if (this.#isDisposed || this.#reconnectTimeoutId !== null) {
+		if (
+			this.#isDisposed ||
+			this.#isSuspended ||
+			this.#reconnectTimeoutId !== null
+		) {
 			return;
 		}
 
@@ -241,11 +264,11 @@ export class ReconnectingWebSocket<TData = unknown>
 	}
 
 	/**
-	 * Checks if the error is unrecoverable and disposes the connection,
+	 * Checks if the error is unrecoverable and suspends the connection,
 	 * otherwise schedules a reconnect.
 	 */
 	private handleConnectionError(error: unknown): void {
-		if (this.#isDisposed) {
+		if (this.#isDisposed || this.#isSuspended) {
 			return;
 		}
 
@@ -254,7 +277,7 @@ export class ReconnectingWebSocket<TData = unknown>
 				`Unrecoverable HTTP error during connection for ${this.#apiRoute}`,
 				error,
 			);
-			this.dispose();
+			this.suspend();
 			return;
 		}
 
@@ -284,7 +307,16 @@ export class ReconnectingWebSocket<TData = unknown>
 		}
 
 		this.#isDisposed = true;
+		this.clearCurrentSocket(code, reason);
 
+		for (const set of Object.values(this.#eventHandlers)) {
+			set.clear();
+		}
+
+		this.#onDispose?.();
+	}
+
+	private clearCurrentSocket(code?: number, reason?: string): void {
 		if (this.#reconnectTimeoutId !== null) {
 			clearTimeout(this.#reconnectTimeoutId);
 			this.#reconnectTimeoutId = null;
@@ -294,11 +326,5 @@ export class ReconnectingWebSocket<TData = unknown>
 			this.#currentSocket.close(code, reason);
 			this.#currentSocket = null;
 		}
-
-		for (const set of Object.values(this.#eventHandlers)) {
-			set.clear();
-		}
-
-		this.#onDispose?.();
 	}
 }

test/unit/api/coderApi.test.ts

@@ -435,6 +435,19 @@ describe("CoderApi", () => {
 			expect(sockets[0].close).not.toHaveBeenCalled();
 			expect(sockets).toHaveLength(1);
 		});
+
+		it("suspends sockets when host is set to empty string (logout)", async () => {
+			const sockets = setupAutoOpeningWebSocket();
+			api = createApi(CODER_URL, AXIOS_TOKEN);
+			await api.watchAgentMetadata(AGENT_ID);
+
+			// Setting host to empty string (logout) should suspend (not permanently close)
+			api.setHost("");
+			await new Promise((resolve) => setImmediate(resolve));
+
+			expect(sockets[0].close).toHaveBeenCalledWith(1000, "Host cleared");
+			expect(sockets).toHaveLength(1);
+		});
 	});
 
 	describe("Error Handling", () => {