diff --git a/src/api-helper.ts b/src/api-helper.ts index f33ae3f8..8e5f3074 100644 --- a/src/api-helper.ts +++ b/src/api-helper.ts @@ -1,12 +1,26 @@ import { Workspace, WorkspaceAgent } from "coder/site/src/api/typesGenerated" import { z } from "zod" +export function errToStr(error: unknown, def: string) { + if (error instanceof Error && error.message) { + return error.message + } + if (typeof error === "string" && error.trim().length > 0) { + return error + } + return def +} + +export function extractAllAgents(workspaces: Workspace[]): WorkspaceAgent[] { + return workspaces.reduce((acc, workspace) => { + return acc.concat(extractAgents(workspace)) + }, [] as WorkspaceAgent[]) +} + export function extractAgents(workspace: Workspace): WorkspaceAgent[] { - const agents = workspace.latest_build.resources.reduce((acc, resource) => { + return workspace.latest_build.resources.reduce((acc, resource) => { return acc.concat(resource.agents || []) }, [] as WorkspaceAgent[]) - - return agents } export const AgentMetadataEventSchema = z.object({ diff --git a/src/workspacesProvider.ts b/src/workspacesProvider.ts index 36d59282..79c4b652 100644 --- a/src/workspacesProvider.ts +++ b/src/workspacesProvider.ts @@ -3,7 +3,13 @@ import { Workspace, WorkspaceAgent } from "coder/site/src/api/typesGenerated" import EventSource from "eventsource" import * as path from "path" import * as vscode from "vscode" -import { AgentMetadataEvent, AgentMetadataEventSchemaArray, extractAgents } from "./api-helper" +import { + AgentMetadataEvent, + AgentMetadataEventSchemaArray, + extractAllAgents, + extractAgents, + errToStr, +} from "./api-helper" import { Storage } from "./storage" export enum WorkspaceQuery { @@ -11,7 +17,12 @@ export enum WorkspaceQuery { All = "", } -type AgentWatcher = { dispose: () => void; metadata?: AgentMetadataEvent[] } +type AgentWatcher = { + onChange: vscode.EventEmitter["event"] + dispose: () => void + metadata?: AgentMetadataEvent[] + error?: unknown +} export class WorkspaceProvider implements vscode.TreeDataProvider { private workspaces: WorkspaceTreeItem[] = [] @@ -39,9 +50,6 @@ export class WorkspaceProvider implements vscode.TreeDataProvider watcher.dispose()) - // It is possible we called fetchAndRefresh() manually (through the button // for example), in which case we might still have a pending refresh that // needs to be cleared. @@ -93,12 +101,38 @@ export class WorkspaceProvider implements vscode.TreeDataProvider { - const showMetadata = this.getWorkspacesQuery === WorkspaceQuery.Mine - if (showMetadata) { - const agents = extractAgents(workspace) - agents.forEach((agent) => this.monitorMetadata(agent.id, url, token2)) // monitor metadata for all agents + const oldWatcherIds = Object.keys(this.agentWatchers) + const reusedWatcherIds: string[] = [] + + // TODO: I think it might make more sense for the tree items to contain + // their own watchers, rather than recreate the tree items every time and + // have this separate map held outside the tree. + const showMetadata = this.getWorkspacesQuery === WorkspaceQuery.Mine + if (showMetadata) { + const agents = extractAllAgents(resp.workspaces) + agents.forEach((agent) => { + // If we have an existing watcher, re-use it. + if (this.agentWatchers[agent.id]) { + reusedWatcherIds.push(agent.id) + return this.agentWatchers[agent.id] + } + // Otherwise create a new watcher. + const watcher = monitorMetadata(agent.id, url, token2) + watcher.onChange(() => this.refresh()) + this.agentWatchers[agent.id] = watcher + return watcher + }) + } + + // Dispose of watchers we ended up not reusing. + oldWatcherIds.forEach((id) => { + if (!reusedWatcherIds.includes(id)) { + this.agentWatchers[id].dispose() + delete this.agentWatchers[id] } + }) + + return resp.workspaces.map((workspace) => { return new WorkspaceTreeItem(workspace, this.getWorkspacesQuery === WorkspaceQuery.All, showMetadata) }) } @@ -157,7 +191,11 @@ export class WorkspaceProvider implements vscode.TreeDataProvider new AgentMetadataTreeItem(metadata))) } @@ -165,53 +203,57 @@ export class WorkspaceProvider implements vscode.TreeDataProvider { - if (!disposed) { - delete this.agentWatchers[agentId] - agentMetadataEventSource.close() - disposed = true - } - }, - } +// 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"], url: string, token: string): AgentWatcher { + const metadataUrl = new URL(`${url}/api/v2/workspaceagents/${agentId}/watch-metadata`) + const eventSource = new EventSource(metadataUrl.toString(), { + headers: { + "Coder-Session-Token": token, + }, + }) + + let disposed = false + const onChange = new vscode.EventEmitter() + const watcher: AgentWatcher = { + onChange: onChange.event, + dispose: () => { + if (!disposed) { + eventSource.close() + disposed = true + } + }, + } - this.agentWatchers[agentId] = watcher + eventSource.addEventListener("data", (event) => { + try { + const dataEvent = JSON.parse(event.data) + const metadata = AgentMetadataEventSchemaArray.parse(dataEvent) - agentMetadataEventSource.addEventListener("data", (event) => { - try { - const dataEvent = JSON.parse(event.data) - const agentMetadata = 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) + } + }) - if (agentMetadata.length === 0) { - watcher.dispose() - } + return watcher +} - // Overwrite metadata if it changed. - if (JSON.stringify(watcher.metadata) !== JSON.stringify(agentMetadata)) { - watcher.metadata = agentMetadata - this.refresh() - } - } catch (error) { - watcher.dispose() - } - }) +class ErrorTreeItem extends vscode.TreeItem { + constructor(error: unknown) { + super("Failed to query metadata: " + errToStr(error, "no error provided"), vscode.TreeItemCollapsibleState.None) + this.contextValue = "coderAgentMetadata" } } -type CoderTreeItemType = "coderWorkspaceSingleAgent" | "coderWorkspaceMultipleAgents" | "coderAgent" - class AgentMetadataTreeItem extends vscode.TreeItem { constructor(metadataEvent: AgentMetadataEvent) { const label = @@ -225,6 +267,8 @@ class AgentMetadataTreeItem extends vscode.TreeItem { } } +type CoderOpenableTreeItemType = "coderWorkspaceSingleAgent" | "coderWorkspaceMultipleAgents" | "coderAgent" + export class OpenableTreeItem extends vscode.TreeItem { constructor( label: string, @@ -236,7 +280,7 @@ export class OpenableTreeItem extends vscode.TreeItem { public readonly workspaceAgent: string | undefined, public readonly workspaceFolderPath: string | undefined, - contextValue: CoderTreeItemType, + contextValue: CoderOpenableTreeItemType, ) { super(label, collapsibleState) this.contextValue = contextValue