Pass the header command output to WebSocket creation

Author: EhabY

src/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,14 +128,23 @@ 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,
 			options: {
 				agent: httpAgent,
 				followRedirects: true,
 				headers: {
+					...headers,
 					...(token ? { [coderSessionTokenHeader]: token } : {}),
 					...configs.options?.headers,
 				},
@@ -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/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;
 		}
 	}
 }

src/remote/remote.ts

@@ -18,7 +18,7 @@ import {
 	getEventValue,
 	formatEventLabel,
 	formatMetadataError,
-} from "../agentMetadataHelper";
+} from "../api/agentMetadataHelper";
 import { createWorkspaceIdentifier, extractAgents } from "../api/api-helper";
 import { CoderApi } from "../api/coderApi";
 import { needToken } from "../api/utils";
@@ -543,7 +543,7 @@ export class Remote {
 			}
 
 			// Watch the workspace for changes.
-			const monitor = new WorkspaceMonitor(
+			const monitor = await WorkspaceMonitor.create(
 				workspace,
 				workspaceClient,
 				this.logger,
@@ -556,7 +556,7 @@ export class Remote {
 			);
 
 			// Watch coder inbox for messages
-			const inbox = new Inbox(workspace, workspaceClient, this.logger);
+			const inbox = await Inbox.create(workspace, workspaceClient, this.logger);
 			disposables.push(inbox);
 
 			// Wait for the agent to connect.
@@ -668,7 +668,7 @@ export class Remote {
 						agent.name,
 					);
 				}),
-				...this.createAgentMetadataStatusBar(agent, workspaceClient),
+				...(await this.createAgentMetadataStatusBar(agent, workspaceClient)),
 			);
 		} catch (ex) {
 			// Whatever error happens, make sure we clean up the disposables in case of failure
@@ -858,8 +858,7 @@ export class Remote {
 			"UserKnownHostsFile",
 			"StrictHostKeyChecking",
 		];
-		for (let i = 0; i < keysToMatch.length; i++) {
-			const key = keysToMatch[i];
+		for (const key of keysToMatch) {
 			if (computedProperties[key] === sshValues[key]) {
 				continue;
 			}
@@ -1005,7 +1004,7 @@ export class Remote {
 			// this to find the SSH process that is powering this connection. That SSH
 			// process will be logging network information periodically to a file.
 			const text = await fs.readFile(logPath, "utf8");
-			const port = await findPort(text);
+			const port = findPort(text);
 			if (!port) {
 				return;
 			}
@@ -1064,16 +1063,16 @@ export class Remote {
 	 * The status bar item updates dynamically based on changes to the agent's metadata,
 	 * and hides itself if no metadata is available or an error occurs.
 	 */
-	private createAgentMetadataStatusBar(
+	private async createAgentMetadataStatusBar(
 		agent: WorkspaceAgent,
 		client: CoderApi,
-	): vscode.Disposable[] {
+	): Promise<vscode.Disposable[]> {
 		const statusBarItem = vscode.window.createStatusBarItem(
 			"agentMetadata",
 			vscode.StatusBarAlignment.Left,
 		);
 
-		const agentWatcher = createAgentMetadataWatcher(agent.id, client);
+		const agentWatcher = await createAgentMetadataWatcher(agent.id, client);
 
 		const onChangeDisposable = agentWatcher.onChange(() => {
 			if (agentWatcher.error) {