Skip to content

Commit 4ba5773

Browse files
EhabYEhabY
committed
Add agent metadata statusbar to monitor resource usage
1 parent 2d7dac8 commit 4ba5773

File tree

3 files changed

+143
-69
lines changed

3 files changed

+143
-69
lines changed

‎src/agentMetadataHelper.ts‎

Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,81 @@
1+
import { Api } from "coder/site/src/api/api";
2+
import { WorkspaceAgent } from "coder/site/src/api/typesGenerated";
3+
import { EventSource } from "eventsource";
4+
import * as vscode from "vscode";
5+
import { createStreamingFetchAdapter } from "./api";
6+
import {
7+
AgentMetadataEvent,
8+
AgentMetadataEventSchemaArray,
9+
errToStr,
10+
} from "./api-helper";
11+
12+
export type AgentMetadataWatcher = {
13+
onChange: vscode.EventEmitter<null>["event"];
14+
dispose: () => void;
15+
metadata?: AgentMetadataEvent[];
16+
error?: unknown;
17+
};
18+
19+
/**
20+
* Opens an SSE connection to watch metadata for a given workspace agent.
21+
* Emits onChange when metadata updates or an error occurs.
22+
*/
23+
export function createAgentMetadataWatcher(
24+
agentId: WorkspaceAgent["id"],
25+
restClient: Api,
26+
): AgentMetadataWatcher {
27+
// TODO: Is there a better way to grab the url and token?
28+
const url = restClient.getAxiosInstance().defaults.baseURL;
29+
const metadataUrl = new URL(
30+
`${url}/api/v2/workspaceagents/${agentId}/watch-metadata`,
31+
);
32+
const eventSource = new EventSource(metadataUrl.toString(), {
33+
fetch: createStreamingFetchAdapter(restClient.getAxiosInstance()),
34+
});
35+
36+
let disposed = false;
37+
const onChange = new vscode.EventEmitter<null>();
38+
const watcher: AgentMetadataWatcher = {
39+
onChange: onChange.event,
40+
dispose: () => {
41+
if (!disposed) {
42+
eventSource.close();
43+
disposed = true;
44+
}
45+
},
46+
};
47+
48+
eventSource.addEventListener("data", (event) => {
49+
try {
50+
const dataEvent = JSON.parse(event.data);
51+
const metadata = AgentMetadataEventSchemaArray.parse(dataEvent);
52+
53+
// Overwrite metadata if it changed.
54+
if (JSON.stringify(watcher.metadata) !== JSON.stringify(metadata)) {
55+
watcher.metadata = metadata;
56+
onChange.fire(null);
57+
}
58+
} catch (error) {
59+
watcher.error = error;
60+
onChange.fire(null);
61+
}
62+
});
63+
64+
return watcher;
65+
}
66+
67+
export function formatMetadataError(error: unknown): string {
68+
return "Failed to query metadata: " + errToStr(error, "no error provided");
69+
}
70+
71+
export function formatEventLabel(metadataEvent: AgentMetadataEvent): string {
72+
return getEventName(metadataEvent) + ": " + getEventValue(metadataEvent);
73+
}
74+
75+
export function getEventName(metadataEvent: AgentMetadataEvent): string {
76+
return metadataEvent.description.display_name.trim();
77+
}
78+
79+
export function getEventValue(metadataEvent: AgentMetadataEvent): string {
80+
return metadataEvent.result.value.replace(/\n/g, "").trim();
81+
}

‎src/remote.ts‎

