fix: slow workspace list refresh when many workspaces are stopped

CHANGELOG.md

@@ -11,6 +11,7 @@
 
 - snackbar dismissal no longer cancels the calling coroutine, so error popups during URI handling don't leave the page
   stuck in a busy state
+- faster workspace list refresh: agents are no longer fetched for stopped workspaces
 
 ### Changed
 

README.md

@@ -485,7 +485,9 @@ Changing the search or a dropdown refreshes the workspace list and regenerates S
 workspace
 set.
 With wildcard SSH enabled, the generated block uses the deployment wildcard host pattern. With wildcard SSH disabled,
-the generated block contains entries for the currently resolved workspace/agent pairs.
+the generated block contains entries for the currently resolved workspace/agent pairs. Agents are resolved only for
+running workspaces, so stopped workspaces have no SSH entries until they are started and picked up by the next
+workspace refresh.
 
 SSH hostnames include the workspace owner so workspaces with the same name owned by different users remain distinct.
 With wildcard SSH enabled, the SSH config contains a deployment-wide `Host coder-jetbrains-toolbox-<host>--*` entry, and

src/main/kotlin/com/coder/toolbox/CoderRemoteEnvironment.kt

@@ -13,6 +13,7 @@ import com.coder.toolbox.util.waitForFalseWithTimeout
 import com.coder.toolbox.util.withPath
 import com.coder.toolbox.views.Action
 import com.coder.toolbox.views.CoderDelimiter
+import com.coder.toolbox.views.EmptyWorkspaceContentView
 import com.coder.toolbox.views.EnvironmentView
 import com.jetbrains.toolbox.api.localization.LocalizableString
 import com.jetbrains.toolbox.api.remoteDev.AfterDisconnectHook
@@ -28,6 +29,7 @@ import com.squareup.moshi.Moshi
 import kotlinx.coroutines.CoroutineName
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.Job
+import kotlinx.coroutines.channels.Channel
 import kotlinx.coroutines.delay
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.StateFlow
@@ -42,21 +44,25 @@ import kotlin.time.Duration.Companion.seconds
 
 private val POLL_INTERVAL = 5.seconds
 
