Reuse agent metadata watchers

This has two effects:

1. Lower overhead from establishing the watchers over and over.
2. No more flashing in the tree as the watchers get recreated and we
   wait for the first update.

Author: code-asher

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({

src/workspacesProvider.ts

@@ -3,15 +3,26 @@ 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 {
   Mine = "owner:me",
   All = "",
 }
 
-type AgentWatcher = { dispose: () => void; metadata?: AgentMetadataEvent[] }
+type AgentWatcher = {
+  onChange: vscode.EventEmitter<null>["event"]
+  dispose: () => void
+  metadata?: AgentMetadataEvent[]
+  error?: unknown
+}
 
 export class WorkspaceProvider implements vscode.TreeDataProvider<vscode.TreeItem> {
   private workspaces: WorkspaceTreeItem[] = []
@@ -39,9 +50,6 @@ export class WorkspaceProvider implements vscode.TreeDataProvider<vscode.TreeIte
     }
     this.fetching = true
 
-    // TODO: It would be better to reuse these.
-    Object.values(this.agentWatchers).forEach((watcher) => 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,37 @@ export class WorkspaceProvider implements vscode.TreeDataProvider<vscode.TreeIte
       return this.fetch()
     }
 
-    return resp.workspaces.map((workspace) => {
-      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()
       }
+    })
+
+    return resp.workspaces.map((workspace) => {
       return new WorkspaceTreeItem(workspace, this.getWorkspacesQuery === WorkspaceQuery.All, showMetadata)
     })
   }
@@ -157,61 +190,71 @@ export class WorkspaceProvider implements vscode.TreeDataProvider<vscode.TreeIte
         )
         return Promise.resolve(agentTreeItems)
       } else if (element instanceof AgentTreeItem) {
-        const savedMetadata = this.agentWatchers[element.agent.id]?.metadata || []
+        const watcher = this.agentWatchers[element.agent.id]
+        if (watcher?.error) {
+          return Promise.resolve([new ErrorTreeItem(watcher.error)])
+        }
+        const savedMetadata = watcher?.metadata || []
         return Promise.resolve(savedMetadata.map((metadata) => new AgentMetadataTreeItem(metadata)))
       }
 
       return Promise.resolve([])
     }
     return Promise.resolve(this.workspaces)
   }
+}
 
-  // monitorMetadata opens an SSE endpoint to monitor metadata on the specified
-  // agent and registers a disposer that can be used to stop the watch.
-  monitorMetadata(agentId: WorkspaceAgent["id"], url: string, token: string): void {
-    const agentMetadataURL = new URL(`${url}/api/v2/workspaceagents/${agentId}/watch-metadata`)
-    const agentMetadataEventSource = new EventSource(agentMetadataURL.toString(), {
-      headers: {
-        "Coder-Session-Token": token,
-      },
-    })
+// 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<null>()
+  const watcher: AgentWatcher = {
+    onChange: onChange.event,
+    dispose: () => {
+      if (!disposed) {
+        eventSource.close()
+        disposed = true
+      }
+    },
+  }
 
-    let disposed = false
-    const watcher: AgentWatcher = {
-      dispose: () => {
-        if (!disposed) {
-          delete this.agentWatchers[agentId]
-          agentMetadataEventSource.close()
-          disposed = true
-        }
-      },
-    }
+  eventSource.addEventListener("data", (event) => {
+    try {
+      const dataEvent = JSON.parse(event.data)
+      const metadata = AgentMetadataEventSchemaArray.parse(dataEvent)
 
-    this.agentWatchers[agentId] = watcher
+      // 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)
+    }
+  })
 
-    agentMetadataEventSource.addEventListener("data", (event) => {
-      try {
-        const dataEvent = JSON.parse(event.data)
-        const agentMetadata = AgentMetadataEventSchemaArray.parse(dataEvent)
+  return watcher
+}
 
-        if (agentMetadata.length === 0) {
-          watcher.dispose()
-        }
+type CoderTreeItemType = "coderWorkspaceSingleAgent" | "coderWorkspaceMultipleAgents" | "coderAgent"
 
-        // 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 =