WebSocket/SSE reconnection improvements

Author: EhabY

src/api/coderApi.ts

@@ -214,68 +214,62 @@ export class CoderApi extends Api implements vscode.Disposable {
 				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,
-			);
+			return this.createOneWayWebSocket<TData>(socketConfigs);
+		};
 
-			const httpAgent = await createHttpAgent(
-				vscode.workspace.getConfiguration(),
-			);
+		if (enableRetry) {
+			return this.createReconnectingSocket(socketFactory, configs.apiRoute);
+		}
+		return socketFactory();
+	}
 
-			/**
-			 * 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,
-			};
+	private async createOneWayWebSocket<TData>(
+		configs: Omit<OneWayWebSocketInit, "location">,
+	): Promise<OneWayWebSocket<TData>> {
+		const baseUrlRaw = this.getAxiosInstance().defaults.baseURL;
+		if (!baseUrlRaw) {
+			throw new Error("No base URL set on REST client");
+		}
+		const token = this.getAxiosInstance().defaults.headers.common[
+			coderSessionTokenHeader
+		] as string | undefined;
 
-			const webSocket = new OneWayWebSocket<TData>({
-				location: baseUrl,
-				...socketConfigs,
-				options: {
-					...configs.options,
-					agent: httpAgent,
-					followRedirects: true,
-					headers,
-				},
-			});
+		const headersFromCommand = await getHeaders(
+			baseUrlRaw,
+			getHeaderCommand(vscode.workspace.getConfiguration()),
+			this.output,
+		);
 
-			this.attachStreamLogger(webSocket);
-			return webSocket;
-		};
+		const httpAgent = await createHttpAgent(
+			vscode.workspace.getConfiguration(),
+		);
 
-		if (enableRetry) {
-			const reconnectingSocket = await ReconnectingWebSocket.create<TData>(
-				socketFactory,
-				this.output,
-				configs.apiRoute,
-				undefined,
-				() =>
-					this.reconnectingSockets.delete(
-						reconnectingSocket as ReconnectingWebSocket<unknown>,
-					),
-			);
+		/**
+		 * 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.reconnectingSockets.add(
-				reconnectingSocket as ReconnectingWebSocket<unknown>,
-			);
+		const baseUrl = new URL(baseUrlRaw);
+		const ws = new OneWayWebSocket<TData>({
+			location: baseUrl,
+			...configs,
+			options: {
+				...configs.options,
+				agent: httpAgent,
+				followRedirects: true,
+				headers,
+			},
+		});
 
-			return reconnectingSocket;
-		} else {
-			return socketFactory();
-		}
+		this.attachStreamLogger(ws);
+		return ws;
 	}
 
 	private attachStreamLogger<TData>(
@@ -304,45 +298,80 @@ export class CoderApi extends Api implements vscode.Disposable {
 	/**
 	 * Create a WebSocket connection with SSE fallback on 404.
 	 *
+	 * The factory tries WS first, falls back to SSE on 404. Since the factory
+	 * is called on every reconnect.
+	 *
 	 * Note: The fallback on SSE ignores all passed client options except the headers.
 	 */
