feat: multiple agent support (#59)

ytcheng · kylecarbs · rodrimaia · web-flow · commit 4662ebdebef0 · 2023-04-28T18:40:01.000-03:00
Co-authored-by: Kyle Carberry &lt;kyle@carberry.com&gt;
Co-authored-by: Rodrigo Maia &lt;rodrigo.maia.pereira@gmail.com&gt;
diff --git a/package.json b/package.json
@@ -94,6 +94,11 @@
         "title": "Coder: Open Workspace",
         "icon": "$(play)"
       },
+      {
+        "command": "coder.openFromSidebar",
+        "title": "Coder: Open Workspace",
+        "icon": "$(play)"
+      },
       {
         "command": "coder.createWorkspace",
         "title": "Create Workspace",
@@ -147,18 +152,18 @@
       ],
       "view/item/context": [
         {
-          "command": "coder.open",
-          "when": "coder.authenticated && view == myWorkspaces || coder.authenticated && view == allWorkspaces",
+          "command": "coder.openFromSidebar",
+          "when": "coder.authenticated && viewItem == coderWorkspaceSingleAgent || coder.authenticated && viewItem == coderAgent",
           "group": "inline"
         },
         {
           "command": "coder.navigateToWorkspace",
-          "when": "coder.authenticated && view == myWorkspaces || coder.authenticated && view == allWorkspaces",
+          "when": "coder.authenticated && viewItem == coderWorkspaceSingleAgent || coder.authenticated && viewItem == coderWorkspaceMultipleAgents",
           "group": "inline"
         },
         {
           "command": "coder.navigateToWorkspaceSettings",
-          "when": "coder.authenticated && view == myWorkspaces || coder.authenticated && view == allWorkspaces",
+          "when": "coder.authenticated && viewItem == coderWorkspaceSingleAgent || coder.authenticated && viewItem == coderWorkspaceMultipleAgents",
           "group": "inline"
         }
       ]
@@ -223,4 +228,4 @@
     "ws": "^8.11.0",
     "yaml": "^1.10.0"
   }
-}
+}
diff --git a/src/api-helper.ts b/src/api-helper.ts
@@ -1,16 +1,9 @@
 import { Workspace, WorkspaceAgent } from "coder/site/src/api/typesGenerated"
 
-export function extractAgentsAndFolderPath(
-  workspace: Workspace,
-): [agents: WorkspaceAgent[], folderPath: string | undefined] {
-  // TODO: multiple agent support
+export function extractAgents(workspace: Workspace): WorkspaceAgent[] {
   const agents = workspace.latest_build.resources.reduce((acc, resource) => {
     return acc.concat(resource.agents || [])
   }, [] as WorkspaceAgent[])
 
-  let folderPath = undefined
-  if (agents.length === 1) {
-    folderPath = agents[0].expanded_directory
-  }
-  return [agents, folderPath]
+  return agents
 }
diff --git a/src/commands.ts b/src/commands.ts
@@ -1,8 +1,8 @@
 import axios from "axios"
 import { getAuthenticatedUser, getWorkspaces, updateWorkspaceVersion } from "coder/site/src/api/api"
-import { Workspace } from "coder/site/src/api/typesGenerated"
+import { Workspace, WorkspaceAgent } from "coder/site/src/api/typesGenerated"
 import * as vscode from "vscode"
-import { extractAgentsAndFolderPath } from "./api-helper"
+import { extractAgents } from "./api-helper"
 import { Remote } from "./remote"
 import { Storage } from "./storage"
 import { WorkspaceTreeItem } from "./workspacesProvider"
@@ -143,9 +143,21 @@ export class Commands {
     }
   }
 
