fix: rendering glitches when a Workspace is stopped

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) {
+        [email protected] { !it }
+    }
+}