From 4a7a287121e3dfaf7284e46ca2ff83d97b980ffd Mon Sep 17 00:00:00 2001
From: Faur Ioan-Aurel <fioan89@gmail.com>
Date: Wed, 9 Apr 2025 22:53:20 +0300
Subject: [PATCH] fix: ssh connection to a workspace is established only once

- right now ssh automatically connects to the workspace when a user selects it from the main env. page.
  However, if the user manually disconnects from the hamburger menu, the SSH is no longer automatically
  established afterward.
- the reason for this behavior was that the connection request flag that triggered the ssh connection behind
  the scenes was never reset. In order to reset we have to know when a user hits the disconnect button.
- fortunately Toolbox provides two callbacks, one before establishing the ssh connection and another one right
  after disconnect. Even though we don't know the reason for disconnect it's still a good opportunity to
  reset the request flag.

- resolves #38
---
 CHANGELOG.md                                  |  4 +++
 .../coder/toolbox/CoderRemoteEnvironment.kt   | 27 +++++++++++++++----
 2 files changed, 26 insertions(+), 5 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index be5a605..b37cfe1 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,10 @@
 
 ## Unreleased
 
+### Fixed
+
+- SSH connection to a Workspace is no longer established only once 
+
 ### Changed
 
 - action buttons on the token input step were swapped to achieve better keyboard navigation 
diff --git a/src/main/kotlin/com/coder/toolbox/CoderRemoteEnvironment.kt b/src/main/kotlin/com/coder/toolbox/CoderRemoteEnvironment.kt
index b7184a0..fe5f039 100644
--- a/src/main/kotlin/com/coder/toolbox/CoderRemoteEnvironment.kt
+++ b/src/main/kotlin/com/coder/toolbox/CoderRemoteEnvironment.kt
@@ -9,6 +9,8 @@ import com.coder.toolbox.sdk.v2.models.WorkspaceAgent
 import com.coder.toolbox.util.withPath
 import com.coder.toolbox.views.Action
 import com.coder.toolbox.views.EnvironmentView
+import com.jetbrains.toolbox.api.remoteDev.AfterDisconnectHook
+import com.jetbrains.toolbox.api.remoteDev.BeforeConnectionHook
 import com.jetbrains.toolbox.api.remoteDev.DeleteEnvironmentConfirmationParams
 import com.jetbrains.toolbox.api.remoteDev.EnvironmentVisibilityState
 import com.jetbrains.toolbox.api.remoteDev.RemoteProviderEnvironment
@@ -35,7 +37,7 @@ class CoderRemoteEnvironment(
     private val client: CoderRestClient,
     private var workspace: Workspace,
     private var agent: WorkspaceAgent,
-) : RemoteProviderEnvironment("${workspace.name}.${agent.name}") {
+) : RemoteProviderEnvironment("${workspace.name}.${agent.name}"), BeforeConnectionHook, AfterDisconnectHook {
     private var wsRawStatus = WorkspaceAndAgentStatus.from(workspace, agent)
 
     override var name: String = "${workspace.name}.${agent.name}"
@@ -109,6 +111,21 @@ class CoderRemoteEnvironment(
         return actions
     }
 
+    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
+    }
+
+    override fun afterDisconnect() {
+        this.connectionRequest.update { false }
+        this.isConnected = false
+        context.logger.info("Disconnected from $id")
+    }
+
     /**
      * Update the workspace/agent status to the listeners, if it has changed.
      */
@@ -140,7 +157,8 @@ class CoderRemoteEnvironment(
         agent
     )
 
-    override val connectionRequest: MutableStateFlow<Boolean>? = MutableStateFlow(false)
+    private var isConnected = false
+    override val connectionRequest: MutableStateFlow<Boolean> = MutableStateFlow(false)
 
     /**
      * Does nothing.  In theory, we could do something like start the workspace
@@ -149,10 +167,9 @@ class CoderRemoteEnvironment(
      * to be much value.
      */
     override fun setVisible(visibilityState: EnvironmentVisibilityState) {
-        if (wsRawStatus.ready() && visibilityState.contentsVisible == true && visibilityState.isBackendConnected == false) {
-            context.logger.info("Connecting to $id...")
+        if (wsRawStatus.ready() && visibilityState.contentsVisible == true && isConnected == false) {
             context.cs.launch {
-                connectionRequest?.update {
+                connectionRequest.update {
                     true
                 }
             }