+  public async openFromSidebar(treeItem: WorkspaceTreeItem) {
+    if (treeItem) {
+      await openWorkspace(
+        treeItem.workspaceOwner,
+        treeItem.workspaceName,
+        treeItem.workspaceAgent,
+        treeItem.workspaceFolderPath,
+      )
+    }
+  }
+
   public async open(...args: unknown[]): Promise<void> {
     let workspaceOwner: string
     let workspaceName: string
+    let workspaceAgent: string | undefined
     let folderPath: string | undefined
 
     if (args.length === 0) {
@@ -200,83 +212,61 @@ export class Commands {
       workspaceOwner = workspace.owner_name
       workspaceName = workspace.name
 
-      const [, folderPathExtracted] = extractAgentsAndFolderPath(workspace)
-      folderPath = folderPathExtracted
-    } else if (args.length === 2) {
-      // opening a workspace from the sidebar
-      const workspaceTreeItem = args[0] as WorkspaceTreeItem
-      workspaceOwner = workspaceTreeItem.workspaceOwner
-      workspaceName = workspaceTreeItem.workspaceName
-      folderPath = workspaceTreeItem.workspaceFolderPath
-    } else {
-      workspaceOwner = args[0] as string
-      workspaceName = args[1] as string
-      // workspaceAgent is reserved for args[2], but multiple agents aren't supported yet.
-      folderPath = args[3] as string | undefined
-    }
-
-    // A workspace can have multiple agents, but that's handled
-    // when opening a workspace unless explicitly specified.
-    const remoteAuthority = `ssh-remote+${Remote.Prefix}${workspaceOwner}--${workspaceName}`
+      const agents = extractAgents(workspace)
 
-    let newWindow = true
-    // Open in the existing window if no workspaces are open.
-    if (!vscode.workspace.workspaceFolders?.length) {
-      newWindow = false
-    }
+      if (agents.length === 1) {
+        folderPath = agents[0].expanded_directory
+        workspaceAgent = agents[0].name
+      } else {
+        const agentQuickPick = vscode.window.createQuickPick()
+        agentQuickPick.title = `Select an agent`
 
-    // If a folder isn't specified, we can try to open a recently opened folder.
-    if (!folderPath) {
-      const output: {
-        workspaces: { folderUri: vscode.Uri; remoteAuthority: string }[]
-      } = await vscode.commands.executeCommand("_workbench.getRecentlyOpened")
-      const opened = output.workspaces.filter(
-        // Filter out `/` since that's added below.
-        (opened) => opened.folderUri?.authority === remoteAuthority,
-      )
-      if (opened.length > 0) {
-        let selected: (typeof opened)[0]
+        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()
 
-        if (opened.length > 1) {
-          const items: vscode.QuickPickItem[] = opened.map((folder): vscode.QuickPickItem => {
-            return {
-              label: folder.folderUri.path,
-            }
+        const agent = await new Promise<WorkspaceAgent | undefined>((resolve) => {
+          agentQuickPick.onDidHide(() => {
+            resolve(undefined)
           })
-          const item = await vscode.window.showQuickPick(items, {
-            title: "Select a recently opened folder",
+          agentQuickPick.onDidChangeSelection((selected) => {
+            if (selected.length < 1) {
+              return resolve(undefined)
+            }
+            const agent = lastAgents[agentQuickPick.items.indexOf(selected[0])]
+            resolve(agent)
           })
-          if (!item) {
-            return
-          }
-          selected = opened[items.indexOf(item)]
+        })
+
+        if (agent) {
+          folderPath = agent.expanded_directory
+          workspaceAgent = agent.name
         } else {
-          selected = opened[0]
+          folderPath = ""
+          workspaceAgent = ""
         }
-
-        folderPath = selected.folderUri.path
       }
+    } else {
+      workspaceOwner = args[0] as string
+      workspaceName = args[1] as string
+      // workspaceAgent is reserved for args[2], but multiple agents aren't supported yet.
+      folderPath = args[3] as string | undefined
     }
 
-    if (folderPath) {
-      await vscode.commands.executeCommand(
-        "vscode.openFolder",
-        vscode.Uri.from({
-          scheme: "vscode-remote",
-          authority: remoteAuthority,
-          path: folderPath,
-        }),
-        // Open this in a new window!
-        newWindow,
-      )
-      return
-    }
-
-    // This opens the workspace without an active folder opened.
-    await vscode.commands.executeCommand("vscode.newWindow", {
-      remoteAuthority: remoteAuthority,
-      reuseWindow: !newWindow,
-    })
+    await openWorkspace(workspaceOwner, workspaceName, workspaceAgent, folderPath)
   }
 
   public async updateWorkspace(): Promise<void> {
@@ -297,3 +287,76 @@ export class Commands {
     }
   }
 }
+
+async function openWorkspace(
+  workspaceOwner: string,
+  workspaceName: string,
+  workspaceAgent: string | undefined,
+  folderPath: string | undefined,
+) {
+  // A workspace can have multiple agents, but that's handled
+  // when opening a workspace unless explicitly specified.
+  let remoteAuthority = `ssh-remote+${Remote.Prefix}${workspaceOwner}--${workspaceName}`
+  if (workspaceAgent) {
+    remoteAuthority += `--${workspaceAgent}`
+  }
+
+  let newWindow = true
+  // Open in the existing window if no workspaces are open.
+  if (!vscode.workspace.workspaceFolders?.length) {
+    newWindow = false
+  }
+
+  // If a folder isn't specified, we can try to open a recently opened folder.
+  if (!folderPath) {
+    const output: {
+      workspaces: { folderUri: vscode.Uri; remoteAuthority: string }[]
+    } = await vscode.commands.executeCommand("_workbench.getRecentlyOpened")
+    const opened = output.workspaces.filter(
+      // Filter out `/` since that's added below.
+      (opened) => opened.folderUri?.authority === remoteAuthority,
+    )
+    if (opened.length > 0) {
+      let selected: (typeof opened)[0]
+
+      if (opened.length > 1) {
+        const items: vscode.QuickPickItem[] = opened.map((folder): vscode.QuickPickItem => {
+          return {
+            label: folder.folderUri.path,
+          }
+        })
+        const item = await vscode.window.showQuickPick(items, {
+          title: "Select a recently opened folder",
+        })
+        if (!item) {
+          return
+        }
+        selected = opened[items.indexOf(item)]
+      } else {
+        selected = opened[0]
+      }
+
+      folderPath = selected.folderUri.path
+    }
+  }
+
+  if (folderPath) {
+    await vscode.commands.executeCommand(
+      "vscode.openFolder",
+      vscode.Uri.from({
+        scheme: "vscode-remote",
+        authority: remoteAuthority,
+        path: folderPath,
+      }),
+      // Open this in a new window!
+      newWindow,
+    )
+    return
+  }
+
+  // This opens the workspace without an active folder opened.
+  await vscode.commands.executeCommand("vscode.newWindow", {
+    remoteAuthority: remoteAuthority,
+    reuseWindow: !newWindow,
+  })
+}
diff --git a/src/extension.ts b/src/extension.ts
@@ -87,6 +87,7 @@ export async function activate(ctx: vscode.ExtensionContext): Promise<void> {
   vscode.commands.registerCommand("coder.login", commands.login.bind(commands))
   vscode.commands.registerCommand("coder.logout", commands.logout.bind(commands))
   vscode.commands.registerCommand("coder.open", commands.open.bind(commands))
+  vscode.commands.registerCommand("coder.openFromSidebar", commands.openFromSidebar.bind(commands))
   vscode.commands.registerCommand("coder.workspace.update", commands.updateWorkspace.bind(commands))
   vscode.commands.registerCommand("coder.createWorkspace", commands.createWorkspace.bind(commands))
   vscode.commands.registerCommand("coder.navigateToWorkspace", commands.navigateToWorkspace.bind(commands))
diff --git a/src/workspacesProvider.ts b/src/workspacesProvider.ts
@@ -1,7 +1,8 @@
 import { getWorkspaces } from "coder/site/src/api/api"
+import { WorkspaceAgent } from "coder/site/src/api/typesGenerated"
 import * as path from "path"
 import * as vscode from "vscode"
-import { extractAgentsAndFolderPath } from "./api-helper"
+import { extractAgents } from "./api-helper"
 
 export enum WorkspaceQuery {
   Mine = "owner:me",
@@ -24,7 +25,19 @@ export class WorkspaceProvider implements vscode.TreeDataProvider<WorkspaceTreeI
     return element
   }
 
-  getChildren(): Thenable<WorkspaceTreeItem[]> {
+  getChildren(element?: WorkspaceTreeItem): Thenable<WorkspaceTreeItem[]> {
+    if (element) {
+      if (element.agents.length > 0) {
+        return Promise.resolve(
+          element.agents.map((agent) => {
+            const label = agent.name
+            const detail = `Status: ${agent.status}`
+            return new WorkspaceTreeItem(label, detail, "", "", agent.name, agent.expanded_directory, [], "coderAgent")
+          }),
+        )
+      }
+      return Promise.resolve([])
+    }
     return getWorkspaces({ q: this.getWorkspacesQuery }).then((workspaces) => {
       return workspaces.workspaces.map((workspace) => {
         const status =
@@ -35,22 +48,42 @@ export class WorkspaceProvider implements vscode.TreeDataProvider<WorkspaceTreeI
             ? `${workspace.owner_name} / ${workspace.name}`
             : workspace.name
         const detail = `Template: ${workspace.template_display_name || workspace.template_name} • Status: ${status}`
-        const [, folderPath] = extractAgentsAndFolderPath(workspace)
-        return new WorkspaceTreeItem(label, detail, workspace.owner_name, workspace.name, folderPath)
+        const agents = extractAgents(workspace)
+        return new WorkspaceTreeItem(
+          label,
+          detail,
+          workspace.owner_name,
+          workspace.name,
+          undefined,
+          agents[0]?.expanded_directory,
+          agents,
+          agents.length > 1 ? "coderWorkspaceMultipleAgents" : "coderWorkspaceSingleAgent",
+        )
       })
     })
   }
 }
 
+type CoderTreeItemType = "coderWorkspaceSingleAgent" | "coderWorkspaceMultipleAgents" | "coderAgent"
+
 export class WorkspaceTreeItem extends vscode.TreeItem {
   constructor(
     public readonly label: string,
     public readonly tooltip: string,
     public readonly workspaceOwner: string,
     public readonly workspaceName: string,
+    public readonly workspaceAgent: string | undefined,
     public readonly workspaceFolderPath: string | undefined,
+    public readonly agents: WorkspaceAgent[],
+    contextValue: CoderTreeItemType,
   ) {
-    super(label, vscode.TreeItemCollapsibleState.None)
+    super(
+      label,
+      contextValue === "coderWorkspaceMultipleAgents"
+        ? vscode.TreeItemCollapsibleState.Collapsed
+        : vscode.TreeItemCollapsibleState.None,
+    )
+    this.contextValue = contextValue
   }
 
   iconPath = {

Original file line number	Diff line number	Diff line change
`@@ -94,6 +94,11 @@`
`94`	`94`	`"title": "Coder: Open Workspace",`
`95`	`95`	`"icon": "$(play)"`
`96`	`96`	`},`
	`97`	`+ {`
	`98`	`+ "command": "coder.openFromSidebar",`
	`99`	`+ "title": "Coder: Open Workspace",`
	`100`	`+ "icon": "$(play)"`
	`101`	`+ },`
`97`	`102`	`{`
`98`	`103`	`"command": "coder.createWorkspace",`
`99`	`104`	`"title": "Create Workspace",`
`@@ -147,18 +152,18 @@`
`147`	`152`	`],`
`148`	`153`	`"view/item/context": [`
`149`	`154`	`{`
`150`		`- "command": "coder.open",`
`151`		`- "when": "coder.authenticated && view == myWorkspaces \|\| coder.authenticated && view == allWorkspaces",`
	`155`	`+ "command": "coder.openFromSidebar",`
	`156`	`+ "when": "coder.authenticated && viewItem == coderWorkspaceSingleAgent \|\| coder.authenticated && viewItem == coderAgent",`
`152`	`157`	`"group": "inline"`
`153`	`158`	`},`
`154`	`159`	`{`
`155`	`160`	`"command": "coder.navigateToWorkspace",`
`156`		`- "when": "coder.authenticated && view == myWorkspaces \|\| coder.authenticated && view == allWorkspaces",`
	`161`	`+ "when": "coder.authenticated && viewItem == coderWorkspaceSingleAgent \|\| coder.authenticated && viewItem == coderWorkspaceMultipleAgents",`
`157`	`162`	`"group": "inline"`
`158`	`163`	`},`
`159`	`164`	`{`
`160`	`165`	`"command": "coder.navigateToWorkspaceSettings",`
`161`		`- "when": "coder.authenticated && view == myWorkspaces \|\| coder.authenticated && view == allWorkspaces",`
	`166`	`+ "when": "coder.authenticated && viewItem == coderWorkspaceSingleAgent \|\| coder.authenticated && viewItem == coderWorkspaceMultipleAgents",`
`162`	`167`	`"group": "inline"`
`163`	`168`	`}`
`164`	`169`	`]`
`@@ -223,4 +228,4 @@`
`223`	`228`	`"ws": "^8.11.0",`
`224`	`229`	`"yaml": "^1.10.0"`
`225`	`230`	`}`
`226`		`-}`
	`231`	`+}`