Add method to get agents even for off workspaces

code-asher · code-asher · commit 8b7c049928b6 · 2023-10-18T10:39:38.000-08:00
diff --git a/src/main/kotlin/com/coder/gateway/sdk/CoderRestClientService.kt b/src/main/kotlin/com/coder/gateway/sdk/CoderRestClientService.kt
@@ -1,5 +1,6 @@
 package com.coder.gateway.sdk
 
+import com.coder.gateway.models.WorkspaceAgentModel
 import com.coder.gateway.sdk.convertors.InstantConverter
 import com.coder.gateway.sdk.ex.AuthenticationResponseException
 import com.coder.gateway.sdk.ex.TemplateResponseException
@@ -12,6 +13,7 @@ import com.coder.gateway.sdk.v2.models.User
 import com.coder.gateway.sdk.v2.models.Workspace
 import com.coder.gateway.sdk.v2.models.WorkspaceBuild
 import com.coder.gateway.sdk.v2.models.WorkspaceTransition
+import com.coder.gateway.sdk.v2.models.toAgentModels
 import com.google.gson.Gson
 import com.google.gson.GsonBuilder
 import com.intellij.ide.plugins.PluginManagerCore
@@ -105,6 +107,23 @@ class CoderRestClient(var url: URL, var token: String, var headerCommand: String
         return workspacesResponse.body()!!.workspaces
     }
 
+    /**
+     * Retrieves agents for the specified workspaces.  Since the workspaces
+     * response does not include agents when the workspace is off, this fires
+     * off separate queries to get the agents for each workspace, just like
+     * `coder config-ssh` does (otherwise we risk removing hosts from the SSH
+     * config when they are off).
+     */
+    fun agents(workspaces: List<Workspace>): List<WorkspaceAgentModel> {
+        return workspaces.flatMap {
+            val resourcesResponse = retroRestClient.templateVersionResources(it.latestBuild.templateVersionID).execute()
+            if (!resourcesResponse.isSuccessful) {
+                throw WorkspaceResponseException("Unable to retrieve template resources for ${it.name} from $url: code ${resourcesResponse.code()}, reason: ${resourcesResponse.message().ifBlank { "no reason provided" }}")
+            }
+            it.toAgentModels(resourcesResponse.body()!!)
+        }
+    }
+
     fun buildInfo(): BuildInfo {
         val buildInfoResponse = retroRestClient.buildInfo().execute()
         if (!buildInfoResponse.isSuccessful) {
diff --git a/src/main/kotlin/com/coder/gateway/sdk/v2/CoderV2RestFacade.kt b/src/main/kotlin/com/coder/gateway/sdk/v2/CoderV2RestFacade.kt
@@ -5,6 +5,7 @@ import com.coder.gateway.sdk.v2.models.CreateWorkspaceBuildRequest
 import com.coder.gateway.sdk.v2.models.Template
 import com.coder.gateway.sdk.v2.models.User
 import com.coder.gateway.sdk.v2.models.WorkspaceBuild
+import com.coder.gateway.sdk.v2.models.WorkspaceResource
 import com.coder.gateway.sdk.v2.models.WorkspacesResponse
 import retrofit2.Call
 import retrofit2.http.Body
@@ -39,4 +40,7 @@ interface CoderV2RestFacade {
 
     @GET("api/v2/templates/{templateID}")
     fun template(@Path("templateID") templateID: UUID): Call<Template>
-}
+
+    @GET("api/v2/templateversions/{templateID}/resources")
+    fun templateVersionResources(@Path("templateID") templateID: UUID): Call<List<WorkspaceResource>>
+}
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
@@ -31,8 +31,8 @@ data class Workspace(
     @SerializedName("last_used_at") val lastUsedAt: Instant,
 )
 
-fun Workspace.toAgentModels(): Set<WorkspaceAgentModel> {
-    val wam = this.latestBuild.resources.filter { it.agents != null }.flatMap { it.agents!! }.map { agent ->
+fun Workspace.toAgentModels(resources: List<WorkspaceResource> = this.latestBuild.resources): Set<WorkspaceAgentModel> {
+    val wam = resources.filter { it.agents != null }.flatMap { it.agents!! }.map { agent ->
         val workspaceWithAgentName = "${this.name}.${agent.name}"
         val wm = WorkspaceAgentModel(
             this.id,
