fix: rendering glitches when a Workspace is stopped while an SSH connection is alive

Toolbox raises a class cast exception when Workspaces are stopped while the SSH connection
is running. After the workspace was stopped Toolbox refused to show widget with some weird glitches
on the screen.
The fix in this case is to safely disconnect the SSH before sending the stop command to the workspace.
The code will wait at most 10 seconds for the disconnect to happen, and only after that send the stop.

Author: fioan89

CHANGELOG.md

@@ -6,6 +6,10 @@
 
 - ssh configuration is simplified, background hostnames have been discarded.
 
+### Fixed
+
+- rendering glitches when a Workspace is stopped while SSH connection is alive
+
 ## 0.2.0 - 2025-04-24
 
 ### Added

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

@@ -7,6 +7,7 @@ import com.coder.toolbox.sdk.CoderRestClient
 import com.coder.toolbox.sdk.ex.APIResponseException
 import com.coder.toolbox.sdk.v2.models.Workspace
 import com.coder.toolbox.sdk.v2.models.WorkspaceAgent
+import com.coder.toolbox.util.waitForFalseWithTimeout
 import com.coder.toolbox.util.withPath
 import com.coder.toolbox.views.Action
 import com.coder.toolbox.views.EnvironmentView
@@ -43,6 +44,10 @@ class CoderRemoteEnvironment(
     private var wsRawStatus = WorkspaceAndAgentStatus.from(workspace, agent)
 
     override var name: String = "${workspace.name}.${agent.name}"
+
+    private var isConnected: MutableStateFlow<Boolean> = MutableStateFlow(false)
+    override val connectionRequest: MutableStateFlow<Boolean> = MutableStateFlow(false)
+
     override val state: MutableStateFlow<RemoteEnvironmentState> =
         MutableStateFlow(wsRawStatus.toRemoteEnvironmentState(context))
     override val description: MutableStateFlow<EnvironmentDescription> =
@@ -106,6 +111,16 @@ class CoderRemoteEnvironment(
             } else {
                 actions.add(Action(context.i18n.ptrl("Stop")) {
                     context.cs.launch {
+                        if (isConnected.value) {
+                            connectionRequest.update {
+                                false
+                            }
+
+                            if (isConnected.waitForFalseWithTimeout(10.seconds) == null) {
+                                context.logger.warn("The SSH connection to workspace $name could not be dropped in time, going to stop the workspace while the SSH connection is live")
+                            }
+                        }
+
                         val build = client.stopWorkspace(workspace)
                         update(workspace.copy(latestBuild = build), agent)
                     }
@@ -121,12 +136,12 @@ class CoderRemoteEnvironment(
 
     override fun beforeConnection() {
         context.logger.info("Connecting to $id...")
-        this.isConnected = true
+        isConnected.update { true }
     }
 
     override fun afterDisconnect() {
         this.connectionRequest.update { false }
-        this.isConnected = false
+        isConnected.update { false }
         context.logger.info("Disconnected from $id")
     }
 
@@ -161,17 +176,14 @@ class CoderRemoteEnvironment(
         agent
     )
 
-    private var isConnected = false
-    override val connectionRequest: MutableStateFlow<Boolean> = MutableStateFlow(false)
-
     /**
      * Does nothing.  In theory, we could do something like start the workspace
      * when you click into the workspace, but you would still need to press
      * "connect" anyway before the content is populated so there does not seem
      * to be much value.
      */
     override fun setVisible(visibilityState: EnvironmentVisibilityState) {
-        if (wsRawStatus.ready() && visibilityState.contentsVisible == true && isConnected == false) {
+        if (wsRawStatus.ready() && visibilityState.contentsVisible == true && isConnected.value == false) {
             context.cs.launch {
                 connectionRequest.update {
                     true

src/main/kotlin/com/coder/toolbox/util/CoderProtocolHandler.kt

@@ -13,7 +13,6 @@ import com.jetbrains.toolbox.api.localization.LocalizableString
 import kotlinx.coroutines.TimeoutCancellationException
 import kotlinx.coroutines.delay
 import kotlinx.coroutines.flow.StateFlow
-import kotlinx.coroutines.flow.first
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.time.withTimeout
 import java.net.HttpURLConnection
@@ -331,9 +330,4 @@ private fun CoderToolboxContext.popupPluginMainPage() {
     this.envPageManager.showPluginEnvironmentsPage(true)
 }
 
-/**
- * Suspends the coroutine until first true value is received.
- */
-suspend fun StateFlow<Boolean>.waitForTrue() = this.first { it }
-
 class MissingArgumentException(message: String, ex: Throwable? = null) : IllegalArgumentException(message, ex)

src/main/kotlin/com/coder/toolbox/util/StateFlowExtensions.kt

@@ -0,0 +1,22 @@
+package com.coder.toolbox.util
+
+import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.first
+import kotlinx.coroutines.withTimeoutOrNull
+import kotlin.time.Duration
+
+/**
+ * Suspends the coroutine until first true value is received.
+ */
+suspend fun StateFlow<Boolean>.waitForTrue() = this.first { it }
+
+/**
+ * Suspends the coroutine until first false value is received.
+ */
+suspend fun StateFlow<Boolean>.waitForFalseWithTimeout(duration: Duration): Boolean? {
+    if (!this.value) return false
+
+    return withTimeoutOrNull(duration) {
+        this@waitForFalseWithTimeout.first { !it }
+    }
+}