Consistently handle websocket errors

Author: EhabY

src/api/coderApi.ts

@@ -311,7 +311,9 @@ export class CoderApi extends Api implements vscode.Disposable {
 		});
 
 		this.attachStreamLogger(ws);
-		return ws;
+
+		// Wait for connection to open before returning
+		return await this.waitForOpen(ws);
 	}
 
 	private attachStreamLogger<TData>(
@@ -351,9 +353,8 @@ export class CoderApi extends Api implements vscode.Disposable {
 	): Promise<UnidirectionalStream<ServerSentEvent>> {
 		const { fallbackApiRoute, ...socketConfigs } = configs;
 		try {
-			const ws =
-				await this.createOneWayWebSocket<ServerSentEvent>(socketConfigs);
-			return await this.waitForOpen(ws);
+			// createOneWayWebSocket already waits for open
+			return await this.createOneWayWebSocket<ServerSentEvent>(socketConfigs);
 		} catch (error) {
 			if (this.is404Error(error)) {
 				this.output.warn(
@@ -398,10 +399,11 @@ export class CoderApi extends Api implements vscode.Disposable {
 
 	/**
 	 * Wait for a connection to open. Rejects on error.
+	 * Preserves the specific connection type (e.g., OneWayWebSocket, SseConnection).
 	 */
-	private waitForOpen<TData>(
-		connection: UnidirectionalStream<TData>,
-	): Promise<UnidirectionalStream<TData>> {
+	private waitForOpen<T extends UnidirectionalStream<unknown>>(
+		connection: T,
+	): Promise<T> {
 		return new Promise((resolve, reject) => {
 			const cleanup = () => {
 				connection.removeEventListener("open", handleOpen);
@@ -416,7 +418,10 @@ export class CoderApi extends Api implements vscode.Disposable {
 			const handleError = (event: ErrorEvent) => {
 				cleanup();
 				connection.close();
-				const error = toError(event.error, "WebSocket connection error");
+				const error = toError(
+					event.error,
+					event.message || "WebSocket connection error",
+				);
 				reject(error);
 			};
 

src/error/certificateError.ts

@@ -21,7 +21,7 @@ export abstract class CertificateError extends Error {
 	): Promise<T | undefined> {
 		const modal = options?.modal ?? false;
 		const message =
-			!modal && title ? `${title}\n\n${this.detail}` : title || this.detail;
+			!modal && title ? `${title}: ${this.detail}` : title || this.detail;
 
 		return await vscode.window.showErrorMessage(
 			message,

src/websocket/reconnectingWebSocket.ts

@@ -53,6 +53,7 @@ export class ReconnectingWebSocket<
 	#isDisposed = false; // Permanent disposal, cannot be resumed
 	#isConnecting = false;
 	#pendingReconnect = false;
+	#certRefreshAttempted = false; // Tracks if cert refresh was already attempted this connection cycle
 	readonly #onDispose?: () => void;
 
 	private constructor(
@@ -86,45 +87,23 @@ export class ReconnectingWebSocket<
 			onDispose,
 		);
 
-		try {
-			await instance.connect();
-		} catch (error) {
-			const certError = ClientCertificateError.fromError(error);
-			if (!certError) {
-				throw error;
-			}
-
-			const refreshed = await instance.handleClientCertificateError(certError);
-			if (refreshed) {
-				try {
-					await instance.connect();
-					return instance;
-				} catch (retryError) {
-					// Check if still a cert error after refresh
-					const retryCertError = ClientCertificateError.fromError(retryError);
-					if (retryCertError) {
-						void retryCertError.showNotification();
-					} else {
-						instance.#logger.error(
-							"Connection failed after certificate refresh",
-							retryError,
-						);
-					}
-				}
-			}
-
-			// Either refresh failed or retry failed - return disconnected instance
-			instance.disconnect();
-			return instance;
-		}
-
+		// connect() handles all errors internally
+		await instance.connect();
 		return instance;
 	}
 
 	get url(): string {
 		return this.#currentSocket?.url ?? "";
 	}
 
+	/**
+	 * Returns true if the socket is temporarily disconnected and not attempting to reconnect.
+	 * Use reconnect() to resume.
+	 */
+	get isDisconnected(): boolean {
+		return this.#isDisconnected;
+	}
+
 	/**
 	 * Extract the route (pathname + search) from the current socket URL for logging.
 	 * Falls back to the last known route when the socket is closed.
@@ -160,6 +139,7 @@ export class ReconnectingWebSocket<
 		if (this.#isDisconnected) {
 			this.#isDisconnected = false;
 			this.#backoffMs = this.#options.initialBackoffMs;
+			this.#certRefreshAttempted = false; // User-initiated reconnect, allow retry
 		}
 
 		if (this.#isDisposed) {
@@ -171,14 +151,13 @@ export class ReconnectingWebSocket<
 			this.#reconnectTimeoutId = null;
 		}
 
-		// If already connecting, schedule reconnect after current attempt
 		if (this.#isConnecting) {
 			this.#pendingReconnect = true;
 			return;
 		}
 
-		// connect() will close any existing socket
-		this.connect().catch((error) => this.handleConnectionError(error));
+		// connect() handles all errors internally
+		void this.connect();
 	}
 
 	/**
@@ -239,6 +218,7 @@ export class ReconnectingWebSocket<
 
 			socket.addEventListener("open", (event) => {
 				this.#backoffMs = this.#options.initialBackoffMs;
+				this.#certRefreshAttempted = false; // Reset on successful connection
 				this.executeHandlers("open", event);
 			});
 
@@ -248,17 +228,11 @@ export class ReconnectingWebSocket<
 
 			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 = toError(event.error, event.message).message;
-				if (this.isUnrecoverableHttpError(errorMessage)) {
-					this.#logger.error(
-						`Unrecoverable HTTP error for ${this.#route}: ${errorMessage}`,
-					);
-					this.disconnect();
-				}
+				// Errors during initial connection are caught by the factory (waitForOpen).
+				// This handler is for errors AFTER successful connection.
+				// Route through handleConnectionError for consistent handling.
+				const error = toError(event.error, event.message);
+				void this.handleConnectionError(error);
 			});
 
 			socket.addEventListener("close", (event) => {
@@ -272,19 +246,18 @@ export class ReconnectingWebSocket<
 					this.#logger.error(
 						`WebSocket connection closed with unrecoverable error code ${event.code}`,
 					);
-					// Suspend instead of dispose - allows recovery when credentials change
 					this.disconnect();
 					return;
 				}
 
-				// Don't reconnect on normal closure
 				if (NORMAL_CLOSURE_CODES.has(event.code)) {
 					return;
 				}
 
-				// Reconnect on abnormal closures (e.g., 1006) or other unexpected codes
 				this.scheduleReconnect();
 			});
+		} catch (error) {
+			await this.handleConnectionError(error);
 		} finally {
 			this.#isConnecting = false;
 
@@ -314,7 +287,8 @@ export class ReconnectingWebSocket<
 
 		this.#reconnectTimeoutId = setTimeout(() => {
 			this.#reconnectTimeoutId = null;
-			this.connect().catch((error) => this.handleConnectionError(error));
+			// connect() handles all errors internally
+			void this.connect();
 		}, delayMs);
 
 		this.#backoffMs = Math.min(this.#backoffMs * 2, this.#options.maxBackoffMs);
@@ -343,7 +317,15 @@ export class ReconnectingWebSocket<
 	private async handleClientCertificateError(
 		certError: ClientCertificateError,
 	): Promise<boolean> {
+		// Only attempt refresh once per connection cycle
+		if (this.#certRefreshAttempted) {
+			this.#logger.warn("Certificate refresh already attempted, not retrying");
+			void certError.showNotification();
+			return false;
+		}
+
 		if (certError.isRefreshable) {
+			this.#certRefreshAttempted = true; // Mark that we're attempting
 			this.#logger.info(
 				`Client certificate error (alert ${certError.alertCode}), attempting refresh...`,
 			);

test/unit/api/coderApi.test.ts

@@ -217,7 +217,15 @@ describe("CoderApi", () => {
 
 		beforeEach(() => {
 			api = createApi(CODER_URL, AXIOS_TOKEN);
-			const mockWs = createMockWebSocket(wsUrl);
+			// createOneWayWebSocket waits for open, so we need to fire it
+			const mockWs = createMockWebSocket(wsUrl, {
+				on: vi.fn((event, handler) => {
+					if (event === "open") {
+						setImmediate(() => handler());
+					}
+					return mockWs as Ws;
+				}),
+			});
 			setupWebSocketMock(mockWs);
 		});
 
@@ -390,15 +398,18 @@ describe("CoderApi", () => {
 			expect(EventSource).toHaveBeenCalled();
 		});
 
-		it("throws non-404 errors without SSE fallback", async () => {
+		it("handles non-404 errors without SSE fallback", async () => {
 			vi.mocked(Ws).mockImplementation(function () {
 				throw new Error("Network error");
 			});
 
-			await expect(api.watchAgentMetadata(AGENT_ID)).rejects.toThrow(
-				"Network error",
-			);
+			// Non-404 errors are handled by ReconnectingWebSocket internally
+			// (schedules reconnect with backoff) instead of being thrown
+			const connection = await api.watchAgentMetadata(AGENT_ID);
+			expect(connection).toBeInstanceOf(ReconnectingWebSocket);
 			expect(EventSource).not.toHaveBeenCalled();
+
+			connection.close();
 		});
 
 		describe("reconnection after fallback", () => {
@@ -572,6 +583,45 @@ describe("CoderApi", () => {
 				"No base URL set on REST client",
 			);
 		});
+
+		interface WebSocketErrorConversionTestCase {
+			description: string;
+			eventMessage: string;
+			expectedMessage: string;
+		}
+
+		it.each<WebSocketErrorConversionTestCase>([
+			{
+				description: "event.message when available",
+				eventMessage: "Custom error message",
+				expectedMessage: "Custom error message",
+			},
+			{
+				description: "fallback when event.message is empty",
+				eventMessage: "",
+				expectedMessage: "WebSocket connection error",
+			},
+		])(
+			"uses $desc for WebSocket error",
+			async ({ eventMessage, expectedMessage }) => {
+				api = createApi();
+				const mockWs = createMockWebSocket("wss://test", {
+					on: vi.fn((event: string, handler: (e: unknown) => void) => {
+						if (event === "error") {
+							setImmediate(() =>
+								handler({ error: undefined, message: eventMessage }),
+							);
+						}
+						return mockWs as Ws;
+					}),
+				});
+				setupWebSocketMock(mockWs);
+
+				await expect(api.watchBuildLogsByBuildId(BUILD_ID, [])).rejects.toThrow(
+					expectedMessage,
+				);
+			},
+		);
 	});
 
 	describe("getHost/getSessionToken", () => {

test/unit/websocket/reconnectingWebSocket.test.ts

@@ -90,13 +90,20 @@ describe("ReconnectingWebSocket", () => {
 					);
 				});
 
-				await expect(
-					ReconnectingWebSocket.create(factory, createMockLogger()),
-				).rejects.toThrow(`Unexpected server response: ${statusCode}`);
+				// create() returns a disconnected instance instead of throwing
+				const ws = await ReconnectingWebSocket.create(
+					factory,
+					createMockLogger(),
+				);
+
+				// Should be disconnected after unrecoverable HTTP error
+				expect(ws.isDisconnected).toBe(true);
 
 				// Should not retry after unrecoverable HTTP error
 				await vi.advanceTimersByTimeAsync(10000);
 				expect(socketCreationAttempts).toBe(1);
+
+				ws.close();
 			},
 		);