Support agent ID parameter

src/main/kotlin/com/coder/gateway/CoderGatewayConnectionProvider.kt

@@ -3,10 +3,12 @@
 package com.coder.gateway
 
 import com.coder.gateway.models.TokenSource
+import com.coder.gateway.models.WorkspaceAgentModel
 import com.coder.gateway.sdk.CoderCLIManager
 import com.coder.gateway.sdk.CoderRestClient
 import com.coder.gateway.sdk.ex.AuthenticationResponseException
 import com.coder.gateway.sdk.toURL
+import com.coder.gateway.sdk.v2.models.Workspace
 import com.coder.gateway.sdk.v2.models.WorkspaceStatus
 import com.coder.gateway.sdk.v2.models.toAgentModels
 import com.coder.gateway.sdk.withPath
@@ -23,7 +25,8 @@ import java.net.URL
 private const val URL = "url"
 private const val TOKEN = "token"
 private const val WORKSPACE = "workspace"
-private const val AGENT = "agent"
+private const val AGENT_NAME = "agent"
+private const val AGENT_ID = "agent_id"
 private const val FOLDER = "folder"
 private const val IDE_DOWNLOAD_LINK = "ide_download_link"
 private const val IDE_PRODUCT_CODE = "ide_product_code"
@@ -66,21 +69,8 @@ class CoderGatewayConnectionProvider : GatewayConnectionProvider {
                 WorkspaceStatus.RUNNING -> Unit // All is well
             }
 
