Review comments: better error handling and documentation

Author: EhabY

src/api/coderApi.ts

@@ -144,6 +144,12 @@ export class CoderApi extends Api {
 			vscode.workspace.getConfiguration(),
 		);
 
+		/**
+		 * 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,
@@ -189,7 +195,9 @@ export class CoderApi extends Api {
 	}
 
 	/**
-	 * Create a WebSocket connection with SSE fallback on 404
+	 * Create a WebSocket connection with SSE fallback on 404.
+	 *
+	 * Note: The fallback on SSE ignores all passed client options except the headers.
 	 */
 	private async createWebSocketWithFallback<TData = unknown>(configs: {
 		apiRoute: string;
@@ -209,13 +217,15 @@ export class CoderApi extends Api {
 			return this.createSseFallback<TData>(
 				configs.fallbackApiRoute,
 				configs.searchParams,
+				configs.options?.headers,
 			);
 		}
 
 		return this.waitForConnection(webSocket, () =>
 			this.createSseFallback<TData>(
 				configs.fallbackApiRoute,
 				configs.searchParams,
+				configs.options?.headers,
 			),
 		);
 	}
@@ -260,6 +270,7 @@ export class CoderApi extends Api {
 	private async createSseFallback<TData = unknown>(
 		apiRoute: string,
 		searchParams?: Record<string, string> | URLSearchParams,
+		optionsHeaders?: Record<string, string>,
 	): Promise<UnidirectionalStream<TData>> {
 		this.output.warn(`WebSocket failed, using SSE fallback: ${apiRoute}`);
 
@@ -274,6 +285,8 @@ export class CoderApi extends Api {
 			apiRoute,
 			searchParams,
 			axiosInstance: this.getAxiosInstance(),
+			optionsHeaders: optionsHeaders,
+			logger: this.output,
 		});
 
 		this.attachStreamLogger(sseConnection);

src/api/streamingFetchAdapter.ts

@@ -8,6 +8,7 @@ import { type IncomingMessage } from "http";
  */
 export function createStreamingFetchAdapter(
 	axiosInstance: AxiosInstance,
+	configHeaders?: Record<string, string>,
 ): (url: string | URL, init?: FetchLikeInit) => Promise<FetchLikeResponse> {
 	return async (
 		url: string | URL,
@@ -18,19 +19,27 @@ export function createStreamingFetchAdapter(
 		const response = await axiosInstance.request<IncomingMessage>({
 			url: urlStr,
 			signal: init?.signal,
-			headers: init?.headers,
+			headers: { ...init?.headers, ...configHeaders },
 			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);
+					try {
+						controller.enqueue(chunk);
+					} catch {
+						// Stream already closed or errored, ignore
+					}
 				});
 
 				response.data.on("end", () => {
-					controller.close();
+					try {
+						controller.close();
+					} catch {
+						// Stream already closed, ignore
+					}
 				});
 
 				response.data.on("error", (err: Error) => {

src/websocket/sseConnection.ts

@@ -4,6 +4,7 @@ import { type WebSocketEventType } from "coder/site/src/utils/OneWayWebSocket";
 import { EventSource } from "eventsource";
 
 import { createStreamingFetchAdapter } from "../api/streamingFetchAdapter";
+import { type Logger } from "../logging/logger";
 
 import { getQueryString } from "./utils";
 
@@ -24,11 +25,14 @@ export type SseConnectionInit = {
 	location: { protocol: string; host: string };
 	apiRoute: string;
 	searchParams?: Record<string, string> | URLSearchParams;
+	optionsHeaders?: Record<string, string>;
 	axiosInstance: AxiosInstance;
+	logger: Logger;
 };
 
 export class SseConnection implements UnidirectionalStream<ServerSentEvent> {
 	private readonly eventSource: EventSource;
+	private readonly logger: Logger;
 	private readonly callbacks = {
 		open: new Set<EventHandler<ServerSentEvent, "open">>(),
 		close: new Set<EventHandler<ServerSentEvent, "close">>(),
@@ -43,9 +47,13 @@ export class SseConnection implements UnidirectionalStream<ServerSentEvent> {
 	public readonly url: string;
 
 	public constructor(init: SseConnectionInit) {
+		this.logger = init.logger;
 		this.url = this.buildUrl(init);
 		this.eventSource = new EventSource(this.url, {
-			fetch: createStreamingFetchAdapter(init.axiosInstance),
+			fetch: createStreamingFetchAdapter(
+				init.axiosInstance,
+				init.optionsHeaders,
+			),
 		});
 		this.setupEventHandlers();
 	}
@@ -58,28 +66,48 @@ export class SseConnection implements UnidirectionalStream<ServerSentEvent> {
 
 	private setupEventHandlers(): void {
 		this.eventSource.addEventListener("open", () =>
-			this.callbacks.open.forEach((cb) => cb({} as WsEvent)),
+			this.invokeCallbacks(this.callbacks.open, {} as WsEvent, "open"),
 		);
 
 		this.eventSource.addEventListener("data", (event: MessageEvent) => {
-			[...this.messageWrappers.values()].forEach((wrapper) => wrapper(event));
+			this.invokeCallbacks(this.messageWrappers.values(), event, "message");
 		});
 
 		this.eventSource.addEventListener("error", (error: Event | ErrorEvent) => {
-			this.callbacks.error.forEach((cb) => cb(this.createErrorEvent(error)));
+			this.invokeCallbacks(
+				this.callbacks.error,
+				this.createErrorEvent(error),
+				"error",
+			);
 
 			if (this.eventSource.readyState === EventSource.CLOSED) {
-				this.callbacks.close.forEach((cb) =>
-					cb({
+				this.invokeCallbacks(
+					this.callbacks.close,
+					{
 						code: 1006,
 						reason: "Connection lost",
 						wasClean: false,
-					} as WsCloseEvent),
+					} as WsCloseEvent,
+					"close",
 				);
 			}
 		});
 	}
 
+	private invokeCallbacks<T>(
+		callbacks: Iterable<(event: T) => void>,
+		event: T,
+		eventType: string,
+	): void {
+		for (const cb of callbacks) {
+			try {
+				cb(event);
+			} catch (err) {
+				this.logger.error(`Error in SSE ${eventType} callback:`, err);
+			}
+		}
+	}
+
 	private createErrorEvent(event: Event | ErrorEvent): WsErrorEvent {
 		const errorMessage =
 			event instanceof ErrorEvent && event.message
@@ -177,12 +205,14 @@ export class SseConnection implements UnidirectionalStream<ServerSentEvent> {
 
 	public close(code?: number, reason?: string): void {
 		this.eventSource.close();
-		this.callbacks.close.forEach((cb) =>
-			cb({
+		this.invokeCallbacks(
+			this.callbacks.close,
+			{
 				code: code ?? 1000,
 				reason: reason ?? "Normal closure",
 				wasClean: true,
-			} as WsCloseEvent),
+			} as WsCloseEvent,
+			"close",
 		);
 
 		Object.values(this.callbacks).forEach((callbackSet) => callbackSet.clear());