Support agent ID parameter

code-asher · code-asher · commit c8e503ee04c8 · 2023-10-18T10:39:38.000-08:00
diff --git a/src/main/kotlin/com/coder/gateway/CoderGatewayConnectionProvider.kt b/src/main/kotlin/com/coder/gateway/CoderGatewayConnectionProvider.kt
@@ -23,7 +23,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"
@@ -72,14 +73,23 @@ class CoderGatewayConnectionProvider : GatewayConnectionProvider {
             }
 
             // 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]}"}
+            // 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 == "$workspaceName.${parameters[AGENT_NAME]}"}
             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")
+                if (parameters[AGENT_ID].isNullOrBlank() && parameters[AGENT_NAME].isNullOrBlank()) {
+                    // TODO: Show a dropdown and ask for an agent.
+                    throw IllegalArgumentException("Unable to determine which agent to connect to; one of \"$AGENT_NAME\" or \"$AGENT_ID\" must be set because \"$workspaceName\" has more than one agent")
+                } else if (parameters[AGENT_ID].isNullOrBlank()) {
+                    throw IllegalArgumentException("The workspace \"$workspaceName\" does not have an agent with ID \"${parameters[AGENT_ID]}\"")
+                } else {
+                    throw IllegalArgumentException("The workspace \"$workspaceName\" does not have an agent named \"${parameters[AGENT_NAME]}\"")
+                }
             }
 
             if (agent.agentStatus.pending()) {
diff --git a/src/main/kotlin/com/coder/gateway/models/WorkspaceAgentModel.kt b/src/main/kotlin/com/coder/gateway/models/WorkspaceAgentModel.kt
@@ -14,6 +14,7 @@ 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.
diff --git a/src/main/kotlin/com/coder/gateway/sdk/v2/models/Workspace.kt b/src/main/kotlin/com/coder/gateway/sdk/v2/models/Workspace.kt
@@ -35,6 +35,7 @@ fun Workspace.toAgentModels(resources: List<WorkspaceResource> = this.latestBuil
     val wam = 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(resources: List<WorkspaceResource> = this.latestBuil
     }.toSet()
     if (wam.isNullOrEmpty()) {
         val wm = WorkspaceAgentModel(
+            null,
             this.id,
             this.name,
             this.name,
diff --git a/src/test/groovy/DataGen.groovy b/src/test/groovy/DataGen.groovy
@@ -7,6 +7,7 @@ import com.coder.gateway.sdk.v2.models.WorkspaceTransition
 class DataGen {
     static WorkspaceAgentModel workspace(String name, String workspaceName = name) {
         return new WorkspaceAgentModel(
+                UUID.randomUUID(),
                 UUID.randomUUID(),
                 workspaceName,
                 name,