-            val agents = workspace.toAgentModels()
-            if (agents.isEmpty()) {
-                throw IllegalArgumentException("The workspace \"$workspaceName\" has no agents")
-            }
-
-            // If the agent is missing and the workspace has only one, use that.
-            val agent = if (!parameters[AGENT].isNullOrBlank())
-                agents.firstOrNull { it.name == "$workspaceName.${parameters[AGENT]}"}
-            else if (agents.size == 1) agents.first()
-            else null
-
-            if (agent == null) {
-                // TODO: Show a dropdown and ask for an agent.
-                throw IllegalArgumentException("Query parameter \"$AGENT\" is missing")
-            }
+            // TODO: Show a dropdown and ask for an agent if missing.
+            val agent = getMatchingAgent(parameters, workspace)
 
             if (agent.agentStatus.pending()) {
                 // TODO: Wait for the agent to be ready.
@@ -201,5 +191,49 @@ class CoderGatewayConnectionProvider : GatewayConnectionProvider {
 
     companion object {
         val logger = Logger.getInstance(CoderGatewayConnectionProvider::class.java.simpleName)
+
+        /**
+         * Return the agent matching the provided agent ID or name in the
+         * parameters.  The name is ignored if the ID is set.  If neither was
+         * supplied and the workspace has only one agent, return that.
+         * Otherwise throw an error.
+         *
+         * @throws [MissingArgumentException, IllegalArgumentException]
+         */
+        @JvmStatic
+        fun getMatchingAgent(parameters: Map<String, String>, workspace: Workspace): WorkspaceAgentModel {
+            // A WorkspaceAgentModel will still be returned if there are no
+            // agents; in this case it represents the workspace instead.
+            // TODO: Seems confusing for something with "agent" in the name to
+            //       potentially not actually be an agent; can we replace
+            //       WorkspaceAgentModel with the original structs from the API?
+            val agents = workspace.toAgentModels()
+            if (agents.isEmpty() || (agents.size == 1 && agents.first().agentID == null)) {
+                throw IllegalArgumentException("The workspace \"${workspace.name}\" has no agents")
+            }
+
+            // If the agent is missing and the workspace has only one, use that.
+            // Prefer the ID over the name if both are set.
+            val agent = if (!parameters[AGENT_ID].isNullOrBlank())
+                agents.firstOrNull { it.agentID.toString() == parameters[AGENT_ID] }
+            else if (!parameters[AGENT_NAME].isNullOrBlank())
+                agents.firstOrNull { it.name == "${workspace.name}.${parameters[AGENT_NAME]}"}
+            else if (agents.size == 1) agents.first()
+            else null
+
+            if (agent == null) {
+                if (!parameters[AGENT_ID].isNullOrBlank()) {
+                    throw IllegalArgumentException("The workspace \"${workspace.name}\" does not have an agent with ID \"${parameters[AGENT_ID]}\"")
+                } else if (!parameters[AGENT_NAME].isNullOrBlank()){
+                    throw IllegalArgumentException("The workspace \"${workspace.name}\"does not have an agent named \"${parameters[AGENT_NAME]}\"")
+                } else {
+                    throw MissingArgumentException("Unable to determine which agent to connect to; one of \"$AGENT_NAME\" or \"$AGENT_ID\" must be set because the workspace \"${workspace.name}\" has more than one agent")
+                }
+            }
+
+            return agent
+        }
     }
 }
+
+class MissingArgumentException(message: String) : IllegalArgumentException(message)

src/main/kotlin/com/coder/gateway/models/WorkspaceAgentModel.kt

@@ -14,9 +14,10 @@ import javax.swing.Icon
 // iterate over the list we can add the workspace row if it has no agents
 // otherwise iterate over the agents and then flatten the result.
 data class WorkspaceAgentModel(
+    val agentID: UUID?,
     val workspaceID: UUID,
     val workspaceName: String,
-    val name: String, // Name of the workspace OR the agent if this is for an agent.
+    val name: String, // Name of the workspace OR workspace.agent if this is for an agent.
     val templateID: UUID,
     val templateName: String,
     val templateIconPath: String,

src/main/kotlin/com/coder/gateway/sdk/v2/models/Workspace.kt

@@ -35,6 +35,7 @@ fun Workspace.toAgentModels(): Set<WorkspaceAgentModel> {
     val wam = this.latestBuild.resources.filter { it.agents != null }.flatMap { it.agents!! }.map { agent ->
         val workspaceWithAgentName = "${this.name}.${agent.name}"
         val wm = WorkspaceAgentModel(
+            agent.id,
             this.id,
             this.name,
             workspaceWithAgentName,
@@ -55,6 +56,7 @@ fun Workspace.toAgentModels(): Set<WorkspaceAgentModel> {
     }.toSet()
     if (wam.isNullOrEmpty()) {
         val wm = WorkspaceAgentModel(
+            null,
             this.id,
             this.name,
             this.name,

src/test/groovy/CoderCLIManagerTest.groovy

@@ -413,7 +413,7 @@ class CoderCLIManagerTest extends Specification {
                 .replace("/tmp/coder-gateway/test.coder.invalid/coder-linux-amd64", CoderCLIManager.escape(ccm.localBinaryPath.toString()))
 
         when:
-        ccm.configSsh(workspaces.collect { DataGen.workspace(it) }, headerCommand)
+        ccm.configSsh(workspaces.collect { DataGen.workspaceAgentModel(it) }, headerCommand)
 
         then:
         sshConfigPath.toFile().text == expectedConf
@@ -473,7 +473,7 @@ class CoderCLIManagerTest extends Specification {
         def ccm = new CoderCLIManager(new URL("https://test.coder.invalid"), tmpdir)
 
         when:
-        ccm.configSsh(["foo", "bar"].collect { DataGen.workspace(it) }, headerCommand)
+        ccm.configSsh(["foo", "bar"].collect { DataGen.workspaceAgentModel(it) }, headerCommand)
 
         then:
         thrown(Exception)

src/test/groovy/CoderGatewayConnectionProviderTest.groovy

@@ -0,0 +1,114 @@
+package com.coder.gateway
+
+import spock.lang.Shared
+import spock.lang.Specification
+import spock.lang.Unroll
+
+@Unroll
+class CoderGatewayConnectionProviderTest extends Specification {
+    @Shared
+    def agents = [
+        agent_name_3: "b0e4c54d-9ba9-4413-8512-11ca1e826a24",
+        agent_name_2: "fb3daea4-da6b-424d-84c7-36b90574cfef",
+        agent_name: "9a920eee-47fb-4571-9501-e4b3120c12f2",
+    ]
+    def oneAgent = [
+        agent_name_3: "b0e4c54d-9ba9-4413-8512-11ca1e826a24"
+    ]
+
+    def "gets matching agent"() {
+        expect:
+        def ws = DataGen.workspace("ws", agents)
+        CoderGatewayConnectionProvider.getMatchingAgent(parameters, ws).agentID == UUID.fromString(expected)
+
+        where:
+        parameters                                         | expected
+        [agent:    "agent_name"]                           | "9a920eee-47fb-4571-9501-e4b3120c12f2"
+        [agent_id: "9a920eee-47fb-4571-9501-e4b3120c12f2"] | "9a920eee-47fb-4571-9501-e4b3120c12f2"
+        [agent:    "agent_name_2"]                         | "fb3daea4-da6b-424d-84c7-36b90574cfef"
+        [agent_id: "fb3daea4-da6b-424d-84c7-36b90574cfef"] | "fb3daea4-da6b-424d-84c7-36b90574cfef"
+        [agent:    "agent_name_3"]                         | "b0e4c54d-9ba9-4413-8512-11ca1e826a24"
+        [agent_id: "b0e4c54d-9ba9-4413-8512-11ca1e826a24"] | "b0e4c54d-9ba9-4413-8512-11ca1e826a24"
+
+        // Prefer agent_id.
+        [agent: "agent_name", agent_id: "b0e4c54d-9ba9-4413-8512-11ca1e826a24"] | "b0e4c54d-9ba9-4413-8512-11ca1e826a24"
+    }
+
+    def "fails to get matching agent"() {
+        when:
+        def ws = DataGen.workspace("ws", agents)
+        CoderGatewayConnectionProvider.getMatchingAgent(parameters, ws)
+
+        then:
+        def err = thrown(expected)
+        err.message.contains(message)
+
+        where:
+        parameters                                         | expected                 | message
+        [:]                                                | MissingArgumentException | "Unable to determine"
+        [agent: ""]                                        | MissingArgumentException | "Unable to determine"
+        [agent_id: ""]                                     | MissingArgumentException | "Unable to determine"
+        [agent: null]                                      | MissingArgumentException | "Unable to determine"
+        [agent_id: null]                                   | MissingArgumentException | "Unable to determine"
+        [agent: "ws"]                                      | IllegalArgumentException | "agent named"
+        [agent: "ws.agent_name"]                           | IllegalArgumentException | "agent named"
+        [agent: "agent_name_4"]                            | IllegalArgumentException | "agent named"
+        [agent_id: "not-a-uuid"]                           | IllegalArgumentException | "agent with ID"
+        [agent_id: "ceaa7bcf-1612-45d7-b484-2e0da9349168"] | IllegalArgumentException | "agent with ID"
+
+        // Will ignore agent if agent_id is set even if agent matches.
+        [agent: "agent_name", agent_id: "ceaa7bcf-1612-45d7-b484-2e0da9349168"] | IllegalArgumentException | "agent with ID"
+    }
+
+    def "gets the first agent when workspace has only one"() {
+        expect:
+        def ws = DataGen.workspace("ws", oneAgent)
+        CoderGatewayConnectionProvider.getMatchingAgent(parameters, ws).agentID == UUID.fromString("b0e4c54d-9ba9-4413-8512-11ca1e826a24")
+
+        where:
+        parameters << [
+            [:],
+            [agent: ""],
+            [agent_id: ""],
+            [agent: null],
+            [agent_id: null],
+        ]
+    }
+
+    def "fails to get agent when workspace has only one"() {
+        when:
+        def ws = DataGen.workspace("ws", oneAgent)
+        CoderGatewayConnectionProvider.getMatchingAgent(parameters, ws)
+
+        then:
+        def err = thrown(expected)
+        err.message.contains(message)
+
+        where:
+        parameters                                         | expected                 | message
+        [agent:    "ws"]                                   | IllegalArgumentException | "agent named"
+        [agent:    "ws.agent_name_3"]                      | IllegalArgumentException | "agent named"
+        [agent:    "agent_name_4"]                         | IllegalArgumentException | "agent named"
+        [agent_id: "ceaa7bcf-1612-45d7-b484-2e0da9349168"] | IllegalArgumentException | "agent with ID"
+    }
+
+    def "fails to get agent from workspace without agents"() {
+        when:
+        def ws = DataGen.workspace("ws")
+        CoderGatewayConnectionProvider.getMatchingAgent(parameters, ws)
+
+        then:
+        def err = thrown(expected)
+        err.message.contains(message)
+
+        where:
+        parameters                                         | expected                 | message
+        [:]                                                | IllegalArgumentException | "has no agents"
+        [agent: ""]                                        | IllegalArgumentException | "has no agents"
+        [agent_id: ""]                                     | IllegalArgumentException | "has no agents"
+        [agent: null]                                      | IllegalArgumentException | "has no agents"
+        [agent_id: null]                                   | IllegalArgumentException | "has no agents"
+        [agent:    "agent_name"]                           | IllegalArgumentException | "has no agents"
+        [agent_id: "9a920eee-47fb-4571-9501-e4b3120c12f2"] | IllegalArgumentException | "has no agents"
+    }
+}

src/test/groovy/CoderWorkspacesStepViewTest.groovy

@@ -9,46 +9,46 @@ class CoderWorkspacesStepViewTest extends Specification {
         def table = new WorkspacesTable()
         table.listTableModel.items = List.of(
                 // An off workspace.
-                DataGen.workspace("ws1", "ws1"),
+                DataGen.workspaceAgentModel("ws1", "ws1"),
 
                 // On workspaces.
-                DataGen.workspace("agent1", "ws2"),
-                DataGen.workspace("agent2", "ws2"),
-                DataGen.workspace("agent3", "ws3"),
+                DataGen.workspaceAgentModel("agent1", "ws2"),
+                DataGen.workspaceAgentModel("agent2", "ws2"),
+                DataGen.workspaceAgentModel("agent3", "ws3"),
 
                 // Another off workspace.
-                DataGen.workspace("ws4", "ws4"),
+                DataGen.workspaceAgentModel("ws4", "ws4"),
 
                 // In practice we do not list both agents and workspaces
                 // together but here test that anyway with an agent first and
                 // then with a workspace first.
-                DataGen.workspace("agent2", "ws5"),
-                DataGen.workspace("ws5", "ws5"),
-                DataGen.workspace("ws6", "ws6"),
-                DataGen.workspace("agent3", "ws6"),
+                DataGen.workspaceAgentModel("agent2", "ws5"),
+                DataGen.workspaceAgentModel("ws5", "ws5"),
+                DataGen.workspaceAgentModel("ws6", "ws6"),
+                DataGen.workspaceAgentModel("agent3", "ws6"),
         )
 
         expect:
         table.getNewSelection(selected) == expected
 
         where:
-        selected                           | expected
-        null                               | -1 // No selection.
-        DataGen.workspace("gone", "gone")  | -1 // No workspace that matches.
-        DataGen.workspace("ws1", "ws1")    | 0  // Workspace exact match.
-        DataGen.workspace("gone", "ws1")   | 0  // Agent gone, select workspace.
-        DataGen.workspace("ws2", "ws2")    | 1  // Workspace gone, select first agent.
-        DataGen.workspace("agent1", "ws2") | 1  // Agent exact match.
-        DataGen.workspace("agent2", "ws2") | 2  // Agent exact match.
-        DataGen.workspace("ws3", "ws3")    | 3  // Workspace gone, select first agent.
-        DataGen.workspace("agent3", "ws3") | 3  // Agent exact match.
-        DataGen.workspace("gone", "ws4")   | 4  // Agent gone, select workspace.
-        DataGen.workspace("ws4", "ws4")    | 4  // Workspace exact match.
-        DataGen.workspace("agent2", "ws5") | 5  // Agent exact match.
-        DataGen.workspace("gone", "ws5")   | 5  // Agent gone, another agent comes first.
-        DataGen.workspace("ws5", "ws5")    | 6  // Workspace exact match.
-        DataGen.workspace("ws6", "ws6")    | 7  // Workspace exact match.
-        DataGen.workspace("gone", "ws6")   | 7  // Agent gone, workspace comes first.
-        DataGen.workspace("agent3", "ws6") | 8  // Agent exact match.
+        selected                                     | expected
+        null                                         | -1 // No selection.
+        DataGen.workspaceAgentModel("gone", "gone")  | -1 // No workspace that matches.
+        DataGen.workspaceAgentModel("ws1", "ws1")    | 0  // Workspace exact match.
+        DataGen.workspaceAgentModel("gone", "ws1")   | 0  // Agent gone, select workspace.
+        DataGen.workspaceAgentModel("ws2", "ws2")    | 1  // Workspace gone, select first agent.
+        DataGen.workspaceAgentModel("agent1", "ws2") | 1  // Agent exact match.
+        DataGen.workspaceAgentModel("agent2", "ws2") | 2  // Agent exact match.
+        DataGen.workspaceAgentModel("ws3", "ws3")    | 3  // Workspace gone, select first agent.
+        DataGen.workspaceAgentModel("agent3", "ws3") | 3  // Agent exact match.
+        DataGen.workspaceAgentModel("gone", "ws4")   | 4  // Agent gone, select workspace.
+        DataGen.workspaceAgentModel("ws4", "ws4")    | 4  // Workspace exact match.
+        DataGen.workspaceAgentModel("agent2", "ws5") | 5  // Agent exact match.
+        DataGen.workspaceAgentModel("gone", "ws5")   | 5  // Agent gone, another agent comes first.
+        DataGen.workspaceAgentModel("ws5", "ws5")    | 6  // Workspace exact match.
+        DataGen.workspaceAgentModel("ws6", "ws6")    | 7  // Workspace exact match.
+        DataGen.workspaceAgentModel("gone", "ws6")   | 7  // Agent gone, workspace comes first.
+        DataGen.workspaceAgentModel("agent3", "ws6") | 8  // Agent exact match.
     }
 }