-	private async createWebSocketWithFallback(configs: {
-		apiRoute: string;
-		fallbackApiRoute: string;
-		searchParams?: Record<string, string> | URLSearchParams;
-		options?: ClientOptions;
-		enableRetry?: boolean;
-	}): Promise<UnidirectionalStream<ServerSentEvent>> {
-		let webSocket: UnidirectionalStream<ServerSentEvent>;
-		try {
-			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(
-				configs.fallbackApiRoute,
-				configs.searchParams,
-				configs.options?.headers,
+	private async createWebSocketWithFallback(
+		configs: Omit<OneWayWebSocketInit, "location"> & {
+			fallbackApiRoute: string;
+			enableRetry?: boolean;
+		},
+	): Promise<UnidirectionalStream<ServerSentEvent>> {
+		const { fallbackApiRoute, enableRetry, ...socketConfigs } = configs;
+		const socketFactory: SocketFactory<ServerSentEvent> = async () => {
+			try {
+				const ws =
+					await this.createOneWayWebSocket<ServerSentEvent>(socketConfigs);
+				return await this.waitForOpen(ws);
+			} catch (error) {
+				if (this.is404Error(error)) {
+					this.output.warn(
+						`WebSocket failed, using SSE fallback: ${socketConfigs.apiRoute}`,
+					);
+					const sse = this.createSseConnection(
+						fallbackApiRoute,
+						socketConfigs.searchParams,
+						socketConfigs.options?.headers,
+					);
+					return await this.waitForOpen(sse);
+				}
+				throw error;
+			}
+		};
+
+		if (enableRetry) {
+			return this.createReconnectingSocket(
+				socketFactory,
+				socketConfigs.apiRoute,
 			);
 		}
+		return socketFactory();
+	}
 
-		return this.waitForConnection(webSocket, () =>
-			this.createSseFallback(
-				configs.fallbackApiRoute,
-				configs.searchParams,
-				configs.options?.headers,
-			),
-		);
+	/**
+	 * Create an SSE connection without waiting for connection.
+	 */
+	private createSseConnection(
+		apiRoute: string,
+		searchParams?: Record<string, string> | URLSearchParams,
+		optionsHeaders?: Record<string, string>,
+	): SseConnection {
+		const baseUrlRaw = this.getAxiosInstance().defaults.baseURL;
+		if (!baseUrlRaw) {
+			throw new Error("No base URL set on REST client");
+		}
+		const url = new URL(baseUrlRaw);
+		const sse = new SseConnection({
+			location: url,
+			apiRoute,
+			searchParams,
+			axiosInstance: this.getAxiosInstance(),
+			optionsHeaders,
+			logger: this.output,
+		});
+
+		this.attachStreamLogger(sse);
+		return sse;
 	}
 
-	private waitForConnection(
-		connection: UnidirectionalStream<ServerSentEvent>,
-		onNotFound?: () => Promise<UnidirectionalStream<ServerSentEvent>>,
-	): Promise<UnidirectionalStream<ServerSentEvent>> {
+	/**
+	 * Wait for a connection to open. Rejects on error.
+	 */
+	private waitForOpen<TData>(
+		connection: UnidirectionalStream<TData>,
+	): Promise<UnidirectionalStream<TData>> {
 		return new Promise((resolve, reject) => {
 			const cleanup = () => {
 				connection.removeEventListener("open", handleOpen);
@@ -356,21 +385,8 @@ export class CoderApi extends Api implements vscode.Disposable {
 
 			const handleError = (event: ErrorEvent) => {
 				cleanup();
-				const is404 =
-					event.message?.includes(String(HttpStatusCode.NOT_FOUND)) ||
-					event.error?.message?.includes(String(HttpStatusCode.NOT_FOUND));
-
-				if (is404 && onNotFound) {
-					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));
-				}
+				connection.close();
+				reject(event.error || new Error(event.message));
 			};
 
 			connection.addEventListener("open", handleOpen);
@@ -379,32 +395,36 @@ export class CoderApi extends Api implements vscode.Disposable {
 	}
 
 	/**
-	 * Create SSE fallback connection
+	 * Check if an error is a 404 Not Found error.
 	 */
-	private async createSseFallback(
-		apiRoute: string,
-		searchParams?: Record<string, string> | URLSearchParams,
-		optionsHeaders?: Record<string, string>,
-	): Promise<UnidirectionalStream<ServerSentEvent>> {
-		this.output.warn(`WebSocket failed, using SSE fallback: ${apiRoute}`);
-
-		const baseUrlRaw = this.getAxiosInstance().defaults.baseURL;
-		if (!baseUrlRaw) {
-			throw new Error("No base URL set on REST client");
-		}
+	private is404Error(error: unknown): boolean {
+		const msg = error instanceof Error ? error.message : String(error);
+		return msg.includes(String(HttpStatusCode.NOT_FOUND));
+	}
 
-		const baseUrl = new URL(baseUrlRaw);
-		const sseConnection = new SseConnection({
-			location: baseUrl,
+	/**
+	 * Create a ReconnectingWebSocket and track it for lifecycle management.
+	 */
+	private async createReconnectingSocket<TData>(
+		socketFactory: SocketFactory<TData>,
+		apiRoute: string,
+	): Promise<ReconnectingWebSocket<TData>> {
+		const reconnectingSocket = await ReconnectingWebSocket.create<TData>(
+			socketFactory,
+			this.output,
 			apiRoute,
-			searchParams,
-			axiosInstance: this.getAxiosInstance(),
-			optionsHeaders: optionsHeaders,
-			logger: this.output,
-		});
+			undefined,
+			() =>
+				this.reconnectingSockets.delete(
+					reconnectingSocket as ReconnectingWebSocket<unknown>,
+				),
+		);
+
+		this.reconnectingSockets.add(
+			reconnectingSocket as ReconnectingWebSocket<unknown>,
+		);
 
-		this.attachStreamLogger(sseConnection);
-		return this.waitForConnection(sseConnection);
+		return reconnectingSocket;
 	}
 }
 

src/websocket/codes.ts

@@ -19,7 +19,9 @@ export const WebSocketCloseCode = {
 
 /** HTTP status codes used for socket creation and connection logic */
 export const HttpStatusCode = {
-	/** Authentication or permission denied */
+	/** Authentication required */
+	UNAUTHORIZED: 401,
+	/** Permission denied */
 	FORBIDDEN: 403,
 	/** Endpoint not found */
 	NOT_FOUND: 404,
@@ -43,6 +45,7 @@ export const UNRECOVERABLE_WS_CLOSE_CODES = new Set<number>([
  * These appear during socket creation and should stop reconnection attempts.
  */
 export const UNRECOVERABLE_HTTP_CODES = new Set<number>([
+	HttpStatusCode.UNAUTHORIZED,
 	HttpStatusCode.FORBIDDEN,
 	HttpStatusCode.NOT_FOUND,
 	HttpStatusCode.GONE,

src/websocket/reconnectingWebSocket.ts

@@ -186,6 +186,17 @@ export class ReconnectingWebSocket<TData = unknown>
 
 			socket.addEventListener("error", (event) => {
 				this.executeHandlers("error", event);
+
+				// Check for unrecoverable HTTP errors in the error event
+				// HTTP errors during handshake fire 'error' then 'close' with 1006
+				// We need to suspend here to prevent infinite reconnect loops
+				const errorMessage = event.error?.message ?? event.message ?? "";
+				if (this.isUnrecoverableHttpError(errorMessage)) {
+					this.#logger.error(
+						`Unrecoverable HTTP error for ${this.#apiRoute}: ${errorMessage}`,
+					);
+					this.suspend();
+				}
 			});
 
 			socket.addEventListener("close", (event) => {
@@ -289,12 +300,12 @@ export class ReconnectingWebSocket<TData = unknown>
 	}
 
 	/**
-	 * Check if an error contains an unrecoverable HTTP status code.
+	 * Check if an error message contains an unrecoverable HTTP status code.
 	 */
 	private isUnrecoverableHttpError(error: unknown): boolean {
-		const errorMessage = error instanceof Error ? error.message : String(error);
+		const message = (error as { message?: string }).message || String(error);
 		for (const code of UNRECOVERABLE_HTTP_CODES) {
-			if (errorMessage.includes(String(code))) {
+			if (message.includes(String(code))) {
 				return true;
 			}
 		}