Add agent metadata statusbar to monitor resource usage

EhabY · EhabY · commit 7583f2841cf4 · 2025-07-16T23:35:40.000+03:00
diff --git a/src/agentMetadataHelper.ts b/src/agentMetadataHelper.ts
@@ -0,0 +1,81 @@
+import { Api } from "coder/site/src/api/api";
+import { WorkspaceAgent } from "coder/site/src/api/typesGenerated";
+import { EventSource } from "eventsource";
+import * as vscode from "vscode";
+import { createStreamingFetchAdapter } from "./api";
+import {
+	AgentMetadataEvent,
+	AgentMetadataEventSchemaArray,
+	errToStr,
+} from "./api-helper";
+
+export type AgentMetadataWatcher = {
+	onChange: vscode.EventEmitter<null>["event"];
+	dispose: () => void;
+	metadata?: AgentMetadataEvent[];
+	error?: unknown;
+};
+
+/**
+ * Opens an SSE connection to watch metadata for a given workspace agent.
+ * Emits onChange when metadata updates or an error occurs.
+ */
+export function createAgentMetadataWatcher(
+	agentId: WorkspaceAgent["id"],
+	restClient: Api,
+): AgentMetadataWatcher {
+	// TODO: Is there a better way to grab the url and token?
+	const url = restClient.getAxiosInstance().defaults.baseURL;
+	const metadataUrl = new URL(
+		`${url}/api/v2/workspaceagents/${agentId}/watch-metadata`,
+	);
+	const eventSource = new EventSource(metadataUrl.toString(), {
+		fetch: createStreamingFetchAdapter(restClient.getAxiosInstance()),
+	});
+
+	let disposed = false;
+	const onChange = new vscode.EventEmitter<null>();
+	const watcher: AgentMetadataWatcher = {
+		onChange: onChange.event,
+		dispose: () => {
+			if (!disposed) {
+				eventSource.close();
+				disposed = true;
+			}
+		},
+	};
+
+	eventSource.addEventListener("data", (event) => {
+		try {
+			const dataEvent = JSON.parse(event.data);
+			const metadata = AgentMetadataEventSchemaArray.parse(dataEvent);
+
+			// Overwrite metadata if it changed.
+			if (JSON.stringify(watcher.metadata) !== JSON.stringify(metadata)) {
+				watcher.metadata = metadata;
+				onChange.fire(null);
+			}
+		} catch (error) {
+			watcher.error = error;
+			onChange.fire(null);
+		}
+	});
+
+	return watcher;
+}
+
+export function formatMetadataError(error: unknown): string {
+	return "Failed to query metadata: " + errToStr(error, "no error provided");
+}
+
+export function formatEventLabel(metadataEvent: AgentMetadataEvent): string {
+	return getEventName(metadataEvent) + ": " + getEventValue(metadataEvent);
+}
+
+export function getEventName(metadataEvent: AgentMetadataEvent): string {
+	return metadataEvent.description.display_name.trim();
+}
+
+export function getEventValue(metadataEvent: AgentMetadataEvent): string {
+	return metadataEvent.result.value.replace(/\n/g, "").trim();
+}
diff --git a/src/remote.ts b/src/remote.ts
@@ -1,6 +1,6 @@
 import { isAxiosError } from "axios";
 import { Api } from "coder/site/src/api/api";
-import { Workspace } from "coder/site/src/api/typesGenerated";
+import { Workspace, WorkspaceAgent } from "coder/site/src/api/typesGenerated";
 import find from "find-process";
 import * as fs from "fs/promises";
 import * as jsonc from "jsonc-parser";
@@ -9,6 +9,12 @@ import * as path from "path";
 import prettyBytes from "pretty-bytes";
 import * as semver from "semver";
 import * as vscode from "vscode";
+import {
+	createAgentMetadataWatcher,
+	getEventValue,
+	formatEventLabel,
+	formatMetadataError,
+} from "./agentMetadataHelper";
 import {
 	createHttpAgent,
 	makeCoderSdk,
@@ -633,6 +639,8 @@ export class Remote {
 			}),
 		);
 
+		this.createAgentMetadataStatusBar(agent, workspaceRestClient, disposables);
+
 		this.storage.writeToCoderOutputChannel("Remote setup complete");
 
 		// Returning the URL and token allows the plugin to authenticate its own
@@ -974,6 +982,50 @@ export class Remote {
 		return loop();
 	}
 
