Show agent chooser when there are multiple

Author: code-asher

src/commands.ts

@@ -27,6 +27,52 @@ export class Commands {
     private readonly storage: Storage,
   ) {}
 
+  /**
+   * Find the requested agent if specified, otherwise return the agent if there
+   * is only one or ask the user to pick if there are multiple.  Return
+   * undefined if the user cancels.
+   */
+  public async maybeAskAgent(workspace: Workspace, filter?: string): Promise<WorkspaceAgent | undefined> {
+    const agents = extractAgents(workspace)
+    const filteredAgents = filter ? agents.filter((agent) => agent.name === filter) : agents
+    if (filteredAgents.length === 0) {
+      throw new Error("Workspace has no matching agents")
+    } else if (filteredAgents.length === 1) {
+      return filteredAgents[0]
+    } else {
+      const quickPick = vscode.window.createQuickPick()
+      quickPick.title = "Select an agent"
+      quickPick.busy = true
+      const agentItems: vscode.QuickPickItem[] = filteredAgents.map((agent) => {
+        let icon = "$(debug-start)"
+        if (agent.status !== "connected") {
+          icon = "$(debug-stop)"
+        }
+        return {
+          alwaysShow: true,
+          label: `${icon} ${agent.name}`,
+          detail: `${agent.name} • Status: ${agent.status}`,
+        }
+      })
+      quickPick.items = agentItems
+      quickPick.busy = false
+      quickPick.show()
+
+      const selected = await new Promise<WorkspaceAgent | undefined>((resolve) => {
+        quickPick.onDidHide(() => resolve(undefined))
+        quickPick.onDidChangeSelection((selected) => {
+          if (selected.length < 1) {
+            return resolve(undefined)
+          }
+          const agent = filteredAgents[quickPick.items.indexOf(selected[0])]
+          resolve(agent)
+        })
+      })
+      quickPick.dispose()
+      return selected
+    }
+  }
+
   /**
    * Ask the user for the URL, letting them choose from a list of recent URLs or
    * CODER_URL or enter a new one.  Undefined means the user aborted.
@@ -376,58 +422,19 @@ export class Commands {
         })
       })
       if (!workspace) {
+        // User declined to pick a workspace.
         return
       }
       workspaceOwner = workspace.owner_name
       workspaceName = workspace.name
 
-      const agents = extractAgents(workspace)
-
-      if (agents.length === 1) {
-        folderPath = agents[0].expanded_directory
-        workspaceAgent = agents[0].name
-      } else if (agents.length > 0) {
-        const agentQuickPick = vscode.window.createQuickPick()
-        agentQuickPick.title = `Select an agent`
-
-        agentQuickPick.busy = true
-        const lastAgents = agents
-        const agentItems: vscode.QuickPickItem[] = agents.map((agent) => {
-          let icon = "$(debug-start)"
-          if (agent.status !== "connected") {
-            icon = "$(debug-stop)"
-          }
-          return {
-            alwaysShow: true,
-            label: `${icon} ${agent.name}`,
-            detail: `${agent.name} • Status: ${agent.status}`,
-          }
-        })
-        agentQuickPick.items = agentItems
-        agentQuickPick.busy = false
-        agentQuickPick.show()
-
-        const agent = await new Promise<WorkspaceAgent | undefined>((resolve) => {
-          agentQuickPick.onDidHide(() => {
-            resolve(undefined)
-          })
-          agentQuickPick.onDidChangeSelection((selected) => {
-            if (selected.length < 1) {
-              return resolve(undefined)
-            }
-            const agent = lastAgents[agentQuickPick.items.indexOf(selected[0])]
-            resolve(agent)
-          })
-        })
-
-        if (agent) {
-          folderPath = agent.expanded_directory
-          workspaceAgent = agent.name
-        } else {
-          folderPath = ""
-          workspaceAgent = ""
-        }
+      const agent = await this.maybeAskAgent(workspace)
+      if (!agent) {
+        // User declined to pick an agent.
+        return
       }
+      folderPath = agent.expanded_directory
+      workspaceAgent = agent.name
     } else {
       workspaceOwner = args[0] as string
       workspaceName = args[1] as string

src/remote.ts

@@ -1,6 +1,6 @@
 import { isAxiosError } from "axios"
 import { Api } from "coder/site/src/api/api"
-import { Workspace, WorkspaceAgent } from "coder/site/src/api/typesGenerated"
+import { Workspace } from "coder/site/src/api/typesGenerated"
 import EventSource from "eventsource"
 import find from "find-process"
 import * as fs from "fs/promises"
@@ -11,6 +11,7 @@ import prettyBytes from "pretty-bytes"
 import * as semver from "semver"
 import * as vscode from "vscode"
 import { makeCoderSdk, startWorkspace, waitForBuild } from "./api"
+import { extractAgents } from "./api-helper"
 import { Commands } from "./commands"
 import { getHeaderCommand } from "./headers"
 import { SSHConfig, SSHValues, mergeSSHConfigValues } from "./sshConfig"
@@ -287,28 +288,13 @@ export class Remote {
 
     // Pick an agent.
     this.storage.writeToCoderOutputChannel(`Finding agent for ${workspaceName}...`)
-    const agents = workspace.latest_build.resources.reduce((acc, resource) => {
-      return acc.concat(resource.agents || [])
-    }, [] as WorkspaceAgent[])
-
-    // With no agent specified, pick the first one.  Otherwise choose the
-    // matching agent.
-    let agent: WorkspaceAgent
-    if (!parts.agent) {
-      if (agents.length === 1) {
-        agent = agents[0]
-      } else {
-        // TODO: Show the agent selector here instead.
-        throw new Error("Invalid Coder SSH authority. An agent must be specified when there are multiple.")
-      }
-    } else {
-      const matchingAgents = agents.filter((agent) => agent.name === parts.agent)
-      if (matchingAgents.length !== 1) {
-        // TODO: Show the agent selector here instead.
-        throw new Error("Invalid Coder SSH authority. Agent not found.")
-      }
-      agent = matchingAgents[0]
+    const gotAgent = await this.commands.maybeAskAgent(workspace, parts.agent)
+    if (!gotAgent) {
+      // User declined to pick an agent.
+      await this.closeRemote()
+      return
     }
+    let agent = gotAgent // Reassign so it cannot be undefined in callbacks.
     this.storage.writeToCoderOutputChannel(`Found agent ${agent.name} with status ${agent.status}`)
 
     // Do some janky setting manipulation.
@@ -466,17 +452,11 @@ export class Remote {
         async () => {
           await new Promise<void>((resolve) => {
             const updateEvent = workspaceUpdate.event((workspace) => {
-              const agents = workspace.latest_build.resources.reduce((acc, resource) => {
-                return acc.concat(resource.agents || [])
-              }, [] as WorkspaceAgent[])
               if (!agent) {
                 return
               }
+              const agents = extractAgents(workspace)
               const found = agents.find((newAgent) => {
-                if (!agent) {
-                  // This shouldn't be possible... just makes the types happy!
-                  return false
-                }
                 return newAgent.id === agent.id
               })
               if (!found) {
@@ -538,10 +518,9 @@ export class Remote {
     })
 
     // Register the label formatter again because SSH overrides it!
-    const agentName = agents.length > 1 ? agent.name : undefined
     disposables.push(
       vscode.extensions.onDidChange(() => {
-        disposables.push(this.registerLabelFormatter(remoteAuthority, workspace.owner_name, workspace.name, agentName))
+        disposables.push(this.registerLabelFormatter(remoteAuthority, workspace.owner_name, workspace.name, agent.name))
       }),
     )