+private fun environmentId(workspace: Workspace, agent: WorkspaceAgent?): String =
+    agent?.let { "${workspace.name}.${it.name}" } ?: workspace.name
+
 /**
- * Represents an agent and workspace combination.
+ * Represents a workspace, or a workspace and agent combination when an agent is available.
  *
  * Used in the environment list view.
  */
 class CoderRemoteEnvironment(
     private val context: CoderToolboxContext,
     internal var client: CoderRestClient,
     internal var cli: CoderCLIManager,
+    private val workspaceRefreshTrigger: Channel<Boolean>,
     private var workspace: Workspace,
-    private var agent: WorkspaceAgent,
-) : RemoteProviderEnvironment("${workspace.name}.${agent.name}"), BeforeConnectionHook, AfterDisconnectHook {
+    private var agent: WorkspaceAgent?,
+) : RemoteProviderEnvironment(environmentId(workspace, agent)), BeforeConnectionHook, AfterDisconnectHook {
     private var environmentStatus = WorkspaceAndAgentStatus.from(workspace, agent)
 
-    override var name: String = "${workspace.name}.${agent.name}"
+    override var name: String = environmentId(workspace, agent)
     private var isConnected: MutableStateFlow<Boolean> = MutableStateFlow(false)
     override val connectionRequest: MutableStateFlow<Boolean> = MutableStateFlow(false)
 
@@ -79,12 +85,12 @@ class CoderRemoteEnvironment(
         refreshAvailableActions()
     }
 
-    fun asPairOfWorkspaceAndAgent(): Pair<Workspace, WorkspaceAgent> = Pair(workspace, agent)
+    fun toWorkspaceAgentPairOrNull(): Pair<Workspace, WorkspaceAgent>? = agent?.let { Pair(workspace, it) }
 
     private fun refreshAvailableActions() {
         val actions = mutableListOf<ActionDescription>()
         context.logger.debug("Refreshing available actions for workspace $id with status: $environmentStatus")
-        if (environmentStatus.canStop()) {
+        if (environmentStatus.canStop() && agent != null) {
             actions.add(Action(context, "Open web terminal") {
                 context.logger.debug("Launching web terminal for $id...")
                 context.desktop.browse(client.url.withPath("/${workspace.ownerName}/$name/terminal").toString()) {
@@ -122,13 +128,15 @@ class CoderRemoteEnvironment(
                     context.logger.debug("Updating and starting $id...")
                     val build = client.updateWorkspace(workspace)
                     update(workspace.copy(latestBuild = build), agent)
+                    workspaceRefreshTrigger.trySend(true)
                 })
             } else {
                 actions.add(Action(context, "Start") {
                     context.logger.debug("Starting $id... ")
                     context.cs
                         .launch(CoroutineName("Start Workspace Action CLI Runner") + Dispatchers.IO) {
                             cli.startWorkspace(workspace.ownerName, workspace.name)
+                            workspaceRefreshTrigger.trySend(true)
                         }
                     // cli takes 15 seconds to move the workspace in queueing/starting state
                     // while the user won't see anything happening in TBX after start is clicked
@@ -145,6 +153,7 @@ class CoderRemoteEnvironment(
                     context.logger.debug("Updating and re-starting $id...")
                     val build = client.updateWorkspace(workspace)
                     update(workspace.copy(latestBuild = build), agent)
+                    workspaceRefreshTrigger.trySend(true)
                 }
                 )
             }
@@ -204,6 +213,9 @@ class CoderRemoteEnvironment(
     override fun getAfterDisconnectHooks(): List<AfterDisconnectHook> = listOf(this)
 
     override fun beforeConnection() {
+        if (agent == null) {
+            return
+        }
         context.logger.info("Connecting to $id...")
         isConnected.update { true }
         context.settingsStore.updateAutoConnect(this.id, true)
@@ -213,8 +225,9 @@ class CoderRemoteEnvironment(
     private fun pollNetworkMetrics(): Job = context.cs.launch(CoroutineName("Network Metrics Poller")) {
         context.logger.info("Starting the network metrics poll job for $id")
         while (isActive) {
+            val currentAgent = agent ?: break
             context.logger.debug("Searching SSH command's PID for workspace $id...")
-            val pid = proxyCommandHandle.findByWorkspaceAndAgent(workspace, agent)
+            val pid = proxyCommandHandle.findByWorkspaceAndAgent(workspace, currentAgent)
             if (pid == null) {
                 context.logger.debug("No SSH command PID was found for workspace $id")
                 delay(POLL_INTERVAL)
@@ -244,6 +257,20 @@ class CoderRemoteEnvironment(
 
     private fun File.doesNotExists(): Boolean = !this.exists()
 
+    /**
+     * Cancels background work when the provider drops this environment from its
+     * list (e.g. a running workspace stopped and is replaced by a workspace-only
+     * environment with a different id).
+     *
+     * Toolbox reacts to the removal by closing its environment wrapper, which only
+     * cancels the wrapper's own coroutine scope and never invokes the
+     * [AfterDisconnectHook], so the network metrics poller must be stopped here.
+     */
+    fun dispose() {
+        pollJob?.cancel()
+        isConnected.update { false }
+    }
+
     override fun afterDisconnect(isManual: Boolean) {
         context.logger.info("Stopping the network metrics poll job for $id")
         pollJob?.cancel()
@@ -259,20 +286,20 @@ class CoderRemoteEnvironment(
     /**
      * Update the workspace/agent status to the listeners, if it has changed.
      */
-    /**
-     * Update the workspace/agent status to the listeners, if it has changed.
-     */
-    fun update(workspace: Workspace, agent: WorkspaceAgent) {
+    fun update(workspace: Workspace, agent: WorkspaceAgent?) {
         if (this.workspace.latestBuild == workspace.latestBuild) {
             return
         }
         this.workspace = workspace
         this.agent = agent
+        name = environmentId(workspace, agent)
 
         // workspace&agent status can be different from "environment status"
         // which is forced to queued state when a workspace is scheduled to start
         updateStatus(WorkspaceAndAgentStatus.from(workspace, agent))
-        context.connectionMonitoringService.checkConnectionStatus(workspace, agent)
+        if (agent != null) {
+            context.connectionMonitoringService.checkConnectionStatus(workspace, agent)
+        }
 
         // we have to regenerate the action list in order to force a redraw
         // because the actions don't have a state flow on the enabled property
@@ -285,20 +312,33 @@ class CoderRemoteEnvironment(
         state.update {
             environmentStatus.toRemoteEnvironmentState(context)
         }
-        context.logger.info("Overall status for workspace $id is $environmentStatus. Workspace status: ${workspace.latestBuild.status}, agent status: ${agent.status}, agent lifecycle state: ${agent.lifecycleState}, login before ready: ${agent.loginBeforeReady}")
+        context.logger.info(
+            "Overall status for workspace $id is $environmentStatus. " +
+                    "Workspace status: ${workspace.latestBuild.status}, " +
+                    "agent status: ${agent?.status}, " +
+                    "agent lifecycle state: ${agent?.lifecycleState}, " +
+                    "login before ready: ${agent?.loginBeforeReady}"
+        )
     }
 
     /**
      * The contents are provided by the SSH view provided by Toolbox, all we
-     * have to do is provide it a host name.
+     * have to do is provide it a host name. Workspaces without a resolved agent
+     * have no host to connect to yet, so they get an empty contents view instead.
      */
-    override suspend fun getContentsView(): EnvironmentContentsView = EnvironmentView(
-        context,
-        client.url,
-        cli,
-        workspace,
-        agent
-    )
+    override suspend fun getContentsView(): EnvironmentContentsView {
+        val envAgent = agent ?: run {
+            context.logger.info("No agent is available for $id yet, providing an empty contents view")
+            return EmptyWorkspaceContentView
+        }
+        return EnvironmentView(
+            context,
+            client.url,
+            cli,
+            workspace,
+            envAgent
+        )
+    }
 
     /**
      * Automatically launches the SSH connection if the workspace is visible, is ready and there is no

src/main/kotlin/com/coder/toolbox/CoderRemoteProvider.kt

@@ -100,12 +100,12 @@ class CoderRemoteProvider(
             }
         }
     }
-    private val linkHandler = CoderProtocolHandler(context, IdeFeedManager(context))
-
     override val loadingEnvironmentsDescription: LocalizableString = context.i18n.ptrl("Loading workspaces...")
     override val environments: MutableStateFlow<LoadableState<List<CoderRemoteEnvironment>>> = MutableStateFlow(
         LoadableState.Loading
     )
+    private val linkHandler =
+        CoderProtocolHandler(context, IdeFeedManager(context), workspaceRefreshTrigger, environments)
     private val accountDropdownField = dropDownFactory(context.i18n.pnotr("")) {
         logout()
         context.envPageManager.showPluginEnvironmentsPage(false)
@@ -133,10 +133,14 @@ class CoderRemoteProvider(
                         return@launch
                     }
 
+                    // Toolbox closes removed environments without firing their
+                    // disconnect hooks, so stop their background work before dropping them.
+                    lastEnvironments.filter { it !in resolvedEnvironments }.forEach { it.dispose() }
+
                     // Reconfigure if environments changed.
                     if (lastEnvironments.size != resolvedEnvironments.size || lastEnvironments != resolvedEnvironments) {
                         context.logger.info("Workspaces have changed, reconfiguring CLI: $resolvedEnvironments")
-                        cli.configSsh(resolvedEnvironments.map { it.asPairOfWorkspaceAndAgent() }.toSet())
+                        cli.configSsh(resolvedEnvironments.mapNotNull { it.toWorkspaceAgentPairOrNull() }.toSet())
                     }
 
                     environments.update {
@@ -181,7 +185,7 @@ class CoderRemoteProvider(
                     sshConfigTrigger.onReceive { shouldTrigger ->
                         if (shouldTrigger) {
                             context.logger.debug("workspace poller waked up because it should reconfigure the ssh configurations")
-                            cli.configSsh(lastEnvironments.map { it.asPairOfWorkspaceAndAgent() }.toSet())
+                            cli.configSsh(lastEnvironments.mapNotNull { it.toWorkspaceAgentPairOrNull() }.toSet())
                         }
                     }
                     workspaceRefreshTrigger.onReceive { shouldTrigger ->
@@ -203,8 +207,8 @@ class CoderRemoteProvider(
      * Resolves workspace agents into remote environments.
      *
      * For each workspace:
-     * - If running, uses agents from the latest build resources
-     * - If not running, fetches resources separately
+     * - If running, uses agents from the latest build resources.
+     * - If not running, creates a workspace-only environment without resolving agents.
      *
      * @return a sorted list of resolved remote environments
      */
@@ -224,23 +228,23 @@ class CoderRemoteProvider(
         }
         coderHeaderPage.resetError()
         return workspaces.flatMap { ws ->
-            // Agents are not included in workspaces that are off
-            // so fetch them separately.
-            val resources = when (ws.latestBuild.status) {
-                WorkspaceStatus.RUNNING -> ws.latestBuild.resources
-                else -> emptyList()
-            }.ifEmpty {
-                client.resources(ws)
+            if (ws.latestBuild.status != WorkspaceStatus.RUNNING) {
+                return@flatMap listOf(
+                    lastEnvironments.firstOrNull { it.id == ws.name }
+                        ?.also { it.update(ws, null) }
+                        ?: CoderRemoteEnvironment(context, client, cli, workspaceRefreshTrigger, ws, null)
+                )
             }
-            resources
+
+            ws.latestBuild.resources
                 .flatMap { it.agents ?: emptyList() }
                 .distinctBy { it.name }
                 .map { agent ->
                     lastEnvironments.firstOrNull { it.id == "${ws.name}.${agent.name}" }
                         ?.also {
                             // If we have an environment already, update that.
                             it.update(ws, agent)
-                        } ?: CoderRemoteEnvironment(context, client, cli, ws, agent)
+                        } ?: CoderRemoteEnvironment(context, client, cli, workspaceRefreshTrigger, ws, agent)
                 }
 
         }.sortedBy { it.id }
@@ -288,6 +292,7 @@ class CoderRemoteProvider(
         softClose()
         client = null
         cli = null
+        lastEnvironments.forEach { it.dispose() }
         lastEnvironments.clear()
         environments.value = LoadableState.Value(emptyList())
         isInitialized.update { false }
@@ -590,7 +595,11 @@ class CoderRemoteProvider(
                     onTokenRefreshed = ::onTokenRefreshed,
                 )
             } catch (ex: Exception) {
-                context.logAndShowError("Error encountered while setting up Coder", "Failed to set up Coder", ex)
+                context.logAndShowError(
+                    "Error encountered while setting up Coder",
+                    "Failed to set up Coder: ${ex.message}",
+                    ex
+                )
             } finally {
                 firstRun = false
             }

src/main/kotlin/com/coder/toolbox/sdk/CoderRestClient.kt

@@ -21,7 +21,6 @@ import com.coder.toolbox.sdk.v2.models.User
 import com.coder.toolbox.sdk.v2.models.Workspace
 import com.coder.toolbox.sdk.v2.models.WorkspaceBuild
 import com.coder.toolbox.sdk.v2.models.WorkspaceBuildReason
-import com.coder.toolbox.sdk.v2.models.WorkspaceResource
 import com.coder.toolbox.sdk.v2.models.WorkspaceTransition
 import com.coder.toolbox.util.ReloadableTlsContext
 import com.coder.toolbox.views.state.CoderOAuthSessionContext
@@ -202,31 +201,6 @@ open class CoderRestClient(
         }
     }
 
-    /**
-     * Retrieves resources for the specified workspace.  The workspaces response
-     * does not include agents when the workspace is off so this can be used to
-     * get them instead, just like `coder config-ssh` does (otherwise we risk
-     * removing hosts from the SSH config when they are off).
-     * @throws [APIResponseException].
-     */
-    suspend fun resources(workspace: Workspace): List<WorkspaceResource> {
-        val resourcesResponse = callWithRetry {
-            retroRestClient.templateVersionResources(workspace.latestBuild.templateVersionID)
-        }
-        if (!resourcesResponse.isSuccessful) {
-            throw APIResponseException(
-                "retrieve resources for ${workspace.name}",
-                url,
-                resourcesResponse.code(),
-                resourcesResponse.parseErrorBody(moshi)
-            )
-        }
-
-        return requireNotNull(resourcesResponse.body()) {
-            "Successful response returned null body or workspace resources"
-        }
-    }
-
     suspend fun buildInfo(): BuildInfo {
         val buildInfoResponse = callWithRetry { retroRestClient.buildInfo() }
         if (!buildInfoResponse.isSuccessful) {

src/main/kotlin/com/coder/toolbox/sdk/v2/CoderV2RestFacade.kt

@@ -7,7 +7,6 @@ import com.coder.toolbox.sdk.v2.models.Template
 import com.coder.toolbox.sdk.v2.models.User
 import com.coder.toolbox.sdk.v2.models.Workspace
 import com.coder.toolbox.sdk.v2.models.WorkspaceBuild
-import com.coder.toolbox.sdk.v2.models.WorkspaceResource
 import com.coder.toolbox.sdk.v2.models.WorkspacesResponse
 import retrofit2.Response
 import retrofit2.http.Body
@@ -69,8 +68,4 @@ interface CoderV2RestFacade {
         @Path("templateID") templateID: UUID,
     ): Response<Template>
 
-    @GET("api/v2/templateversions/{templateID}/resources")
-    suspend fun templateVersionResources(
-        @Path("templateID") templateID: UUID,
-    ): Response<List<WorkspaceResource>>
 }