Lines changed: 51 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import { isAxiosError } from "axios";
22
import { Api } from "coder/site/src/api/api";
3-
import { Workspace } from "coder/site/src/api/typesGenerated";
3+
import { Workspace, WorkspaceAgent } from "coder/site/src/api/typesGenerated";
44
import find from "find-process";
55
import * as fs from "fs/promises";
66
import * as jsonc from "jsonc-parser";
@@ -9,6 +9,12 @@ import * as path from "path";
99
import prettyBytes from "pretty-bytes";
1010
import * as semver from "semver";
1111
import * as vscode from "vscode";
12+
import {
13+
createAgentMetadataWatcher,
14+
getEventValue,
15+
formatEventLabel,
16+
formatMetadataError,
17+
} from "./agentMetadataHelper";
1218
import {
1319
createHttpAgent,
1420
makeCoderSdk,
@@ -624,6 +630,8 @@ export class Remote {
624630
}),
625631
);
626632

633+
this.createAgentMetadataStatusBar(agent, workspaceRestClient, disposables);
634+
627635
this.storage.output.info("Remote setup complete");
628636

629637
// Returning the URL and token allows the plugin to authenticate its own
@@ -966,6 +974,48 @@ export class Remote {
966974
return loop();
967975
}
968976

977+
/**
978+
* Creates and manages a status bar item that displays metadata information for a given workspace agent.
979+
* The status bar item updates dynamically based on changes to the agent's metadata,
980+
* and hides itself if no metadata is available or an error occurs.
981+
*/
982+
private createAgentMetadataStatusBar(
983+
agent: WorkspaceAgent,
984+
restClient: Api,
985+
disposables: vscode.Disposable[],
986+
): void {
987+
const statusBarItem = vscode.window.createStatusBarItem(
988+
"agentMetadata",
989+
vscode.StatusBarAlignment.Left,
990+
);
991+
disposables.push(statusBarItem);
992+
993+
const agentWatcher = createAgentMetadataWatcher(agent.id, restClient);
994+
disposables.push(agentWatcher);
995+
996+
agentWatcher.onChange(
997+
() => {
998+
if (agentWatcher.error) {
999+
this.storage.output.warn(formatMetadataError(agentWatcher.error));
1000+
statusBarItem.hide();
1001+
return;
1002+
}
1003+
1004+
if (agentWatcher.metadata && agentWatcher.metadata.length > 0) {
1005+
statusBarItem.text = getEventValue(agentWatcher.metadata[0]);
1006+
statusBarItem.tooltip = agentWatcher.metadata
1007+
.map((metadata) => formatEventLabel(metadata))
1008+
.join("\n");
1009+
statusBarItem.show();
1010+
} else {
1011+
statusBarItem.hide();
1012+
}
1013+
},
1014+
undefined,
1015+
disposables,
1016+
);
1017+
}
1018+
9691019
// closeRemote ends the current remote session.
9701020
public async closeRemote() {
9711021
await vscode.commands.executeCommand("workbench.action.remote.close");

‎src/workspacesProvider.ts‎

Lines changed: 11 additions & 68 deletions
Original file line numberDiff line numberDiff line change
@@ -4,16 +4,18 @@ import {
44
WorkspaceAgent,
55
WorkspaceApp,
66
} from "coder/site/src/api/typesGenerated";
7-
import { EventSource } from "eventsource";
87
import * as path from "path";
98
import * as vscode from "vscode";
10-
import { createStreamingFetchAdapter } from "./api";
9+
import {
10+
AgentMetadataWatcher,
11+
createAgentMetadataWatcher,
12+
formatEventLabel,
13+
formatMetadataError,
14+
} from "./agentMetadataHelper";
1115
import {
1216
AgentMetadataEvent,
13-
AgentMetadataEventSchemaArray,
1417
extractAllAgents,
1518
extractAgents,
16-
errToStr,
1719
} from "./api-helper";
1820
import { Storage } from "./storage";
1921

@@ -22,13 +24,6 @@ export enum WorkspaceQuery {
2224
All = "",
2325
}
2426

25-
type AgentWatcher = {
26-
onChange: vscode.EventEmitter<null>["event"];
27-
dispose: () => void;
28-
metadata?: AgentMetadataEvent[];
29-
error?: unknown;
30-
};
31-
3227
/**
3328
* Polls workspaces using the provided REST client and renders them in a tree.
3429
*
@@ -42,7 +37,8 @@ export class WorkspaceProvider
4237
{
4338
// Undefined if we have never fetched workspaces before.
4439
private workspaces: WorkspaceTreeItem[] | undefined;
45-
private agentWatchers: Record<WorkspaceAgent["id"], AgentWatcher> = {};
40+
private agentWatchers: Record<WorkspaceAgent["id"], AgentMetadataWatcher> =
41+
{};
4642
private timeout: NodeJS.Timeout | undefined;
4743
private fetching = false;
4844
private visible = false;
@@ -139,7 +135,7 @@ export class WorkspaceProvider
139135
return this.agentWatchers[agent.id];
140136
}
141137
// Otherwise create a new watcher.
142-
const watcher = monitorMetadata(agent.id, restClient);
138+
const watcher = createAgentMetadataWatcher(agent.id, restClient);
143139
watcher.onChange(() => this.refresh());
144140
this.agentWatchers[agent.id] = watcher;
145141
return watcher;
@@ -313,53 +309,6 @@ export class WorkspaceProvider
313309
}
314310
}
315311

316-
// monitorMetadata opens an SSE endpoint to monitor metadata on the specified
317-
// agent and registers a watcher that can be disposed to stop the watch and
318-
// emits an event when the metadata changes.
319-
function monitorMetadata(
320-
agentId: WorkspaceAgent["id"],
321-
restClient: Api,
322-
): AgentWatcher {
323-
// TODO: Is there a better way to grab the url and token?
324-
const url = restClient.getAxiosInstance().defaults.baseURL;
325-
const metadataUrl = new URL(
326-
`${url}/api/v2/workspaceagents/${agentId}/watch-metadata`,
327-
);
328-
const eventSource = new EventSource(metadataUrl.toString(), {
329-
fetch: createStreamingFetchAdapter(restClient.getAxiosInstance()),
330-
});
331-
332-
let disposed = false;
333-
const onChange = new vscode.EventEmitter<null>();
334-
const watcher: AgentWatcher = {
335-
onChange: onChange.event,
336-
dispose: () => {
337-
if (!disposed) {
338-
eventSource.close();
339-
disposed = true;
340-
}
341-
},
342-
};
343-
344-
eventSource.addEventListener("data", (event) => {
345-
try {
346-
const dataEvent = JSON.parse(event.data);
347-
const metadata = AgentMetadataEventSchemaArray.parse(dataEvent);
348-
349-
// Overwrite metadata if it changed.
350-
if (JSON.stringify(watcher.metadata) !== JSON.stringify(metadata)) {
351-
watcher.metadata = metadata;
352-
onChange.fire(null);
353-
}
354-
} catch (error) {
355-
watcher.error = error;
356-
onChange.fire(null);
357-
}
358-
});
359-
360-
return watcher;
361-
}
362-
363312
/**
364313
* A tree item that represents a collapsible section with child items
365314
*/
@@ -375,20 +324,14 @@ class SectionTreeItem extends vscode.TreeItem {
375324

376325
class ErrorTreeItem extends vscode.TreeItem {
377326
constructor(error: unknown) {
378-
super(
379-
"Failed to query metadata: " + errToStr(error, "no error provided"),
380-
vscode.TreeItemCollapsibleState.None,
381-
);
327+
super(formatMetadataError(error), vscode.TreeItemCollapsibleState.None);
382328
this.contextValue = "coderAgentMetadata";
383329
}
384330
}
385331

386332
class AgentMetadataTreeItem extends vscode.TreeItem {
387333
constructor(metadataEvent: AgentMetadataEvent) {
388-
const label =
389-
metadataEvent.description.display_name.trim() +
390-
": " +
391-
metadataEvent.result.value.replace(/\n/g, "").trim();
334+
const label = formatEventLabel(metadataEvent);
392335

393336
super(label, vscode.TreeItemCollapsibleState.None);
394337
const collected_at = new Date(

0 commit comments

Comments
 (0)