Unify WS implementations across VS Code

Author: EhabY

src/agentMetadataHelper.ts

@@ -1,13 +1,11 @@
-import { Api } from "coder/site/src/api/api";
 import { WorkspaceAgent } from "coder/site/src/api/typesGenerated";
-import { ProxyAgent } from "proxy-agent";
 import * as vscode from "vscode";
 import {
 	AgentMetadataEvent,
 	AgentMetadataEventSchemaArray,
 	errToStr,
 } from "./api-helper";
-import { watchAgentMetadata } from "./websocket/ws-helper";
+import { CoderWebSocketClient } from "./websocket/webSocketClient";
 
 export type AgentMetadataWatcher = {
 	onChange: vscode.EventEmitter<null>["event"];
@@ -22,10 +20,9 @@ export type AgentMetadataWatcher = {
  */
 export function createAgentMetadataWatcher(
 	agentId: WorkspaceAgent["id"],
-	restClient: Api,
-	httpAgent: ProxyAgent,
+	webSocketClient: CoderWebSocketClient,
 ): AgentMetadataWatcher {
-	const socket = watchAgentMetadata(restClient, httpAgent, agentId);
+	const socket = webSocketClient.watchAgentMetadata(agentId);
 
 	let disposed = false;
 	const onChange = new vscode.EventEmitter<null>();

src/api.ts

@@ -1,21 +1,18 @@
 import { AxiosInstance, InternalAxiosRequestConfig, isAxiosError } from "axios";
 import { spawn } from "child_process";
 import { Api } from "coder/site/src/api/api";
-import {
-	ProvisionerJobLog,
-	Workspace,
-} from "coder/site/src/api/typesGenerated";
+import { Workspace } from "coder/site/src/api/typesGenerated";
 import fs from "fs/promises";
 import { ProxyAgent } from "proxy-agent";
 import * as vscode from "vscode";
-import * as ws from "ws";
 import { errToStr } from "./api-helper";
 import { CertificateError } from "./error";
 import { FeatureSet } from "./featureSet";
 import { getHeaderArgs } from "./headers";
 import { getProxyForUrl } from "./proxy";
 import { Storage } from "./storage";
 import { expandPath } from "./util";
+import { CoderWebSocketClient } from "./websocket/webSocketClient";
 
 export const coderSessionTokenHeader = "Coder-Session-Token";
 
@@ -68,8 +65,12 @@ export async function createHttpAgent(): Promise<ProxyAgent> {
 
 /**
  * Create an sdk instance using the provided URL and token and hook it up to
- * configuration.  The token may be undefined if some other form of
+ * configuration. The token may be undefined if some other form of
  * authentication is being used.
+ *
+ * Automatically configures logging interceptors that log:
+ * - Requests and responses at the trace level
+ * - Errors at the error level
  */
 export function makeCoderSdk(
 	baseUrl: string,
@@ -128,14 +129,14 @@ function addLoggingInterceptors(
 ) {
 	client.interceptors.request.use(
 		(config) => {
-			const requestId = crypto.randomUUID();
+			const requestId = crypto.randomUUID().replace(/-/g, "");
 			(config as RequestConfigWithMetadata).metadata = {
 				requestId,
 				startedAt: Date.now(),
 			};
 
 			logger.trace(
-				`Request ${requestId}: ${config.method?.toUpperCase()} ${config.url}`,
+				`req ${requestId}: ${config.method?.toUpperCase()} ${config.url}`,
 				config.data ?? "",
 			);
 
@@ -146,9 +147,9 @@ function addLoggingInterceptors(
 			if (isAxiosError(error)) {
 				const meta = (error.config as RequestConfigWithMetadata)?.metadata;
 				const requestId = meta?.requestId ?? "n/a";
-				message = `Request ${requestId} error`;
+				message = `req ${requestId} error`;
 			}
-			logger.warn(message, error);
+			logger.error(message, error);
 
 			return Promise.reject(error);
 		},
@@ -161,7 +162,7 @@ function addLoggingInterceptors(
 			const ms = startedAt ? Date.now() - startedAt : undefined;
 
 			logger.trace(
-				`Response ${requestId ?? "n/a"}: ${response.status}${
+				`res ${requestId ?? "n/a"}: ${response.status}${
 					ms !== undefined ? ` in ${ms}ms` : ""
 				}`,
 				// { responseBody: response.data }, // TODO too noisy
@@ -177,7 +178,7 @@ function addLoggingInterceptors(
 				const ms = startedAt ? Date.now() - startedAt : undefined;
 
 				const status = error.response?.status ?? "unknown";
-				message = `Response ${requestId}: ${status}${ms !== undefined ? ` in ${ms}ms` : ""}`;
+				message = `res ${requestId}: ${status}${ms !== undefined ? ` in ${ms}ms` : ""}`;
 			}
 			logger.warn(message, error);
 
@@ -264,70 +265,47 @@ export async function startWorkspaceIfStoppedOrFailed(
  */
 export async function waitForBuild(
 	restClient: Api,
+	webSocketClient: CoderWebSocketClient,
 	writeEmitter: vscode.EventEmitter<string>,
 	workspace: Workspace,
 ): Promise<Workspace> {
-	const baseUrlRaw = restClient.getAxiosInstance().defaults.baseURL;
-	if (!baseUrlRaw) {
-		throw new Error("No base URL set on REST client");
-	}
-
 	// This fetches the initial bunch of logs.
 	const logs = await restClient.getWorkspaceBuildLogs(
 		workspace.latest_build.id,
 	);
 	logs.forEach((log) => writeEmitter.fire(log.output + "\r\n"));
 
-	// This follows the logs for new activity!
-	// TODO: watchBuildLogsByBuildId exists, but it uses `location`.
-	//       Would be nice if we could use it here.
-	let path = `/api/v2/workspacebuilds/${workspace.latest_build.id}/logs?follow=true`;
-	if (logs.length) {
-		path += `&after=${logs[logs.length - 1].id}`;
-	}
-
-	const agent = await createHttpAgent();
 	await new Promise<void>((resolve, reject) => {
+		const rejectError = (error: unknown) => {
+			const baseUrlRaw = restClient.getAxiosInstance().defaults.baseURL!;
+			return reject(
+				new Error(
+					`Failed to watch workspace build on ${baseUrlRaw}: ${errToStr(error, "no further details")}`,
+				),
+			);
+		};
+
 		try {
-			// TODO move to `ws-helper`
-			const baseUrl = new URL(baseUrlRaw);
-			const proto = baseUrl.protocol === "https:" ? "wss:" : "ws:";
-			const socketUrlRaw = `${proto}//${baseUrl.host}${path}`;
-			const token = restClient.getAxiosInstance().defaults.headers.common[
-				coderSessionTokenHeader
-			] as string | undefined;
-			const socket = new ws.WebSocket(new URL(socketUrlRaw), {
-				agent: agent,
-				followRedirects: true,
-				headers: token
-					? {
-							[coderSessionTokenHeader]: token,
-						}
-					: undefined,
-			});
-			socket.binaryType = "nodebuffer";
-			socket.on("message", (data) => {
-				const buf = data as Buffer;
-				const log = JSON.parse(buf.toString()) as ProvisionerJobLog;
+			const socket = webSocketClient.watchBuildLogsByBuildId(
+				workspace.latest_build.id,
+				logs,
+			);
+			const closeHandler = () => {
+				resolve();
+			};
+			socket.addEventListener("close", closeHandler);
+			socket.addEventListener("message", (data) => {
+				const log = data.parsedMessage!;
 				writeEmitter.fire(log.output + "\r\n");
 			});
-			socket.on("error", (error) => {
-				reject(
-					new Error(
-						`Failed to watch workspace build using ${socketUrlRaw}: ${errToStr(error, "no further details")}`,
-					),
-				);
-			});
-			socket.on("close", () => {
-				resolve();
+			socket.addEventListener("error", (error) => {
+				socket.removeEventListener("close", closeHandler);
+				socket.close();
+				rejectError(error);
 			});
 		} catch (error) {
 			// If this errors, it is probably a malformed URL.
-			reject(
-				new Error(
-					`Failed to watch workspace build on ${baseUrlRaw}: ${errToStr(error, "no further details")}`,
-				),
-			);
+			reject(rejectError);
 		}
 	});
 

src/extension.ts

@@ -10,6 +10,7 @@ import { CertificateError, getErrorDetail } from "./error";
 import { Remote } from "./remote";
 import { Storage } from "./storage";
 import { toSafeHost } from "./util";
+import { CoderWebSocketClient } from "./websocket/webSocketClient";
 import { WorkspaceQuery, WorkspaceProvider } from "./workspacesProvider";
 
 export async function activate(ctx: vscode.ExtensionContext): Promise<void> {
@@ -69,18 +70,23 @@ export async function activate(ctx: vscode.ExtensionContext): Promise<void> {
 
 	// TODO this won't get updated when users change their settings; Listen to changes and update this
 	const httpAgent = await createHttpAgent();
+	const webSocketClient = new CoderWebSocketClient(
+		restClient,
+		httpAgent,
+		storage,
+	);
 	const myWorkspacesProvider = new WorkspaceProvider(
 		WorkspaceQuery.Mine,
 		restClient,
 		storage,
-		httpAgent,
+		webSocketClient,
 		5,
 	);
 	const allWorkspacesProvider = new WorkspaceProvider(
 		WorkspaceQuery.All,
 		restClient,
 		storage,
-		httpAgent,
+		webSocketClient,
 	);
 
 	// createTreeView, unlike registerTreeDataProvider, gives us the tree view API

src/inbox.ts

@@ -1,14 +1,12 @@
-import { Api } from "coder/site/src/api/api";
 import {
 	Workspace,
 	GetInboxNotificationResponse,
 } from "coder/site/src/api/typesGenerated";
-import { ProxyAgent } from "proxy-agent";
 import * as vscode from "vscode";
 import { errToStr } from "./api-helper";
 import { type Storage } from "./storage";
-import { OneWayCodeWebSocket } from "./websocket/OneWayCodeWebSocket";
-import { watchInboxNotifications } from "./websocket/ws-helper";
+import { OneWayCodeWebSocket } from "./websocket/oneWayCodeWebSocket";
+import { CoderWebSocketClient } from "./websocket/webSocketClient";
 
 // These are the template IDs of our notifications.
 // Maybe in the future we should avoid hardcoding
@@ -23,8 +21,7 @@ export class Inbox implements vscode.Disposable {
 
 	constructor(
 		workspace: Workspace,
-		httpAgent: ProxyAgent,
-		restClient: Api,
+		webSocketClient: CoderWebSocketClient,
 		storage: Storage,
 	) {
 		this.#storage = storage;
@@ -36,9 +33,7 @@ export class Inbox implements vscode.Disposable {
 
 		const watchTargets = [workspace.id];
 
-		this.#socket = watchInboxNotifications(
-			restClient,
-			httpAgent,
+		this.#socket = webSocketClient.watchInboxNotifications(
 			watchTemplates,
 			watchTargets,
 		);
@@ -47,8 +42,8 @@ export class Inbox implements vscode.Disposable {
 			this.#storage.output.info("Listening to Coder Inbox");
 		});
 
-		this.#socket.addEventListener("error", (error) => {
-			this.notifyError(error);
+		this.#socket.addEventListener("error", () => {
+			// Errors are already logged internally
 			this.dispose();
 		});
 

src/remote.ts

@@ -7,7 +7,6 @@ import * as jsonc from "jsonc-parser";
 import * as os from "os";
 import * as path from "path";
 import prettyBytes from "pretty-bytes";
-import { ProxyAgent } from "proxy-agent";
 import * as semver from "semver";
 import * as vscode from "vscode";
 import {
@@ -39,6 +38,7 @@ import {
 	findPort,
 	parseRemoteAuthority,
 } from "./util";
+import { CoderWebSocketClient } from "./websocket/webSocketClient";
 import { WorkspaceMonitor } from "./workspaceMonitor";
 
 export interface RemoteDetails extends vscode.Disposable {
@@ -493,27 +493,27 @@ export class Remote {
 		}
 
 		const httpAgent = await createHttpAgent();
+		const webSocketClient = new CoderWebSocketClient(
+			workspaceRestClient,
+			httpAgent,
+			this.storage,
+		);
 
 		// Watch the workspace for changes.
 		const monitor = new WorkspaceMonitor(
 			workspace,
 			workspaceRestClient,
 			this.storage,
 			this.vscodeProposed,
-			httpAgent,
+			webSocketClient,
 		);
 		disposables.push(monitor);
 		disposables.push(
 			monitor.onChange.event((w) => (this.commands.workspace = w)),
 		);
 
 		// Watch coder inbox for messages
-		const inbox = new Inbox(
-			workspace,
-			httpAgent,
-			workspaceRestClient,
-			this.storage,
-		);
+		const inbox = new Inbox(workspace, webSocketClient, this.storage);
 		disposables.push(inbox);
 
 		// Wait for the agent to connect.
@@ -631,11 +631,7 @@ export class Remote {
 		);
 
 		disposables.push(
-			...this.createAgentMetadataStatusBar(
-				agent,
-				workspaceRestClient,
-				httpAgent,
-			),
+			...this.createAgentMetadataStatusBar(agent, webSocketClient),
 		);
 
 		this.storage.output.info("Remote setup complete");
@@ -987,19 +983,14 @@ export class Remote {
 	 */
 	private createAgentMetadataStatusBar(
 		agent: WorkspaceAgent,
-		restClient: Api,
-		httpAgent: ProxyAgent,
+		webSocketClient: CoderWebSocketClient,
 	): vscode.Disposable[] {
 		const statusBarItem = vscode.window.createStatusBarItem(
 			"agentMetadata",
 			vscode.StatusBarAlignment.Left,
 		);
 
-		const agentWatcher = createAgentMetadataWatcher(
-			agent.id,
-			restClient,
-			httpAgent,
-		);
+		const agentWatcher = createAgentMetadataWatcher(agent.id, webSocketClient);
 
 		const onChangeDisposable = agentWatcher.onChange(() => {
 			if (agentWatcher.error) {