Add SSE fallback to some one way WS connections

Author: EhabY

src/api/coderApi.ts

@@ -6,17 +6,18 @@ import {
 } from "axios";
 import { Api } from "coder/site/src/api/api";
 import {
+	type ServerSentEvent,
 	type GetInboxNotificationResponse,
 	type ProvisionerJobLog,
-	type ServerSentEvent,
 	type Workspace,
 	type WorkspaceAgent,
 } from "coder/site/src/api/typesGenerated";
 import * as vscode from "vscode";
-import { type ClientOptions } from "ws";
+import { type ClientOptions, type CloseEvent, type ErrorEvent } from "ws";
 
 import { CertificateError } from "../error";
 import { getHeaderCommand, getHeaders } from "../headers";
+import { EventStreamLogger } from "../logging/eventStreamLogger";
 import {
 	createRequestMeta,
 	logRequest,
@@ -29,11 +30,12 @@ import {
 	HttpClientLogLevel,
 } from "../logging/types";
 import { sizeOf } from "../logging/utils";
-import { WsLogger } from "../logging/wsLogger";
+import { type UnidirectionalStream } from "../websocket/eventStreamConnection";
 import {
 	OneWayWebSocket,
 	type OneWayWebSocketInit,
 } from "../websocket/oneWayWebSocket";
+import { SseConnection } from "../websocket/sseConnection";
 
 import { createHttpAgent } from "./utils";
 
@@ -84,8 +86,9 @@ export class CoderApi extends Api {
 	};
 
 	watchWorkspace = async (workspace: Workspace, options?: ClientOptions) => {
-		return this.createWebSocket<ServerSentEvent>({
+		return this.createWebSocketWithFallback<ServerSentEvent>({
 			apiRoute: `/api/v2/workspaces/${workspace.id}/watch-ws`,
+			fallbackApiRoute: `/api/v2/workspaces/${workspace.id}/watch`,
 			options,
 		});
 	};
@@ -94,8 +97,9 @@ export class CoderApi extends Api {
 		agentId: WorkspaceAgent["id"],
 		options?: ClientOptions,
 	) => {
-		return this.createWebSocket<ServerSentEvent>({
+		return this.createWebSocketWithFallback<ServerSentEvent>({
 			apiRoute: `/api/v2/workspaceagents/${agentId}/watch-metadata-ws`,
+			fallbackApiRoute: `/api/v2/workspaceagents/${agentId}/watch-metadata`,
 			options,
 		});
 	};
@@ -137,6 +141,7 @@ export class CoderApi extends Api {
 		const httpAgent = await createHttpAgent(
 			vscode.workspace.getConfiguration(),
 		);
+
 		const webSocket = new OneWayWebSocket<TData>({
 			location: baseUrl,
 			...configs,
@@ -152,28 +157,123 @@ export class CoderApi extends Api {
 			},
 		});
 
-		const wsUrl = new URL(webSocket.url);
-		const pathWithQuery = wsUrl.pathname + wsUrl.search;
-		const wsLogger = new WsLogger(this.output, pathWithQuery);
-		wsLogger.logConnecting();
+		this.attachStreamLogger(webSocket);
+		return webSocket;
+	}
 
-		webSocket.addEventListener("open", () => {
-			wsLogger.logOpen();
-		});
+	private attachStreamLogger<TData>(
+		connection: UnidirectionalStream<TData>,
+	): void {
+		const url = new URL(connection.url);
+		const logger = new EventStreamLogger(
+			this.output,
+			url.pathname + url.search,
+			url.protocol.startsWith("http") ? "SSE" : "WS",
+		);
+		logger.logConnecting();
 
-		webSocket.addEventListener("message", (event) => {
-			wsLogger.logMessage(event.sourceEvent.data);
-		});
+		connection.addEventListener("open", () => logger.logOpen());
+		connection.addEventListener("close", (event: CloseEvent) =>
+			logger.logClose(event.code, event.reason),
+		);
+		connection.addEventListener("error", (event: ErrorEvent) =>
+			logger.logError(event.error, event.message),
+		);
+		connection.addEventListener("message", (event) =>
+			logger.logMessage(event.sourceEvent.data),
+		);
+	}
+
+	/**
+	 * Create a WebSocket connection with SSE fallback on 404
+	 */
+	private async createWebSocketWithFallback<TData = unknown>(configs: {
+		apiRoute: string;
+		fallbackApiRoute: string;
+		searchParams?: Record<string, string> | URLSearchParams;
+		options?: ClientOptions;
+	}): Promise<UnidirectionalStream<TData>> {
+		let webSocket: OneWayWebSocket<TData>;
+		try {
+			webSocket = await this.createWebSocket<TData>({
+				apiRoute: configs.apiRoute,
+				searchParams: configs.searchParams,
+				options: configs.options,
+			});
+		} catch {
+			// Failed to create WebSocket, use SSE fallback
+			return this.createSseFallback<TData>(
+				configs.fallbackApiRoute,
+				configs.searchParams,
+			);
+		}
 
-		webSocket.addEventListener("close", (event) => {
-			wsLogger.logClose(event.code, event.reason);
+		return this.waitForConnection(webSocket, () =>
+			this.createSseFallback<TData>(
+				configs.fallbackApiRoute,
+				configs.searchParams,
+			),
+		);
+	}
+
+	private waitForConnection<TData>(
+		connection: UnidirectionalStream<TData>,
+		onNotFound?: () => Promise<UnidirectionalStream<TData>>,
+	): Promise<UnidirectionalStream<TData>> {
+		return new Promise((resolve, reject) => {
+			const cleanup = () => {
+				connection.removeEventListener("open", handleOpen);
+				connection.removeEventListener("error", handleError);
+			};
+
+			const handleOpen = () => {
+				cleanup();
+				resolve(connection);
+			};
+
+			const handleError = (event: ErrorEvent) => {
+				cleanup();
+				const is404 =
+					event.message?.includes("404") ||
+					event.error?.message?.includes("404");
+
+				if (is404 && onNotFound) {
+					connection.close();
+					onNotFound().then(resolve).catch(reject);
+				} else {
+					reject(event.error || new Error(event.message));
+				}
+			};
+
+			connection.addEventListener("open", handleOpen);
+			connection.addEventListener("error", handleError);
 		});
+	}
 
-		webSocket.addEventListener("error", (event) => {
-			wsLogger.logError(event.error, event.message);
+	/**
+	 * Create SSE fallback connection
+	 */
+	private async createSseFallback<TData = unknown>(
+		apiRoute: string,
+		searchParams?: Record<string, string> | URLSearchParams,
+	): Promise<UnidirectionalStream<TData>> {
+		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");
+		}
+
+		const baseUrl = new URL(baseUrlRaw);
+		const sseConnection = new SseConnection({
+			location: baseUrl,
+			apiRoute,
+			searchParams,
+			axiosInstance: this.getAxiosInstance(),
 		});
 
-		return webSocket;
+		this.attachStreamLogger(sseConnection);
+		return this.waitForConnection(sseConnection);
 	}
 }
 

src/api/streamingFetchAdapter.ts

@@ -0,0 +1,62 @@
+import { type AxiosInstance } from "axios";
+import { type FetchLikeInit, type FetchLikeResponse } from "eventsource";
+import { type IncomingMessage } from "http";
+
+/**
+ * Creates a fetch adapter using an Axios instance that returns streaming responses.
+ * This is used by EventSource to make authenticated SSE connections.
+ */
+export function createStreamingFetchAdapter(
+	axiosInstance: AxiosInstance,
+): (url: string | URL, init?: FetchLikeInit) => Promise<FetchLikeResponse> {
+	return async (
+		url: string | URL,
+		init?: FetchLikeInit,
+	): Promise<FetchLikeResponse> => {
+		const urlStr = url.toString();
+
+		const response = await axiosInstance.request<IncomingMessage>({
+			url: urlStr,
+			signal: init?.signal,
+			headers: init?.headers,
+			responseType: "stream",
+			validateStatus: () => true, // Don't throw on any status code
+		});
+
+		const stream = new ReadableStream({
+			start(controller) {
+				response.data.on("data", (chunk: Buffer) => {
+					controller.enqueue(chunk);
+				});
+
+				response.data.on("end", () => {
+					controller.close();
+				});
+
+				response.data.on("error", (err: Error) => {
+					controller.error(err);
+				});
+			},
+
+			cancel() {
+				response.data.destroy();
+				return Promise.resolve();
+			},
+		});
+
+		return {
+			body: {
+				getReader: () => stream.getReader(),
+			},
+			url: urlStr,
+			status: response.status,
+			redirected: response.request?.res?.responseUrl !== urlStr,
+			headers: {
+				get: (name: string) => {
+					const value = response.headers[name.toLowerCase()];
+					return value === undefined ? null : String(value);
+				},
+			},
+		};
+	};
+}

src/logging/wsLogger.ts renamed to src/logging/eventStreamLogger.ts

@@ -12,31 +12,35 @@ const numFormatter = new Intl.NumberFormat("en", {
 	compactDisplay: "short",
 });
 
-export class WsLogger {
+export class EventStreamLogger {
 	private readonly logger: Logger;
 	private readonly url: string;
 	private readonly id: string;
+	private readonly protocol: string;
 	private readonly startedAt: number;
 	private openedAt?: number;
 	private msgCount = 0;
 	private byteCount = 0;
 	private unknownByteCount = false;
 
-	constructor(logger: Logger, url: string) {
+	constructor(logger: Logger, url: string, protocol: "WS" | "SSE") {
 		this.logger = logger;
 		this.url = url;
+		this.protocol = protocol;
 		this.id = createRequestId();
 		this.startedAt = Date.now();
 	}
 
 	logConnecting(): void {
-		this.logger.trace(`→ WS ${shortId(this.id)} ${this.url}`);
+		this.logger.trace(`→ ${this.protocol} ${shortId(this.id)} ${this.url}`);
 	}
 
 	logOpen(): void {
 		this.openedAt = Date.now();
 		const time = formatTime(this.openedAt - this.startedAt);
-		this.logger.trace(`← WS ${shortId(this.id)} connected ${this.url} ${time}`);
+		this.logger.trace(
+			`← ${this.protocol} ${shortId(this.id)} connected ${this.url} ${time}`,
+		);
 	}
 
 	logMessage(data: unknown): void {
@@ -62,15 +66,15 @@ export class WsLogger {
 		const statsStr = ` [${stats.join(", ")}]`;
 
 		this.logger.trace(
-			`▣ WS ${shortId(this.id)} closed ${this.url}${codeStr}${reasonStr}${statsStr}`,
+			`▣ ${this.protocol} ${shortId(this.id)} closed ${this.url}${codeStr}${reasonStr}${statsStr}`,
 		);
 	}
 
 	logError(error: unknown, message: string): void {
 		const time = formatTime(Date.now() - this.startedAt);
 		const errorMsg = message || errToStr(error, "connection error");
 		this.logger.error(
-			`✗ WS ${shortId(this.id)} error ${this.url} ${time} - ${errorMsg}`,
+			`✗ ${this.protocol} ${shortId(this.id)} error ${this.url} ${time} - ${errorMsg}`,
 			error,
 		);
 	}

src/websocket/eventStreamConnection.ts

@@ -0,0 +1,51 @@
+import { type WebSocketEventType } from "coder/site/src/utils/OneWayWebSocket";
+import {
+	type CloseEvent,
+	type Event as WsEvent,
+	type ErrorEvent,
+	type MessageEvent,
+} from "ws";
+
+// Event payload types matching OneWayWebSocket
+export type ParsedMessageEvent<TData> = Readonly<
+	| {
+			sourceEvent: MessageEvent;
+			parsedMessage: TData;
+			parseError: undefined;
+	  }
+	| {
+			sourceEvent: MessageEvent;
+			parsedMessage: undefined;
+			parseError: Error;
+	  }
+>;
+
+export type EventPayloadMap<TData> = {
+	close: CloseEvent;
+	error: ErrorEvent;
+	message: ParsedMessageEvent<TData>;
+	open: WsEvent;
+};
+
+export type EventHandler<TData, TEvent extends WebSocketEventType> = (
+	payload: EventPayloadMap<TData>[TEvent],
+) => void;
+
+/**
+ * Common interface for both WebSocket and SSE connections that handle event streams.
+ * Matches the OneWayWebSocket interface for compatibility.
+ */
+export interface UnidirectionalStream<TData> {
+	readonly url: string;
+	addEventListener<TEvent extends WebSocketEventType>(
+		eventType: TEvent,
+		callback: EventHandler<TData, TEvent>,
+	): void;
+
+	removeEventListener<TEvent extends WebSocketEventType>(
+		eventType: TEvent,
+		callback: EventHandler<TData, TEvent>,
+	): void;
+
+	close(code?: number, reason?: string): void;
+}