+	/**
+	 * Creates and manages a status bar item that displays metadata information for a given workspace agent.
+	 * 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(
+		agent: WorkspaceAgent,
+		restClient: Api,
+		disposables: vscode.Disposable[],
+	): void {
+		const statusBarItem = vscode.window.createStatusBarItem(
+			"agentMetadata",
+			vscode.StatusBarAlignment.Left,
+		);
+		disposables.push(statusBarItem);
+
+		const agentWatcher = createAgentMetadataWatcher(agent.id, restClient);
+		disposables.push(agentWatcher);
+
+		agentWatcher.onChange(
+			() => {
+				if (agentWatcher.error) {
+					this.storage.writeToCoderOutputChannel(
+						formatMetadataError(agentWatcher.error),
+					);
+					statusBarItem.hide();
+					return;
+				}
+
+				if (agentWatcher.metadata && agentWatcher.metadata.length > 0) {
+					statusBarItem.text = getEventValue(agentWatcher.metadata[0]);
+					statusBarItem.tooltip = agentWatcher.metadata
+						.map((metadata) => formatEventLabel(metadata))
+						.join("\n");
+					statusBarItem.show();
+				} else {
+					statusBarItem.hide();
+				}
+			},
+			undefined,
+			disposables,
+		);
+	}
+
 	// closeRemote ends the current remote session.
 	public async closeRemote() {
 		await vscode.commands.executeCommand("workbench.action.remote.close");
diff --git a/src/workspacesProvider.ts b/src/workspacesProvider.ts
@@ -4,16 +4,18 @@ import {
 	WorkspaceAgent,
 	WorkspaceApp,
 } from "coder/site/src/api/typesGenerated";
-import { EventSource } from "eventsource";
 import * as path from "path";
 import * as vscode from "vscode";
-import { createStreamingFetchAdapter } from "./api";
+import {
+	AgentMetadataWatcher,
+	createAgentMetadataWatcher,
+	formatEventLabel,
+	formatMetadataError,
+} from "./agentMetadataHelper";
 import {
 	AgentMetadataEvent,
-	AgentMetadataEventSchemaArray,
 	extractAllAgents,
 	extractAgents,
-	errToStr,
 } from "./api-helper";
 import { Storage } from "./storage";
 
@@ -22,13 +24,6 @@ export enum WorkspaceQuery {
 	All = "",
 }
 
-type AgentWatcher = {
-	onChange: vscode.EventEmitter<null>["event"];
-	dispose: () => void;
-	metadata?: AgentMetadataEvent[];
-	error?: unknown;
-};
-
 /**
  * Polls workspaces using the provided REST client and renders them in a tree.
  *
@@ -42,7 +37,8 @@ export class WorkspaceProvider
 {
 	// Undefined if we have never fetched workspaces before.
 	private workspaces: WorkspaceTreeItem[] | undefined;
-	private agentWatchers: Record<WorkspaceAgent["id"], AgentWatcher> = {};
+	private agentWatchers: Record<WorkspaceAgent["id"], AgentMetadataWatcher> =
+		{};
 	private timeout: NodeJS.Timeout | undefined;
 	private fetching = false;
 	private visible = false;
@@ -139,7 +135,7 @@ export class WorkspaceProvider
 					return this.agentWatchers[agent.id];
 				}
 				// Otherwise create a new watcher.
-				const watcher = monitorMetadata(agent.id, restClient);
+				const watcher = createAgentMetadataWatcher(agent.id, restClient);
 				watcher.onChange(() => this.refresh());
 				this.agentWatchers[agent.id] = watcher;
 				return watcher;
@@ -313,53 +309,6 @@ export class WorkspaceProvider
 	}
 }
 
-// monitorMetadata opens an SSE endpoint to monitor metadata on the specified
-// agent and registers a watcher that can be disposed to stop the watch and
-// emits an event when the metadata changes.
-function monitorMetadata(
-	agentId: WorkspaceAgent["id"],
-	restClient: Api,
-): AgentWatcher {
-	// TODO: Is there a better way to grab the url and token?
-	const url = restClient.getAxiosInstance().defaults.baseURL;
-	const metadataUrl = new URL(
-		`${url}/api/v2/workspaceagents/${agentId}/watch-metadata`,
-	);
-	const eventSource = new EventSource(metadataUrl.toString(), {
-		fetch: createStreamingFetchAdapter(restClient.getAxiosInstance()),
-	});
-
-	let disposed = false;
-	const onChange = new vscode.EventEmitter<null>();
-	const watcher: AgentWatcher = {
-		onChange: onChange.event,
-		dispose: () => {
-			if (!disposed) {
-				eventSource.close();
-				disposed = true;
-			}
-		},
-	};
-
-	eventSource.addEventListener("data", (event) => {
-		try {
-			const dataEvent = JSON.parse(event.data);
-			const metadata = AgentMetadataEventSchemaArray.parse(dataEvent);
-
-			// Overwrite metadata if it changed.
-			if (JSON.stringify(watcher.metadata) !== JSON.stringify(metadata)) {
-				watcher.metadata = metadata;
-				onChange.fire(null);
-			}
-		} catch (error) {
-			watcher.error = error;
-			onChange.fire(null);
-		}
-	});
-
-	return watcher;
-}
-
 /**
  * A tree item that represents a collapsible section with child items
  */
@@ -375,20 +324,14 @@ class SectionTreeItem extends vscode.TreeItem {
 
 class ErrorTreeItem extends vscode.TreeItem {
 	constructor(error: unknown) {
-		super(
-			"Failed to query metadata: " + errToStr(error, "no error provided"),
-			vscode.TreeItemCollapsibleState.None,
-		);
+		super(formatMetadataError(error), vscode.TreeItemCollapsibleState.None);
 		this.contextValue = "coderAgentMetadata";
 	}
 }
 
 class AgentMetadataTreeItem extends vscode.TreeItem {
 	constructor(metadataEvent: AgentMetadataEvent) {
-		const label =
-			metadataEvent.description.display_name.trim() +
-			": " +
-			metadataEvent.result.value.replace(/\n/g, "").trim();
+		const label = formatEventLabel(metadataEvent);
 
 		super(label, vscode.TreeItemCollapsibleState.None);
 		const collected_at = new Date(
