fix: rendering glitches when a Workspace is stopped (#102)

… 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.

- resolves #98

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,8 @@ class CoderRemoteEnvironment(
             } else {
                 actions.add(Action(context.i18n.ptrl("Stop")) {
                     context.cs.launch {
+                        tryStopSshConnection()
+
                         val build = client.stopWorkspace(workspace)
                         update(workspace.copy(latestBuild = build), agent)
                     }
@@ -115,18 +122,30 @@ class CoderRemoteEnvironment(
         return actions
     }
 
+    private suspend fun tryStopSshConnection() {
+        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")
+            }
+        }
+    }
+
     override fun getBeforeConnectionHooks(): List<BeforeConnectionHook> = listOf(this)
 
     override fun getAfterDisconnectHooks(): List<AfterDisconnectHook> = listOf(this)
 
     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 +180,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 }
+    }
+}