Fixed WebSocket connections not receiving headers from the configured header command (coder#619)

Closes coder#618

Author: EhabY

CHANGELOG.md

@@ -2,6 +2,11 @@
 
 ## Unreleased
 
+### Fixed
+
+- Fixed WebSocket connections not receiving headers from the configured header command
+  (`coder.headerCommand`), which could cause authentication failures with remote workspaces.
+
 ## [v1.11.2](https://github.com/coder/vscode-coder/releases/tag/v1.11.2) 2025-10-07
 
 ### Changed

src/agentMetadataHelper.ts src/api/agentMetadataHelper.tssrc/agentMetadataHelper.ts renamed to src/api/agentMetadataHelper.ts

@@ -5,8 +5,8 @@ import {
 	type AgentMetadataEvent,
 	AgentMetadataEventSchemaArray,
 	errToStr,
-} from "./api/api-helper";
-import { type CoderApi } from "./api/coderApi";
+} from "./api-helper";
+import { type CoderApi } from "./coderApi";
 
 export type AgentMetadataWatcher = {
 	onChange: vscode.EventEmitter<null>["event"];
@@ -19,11 +19,11 @@ export type AgentMetadataWatcher = {
  * Opens a websocket connection to watch metadata for a given workspace agent.
  * Emits onChange when metadata updates or an error occurs.
  */
-export function createAgentMetadataWatcher(
+export async function createAgentMetadataWatcher(
 	agentId: WorkspaceAgent["id"],
 	client: CoderApi,
-): AgentMetadataWatcher {
-	const socket = client.watchAgentMetadata(agentId);
+): Promise<AgentMetadataWatcher> {
+	const socket = await client.watchAgentMetadata(agentId);
 
 	let disposed = false;
 	const onChange = new vscode.EventEmitter<null>();

src/api/coderApi.ts

@@ -67,7 +67,7 @@ export class CoderApi extends Api {
 		return client;
 	}
 
-	watchInboxNotifications = (
+	watchInboxNotifications = async (
 		watchTemplates: string[],
 		watchTargets: string[],
 		options?: ClientOptions,
@@ -83,14 +83,14 @@ export class CoderApi extends Api {
 		});
 	};
 
-	watchWorkspace = (workspace: Workspace, options?: ClientOptions) => {
+	watchWorkspace = async (workspace: Workspace, options?: ClientOptions) => {
 		return this.createWebSocket<ServerSentEvent>({
 			apiRoute: `/api/v2/workspaces/${workspace.id}/watch-ws`,
 			options,
 		});
 	};
 
-	watchAgentMetadata = (
+	watchAgentMetadata = async (
 		agentId: WorkspaceAgent["id"],
 		options?: ClientOptions,
 	) => {
@@ -100,21 +100,22 @@ export class CoderApi extends Api {
 		});
 	};
 
-	watchBuildLogsByBuildId = (buildId: string, logs: ProvisionerJobLog[]) => {
+	watchBuildLogsByBuildId = async (
+		buildId: string,
+		logs: ProvisionerJobLog[],
+	) => {
 		const searchParams = new URLSearchParams({ follow: "true" });
 		if (logs.length) {
 			searchParams.append("after", logs[logs.length - 1].id.toString());
 		}
 
-		const socket = this.createWebSocket<ProvisionerJobLog>({
+		return this.createWebSocket<ProvisionerJobLog>({
 			apiRoute: `/api/v2/workspacebuilds/${buildId}/logs`,
 			searchParams,
 		});
-
-		return socket;
 	};
 
-	private createWebSocket<TData = unknown>(
+	private async createWebSocket<TData = unknown>(
 		configs: Omit<OneWayWebSocketInit, "location">,
 	) {
 		const baseUrlRaw = this.getAxiosInstance().defaults.baseURL;
@@ -127,7 +128,15 @@ export class CoderApi extends Api {
 			coderSessionTokenHeader
 		] as string | undefined;
 
-		const httpAgent = createHttpAgent(vscode.workspace.getConfiguration());
+		const headers = await getHeaders(
+			baseUrlRaw,
+			getHeaderCommand(vscode.workspace.getConfiguration()),
+			this.output,
+		);
+
+		const httpAgent = await createHttpAgent(
+			vscode.workspace.getConfiguration(),
+		);
 		const webSocket = new OneWayWebSocket<TData>({
 			location: baseUrl,
 			...configs,
@@ -137,6 +146,7 @@ export class CoderApi extends Api {
 				headers: {
 					...(token ? { [coderSessionTokenHeader]: token } : {}),
 					...configs.options?.headers,
+					...headers,
 				},
 				...configs.options,
 			},
@@ -191,7 +201,7 @@ function setupInterceptors(
 		// Configure proxy and TLS.
 		// Note that by default VS Code overrides the agent. To prevent this, set
 		// `http.proxySupport` to `on` or `off`.
-		const agent = createHttpAgent(vscode.workspace.getConfiguration());
+		const agent = await createHttpAgent(vscode.workspace.getConfiguration());
 		config.httpsAgent = agent;
 		config.httpAgent = agent;
 		config.proxy = false;

src/api/utils.ts

@@ -1,4 +1,4 @@
-import fs from "fs";
+import fs from "fs/promises";
 import { ProxyAgent } from "proxy-agent";
 import { type WorkspaceConfiguration } from "vscode";
 
@@ -23,7 +23,9 @@ export function needToken(cfg: WorkspaceConfiguration): boolean {
  * Create a new HTTP agent based on the current VS Code settings.
  * Configures proxy, TLS certificates, and security options.
  */
-export function createHttpAgent(cfg: WorkspaceConfiguration): ProxyAgent {
+export async function createHttpAgent(
+	cfg: WorkspaceConfiguration,
+): Promise<ProxyAgent> {
 	const insecure = Boolean(cfg.get("coder.insecure"));
 	const certFile = expandPath(
 		String(cfg.get("coder.tlsCertFile") ?? "").trim(),
@@ -32,6 +34,12 @@ export function createHttpAgent(cfg: WorkspaceConfiguration): ProxyAgent {
 	const caFile = expandPath(String(cfg.get("coder.tlsCaFile") ?? "").trim());
 	const altHost = expandPath(String(cfg.get("coder.tlsAltHost") ?? "").trim());
 
+	const [cert, key, ca] = await Promise.all([
+		certFile === "" ? Promise.resolve(undefined) : fs.readFile(certFile),
+		keyFile === "" ? Promise.resolve(undefined) : fs.readFile(keyFile),
+		caFile === "" ? Promise.resolve(undefined) : fs.readFile(caFile),
+	]);
+
 	return new ProxyAgent({
 		// Called each time a request is made.
 		getProxyForUrl: (url: string) => {
@@ -41,9 +49,9 @@ export function createHttpAgent(cfg: WorkspaceConfiguration): ProxyAgent {
 				cfg.get("coder.proxyBypass"),
 			);
 		},
-		cert: certFile === "" ? undefined : fs.readFileSync(certFile),
-		key: keyFile === "" ? undefined : fs.readFileSync(keyFile),
-		ca: caFile === "" ? undefined : fs.readFileSync(caFile),
+		cert,
+		key,
+		ca,
 		servername: altHost === "" ? undefined : altHost,
 		// rejectUnauthorized defaults to true, so we need to explicitly set it to
 		// false if we want to allow self-signed certificates.

src/api/workspace.ts

@@ -95,12 +95,12 @@ export async function waitForBuild(
 	const logs = await client.getWorkspaceBuildLogs(workspace.latest_build.id);
 	logs.forEach((log) => writeEmitter.fire(log.output + "\r\n"));
 
-	await new Promise<void>((resolve, reject) => {
-		const socket = client.watchBuildLogsByBuildId(
-			workspace.latest_build.id,
-			logs,
-		);
+	const socket = await client.watchBuildLogsByBuildId(
+		workspace.latest_build.id,
+		logs,
+	);
 
+	await new Promise<void>((resolve, reject) => {
 		socket.addEventListener("message", (data) => {
 			if (data.parseError) {
 				writeEmitter.fire(

src/headers.ts

@@ -24,7 +24,7 @@ export function getHeaderCommand(
 		config.get<string>("coder.headerCommand")?.trim() ||
 		process.env.CODER_HEADER_COMMAND?.trim();
 
-	return cmd ? cmd : undefined;
+	return cmd || undefined;
 }
 
 export function getHeaderArgs(config: WorkspaceConfiguration): string[] {
@@ -44,16 +44,13 @@ export function getHeaderArgs(config: WorkspaceConfiguration): string[] {
 	return ["--header-command", escapeSubcommand(command)];
 }
 
-// TODO: getHeaders might make more sense to directly implement on Storage
-// but it is difficult to test Storage right now since we use vitest instead of
-// the standard extension testing framework which would give us access to vscode
-// APIs.  We should revert the testing framework then consider moving this.
-
-// getHeaders executes the header command and parses the headers from stdout.
-// Both stdout and stderr are logged on error but stderr is otherwise ignored.
-// Throws an error if the process exits with non-zero or the JSON is invalid.
-// Returns undefined if there is no header command set.  No effort is made to
-// validate the JSON other than making sure it can be parsed.
+/**
+ * getHeaders executes the header command and parses the headers from stdout.
+ * Both stdout and stderr are logged on error but stderr is otherwise ignored.
+ * Throws an error if the process exits with non-zero or the JSON is invalid.
+ * Returns undefined if there is no header command set. No effort is made to
+ * validate the JSON other than making sure it can be parsed.
+ */
 export async function getHeaders(
 	url: string | undefined,
 	command: string | undefined,
@@ -90,8 +87,8 @@ export async function getHeaders(
 			return headers;
 		}
 		const lines = result.stdout.replace(/\r?\n$/, "").split(/\r?\n/);
-		for (let i = 0; i < lines.length; ++i) {
-			const [key, value] = lines[i].split(/=(.*)/);
+		for (const line of lines) {
+			const [key, value] = line.split(/=(.*)/);
 			// Header names cannot be blank or contain whitespace and the Coder CLI
 			// requires that there be an equals sign (the value can be blank though).
 			if (
@@ -100,7 +97,7 @@ export async function getHeaders(
 				typeof value === "undefined"
 			) {
 				throw new Error(
-					`Malformed line from header command: [${lines[i]}] (out: ${result.stdout})`,
+					`Malformed line from header command: [${line}] (out: ${result.stdout})`,
 				);
 			}
 			headers[key] = value;

src/inbox.ts

@@ -16,12 +16,21 @@ const TEMPLATE_WORKSPACE_OUT_OF_MEMORY = "a9d027b4-ac49-4fb1-9f6d-45af15f64e7a";
 const TEMPLATE_WORKSPACE_OUT_OF_DISK = "f047f6a3-5713-40f7-85aa-0394cce9fa3a";
 
 export class Inbox implements vscode.Disposable {
-	readonly #logger: Logger;
-	#disposed = false;
-	#socket: OneWayWebSocket<GetInboxNotificationResponse>;
+	private socket: OneWayWebSocket<GetInboxNotificationResponse> | undefined;
+	private disposed = false;
 
-	constructor(workspace: Workspace, client: CoderApi, logger: Logger) {
-		this.#logger = logger;
+	private constructor(private readonly logger: Logger) {}
+
+	/**
+	 * Factory method to create and initialize an Inbox.
+	 * Use this instead of the constructor to properly handle async websocket initialization.
+	 */
+	static async create(
+		workspace: Workspace,
+		client: CoderApi,
+		logger: Logger,
+	): Promise<Inbox> {
+		const inbox = new Inbox(logger);
 
 		const watchTemplates = [
 			TEMPLATE_WORKSPACE_OUT_OF_DISK,
@@ -30,33 +39,40 @@ export class Inbox implements vscode.Disposable {
 
 		const watchTargets = [workspace.id];
 
-		this.#socket = client.watchInboxNotifications(watchTemplates, watchTargets);
+		const socket = await client.watchInboxNotifications(
+			watchTemplates,
+			watchTargets,
+		);
 
-		this.#socket.addEventListener("open", () => {
-			this.#logger.info("Listening to Coder Inbox");
+		socket.addEventListener("open", () => {
+			logger.info("Listening to Coder Inbox");
 		});
 
-		this.#socket.addEventListener("error", () => {
+		socket.addEventListener("error", () => {
 			// Errors are already logged internally
-			this.dispose();
+			inbox.dispose();
 		});
 
-		this.#socket.addEventListener("message", (data) => {
+		socket.addEventListener("message", (data) => {
 			if (data.parseError) {
-				this.#logger.error("Failed to parse inbox message", data.parseError);
+				logger.error("Failed to parse inbox message", data.parseError);
 			} else {
 				vscode.window.showInformationMessage(
 					data.parsedMessage.notification.title,
 				);
 			}
 		});
+
+		inbox.socket = socket;
+
+		return inbox;
 	}
 
 	dispose() {
-		if (!this.#disposed) {
-			this.#logger.info("No longer listening to Coder Inbox");
-			this.#socket.close();
-			this.#disposed = true;
+		if (!this.disposed) {
+			this.logger.info("No longer listening to Coder Inbox");
+			this.socket?.close();
+			this.disposed = true;
 		}
 	}
 }