impl: remember the ssh connection state

And try to automatically establish the connections after an expired token was refreshed (by going again
through the login sequence)

Users will still have to launch a new a remote IDE if it was opened while a session expired.

- resolves #121

Author: fioan89

CHANGELOG.md

@@ -2,6 +2,11 @@
 
 ## Unreleased
 
+### Changed
+
+- the plugin will now remember the SSH connection state for each workspace, and it will try to automatically
+  establish it after an expired token was refreshed.
+
 ### Fixed
 
 - show errors when the Toolbox is visible again after being minimized.

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

@@ -234,28 +234,38 @@ class CoderRemoteEnvironment(
      * The contents are provided by the SSH view provided by Toolbox, all we
      * have to do is provide it a host name.
      */
-    override suspend
-    fun getContentsView(): EnvironmentContentsView = EnvironmentView(
+    override suspend fun getContentsView(): EnvironmentContentsView = EnvironmentView(
         client.url,
         cli,
         workspace,
         agent
     )
 
     /**
-     * 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.
+     * Automatically launches the SSH connection if the workspace is visible, is ready and there is no
+     * connection already established.
      */
     override fun setVisible(visibilityState: EnvironmentVisibilityState) {
-        if (wsRawStatus.ready() && visibilityState.contentsVisible == true && isConnected.value == false) {
+        if (visibilityState.contentsVisible) {
+            startSshConnection()
+        }
+    }
+
+    /**
+     * Launches the SSH connection if the workspace is ready and there is no connection already established.
+     *
+     * Returns true if the SSH connection was scheduled to start, false otherwise.
+     */
+    fun startSshConnection(): Boolean {
+        if (wsRawStatus.ready() && !isConnected.value) {
             context.cs.launch {
                 connectionRequest.update {
                     true
                 }
             }
+            return true
         }
+        return false
     }
 
     override fun getDeleteEnvironmentConfirmationParams(): DeleteEnvironmentConfirmationParams? {
@@ -298,6 +308,8 @@ class CoderRemoteEnvironment(
         }
     }
 
+    fun isConnected(): Boolean = isConnected.value
+
     /**
      * An environment is equal if it has the same ID.
      */

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

@@ -3,6 +3,7 @@ package com.coder.toolbox
 import com.coder.toolbox.browser.browse
 import com.coder.toolbox.cli.CoderCLIManager
 import com.coder.toolbox.sdk.CoderRestClient
+import com.coder.toolbox.sdk.ex.APIResponseException
 import com.coder.toolbox.sdk.v2.models.WorkspaceStatus
 import com.coder.toolbox.util.CoderProtocolHandler
 import com.coder.toolbox.util.DialogUi
@@ -20,7 +21,6 @@ import com.jetbrains.toolbox.api.core.util.LoadableState
 import com.jetbrains.toolbox.api.localization.LocalizableString
 import com.jetbrains.toolbox.api.remoteDev.ProviderVisibilityState
 import com.jetbrains.toolbox.api.remoteDev.RemoteProvider
-import com.jetbrains.toolbox.api.remoteDev.RemoteProviderEnvironment
 import com.jetbrains.toolbox.api.ui.actions.ActionDelimiter
 import com.jetbrains.toolbox.api.ui.actions.ActionDescription
 import com.jetbrains.toolbox.api.ui.components.UiPage
@@ -66,7 +66,8 @@ class CoderRemoteProvider(
     private val isInitialized: MutableStateFlow<Boolean> = MutableStateFlow(false)
     private var coderHeaderPage = NewEnvironmentPage(context, context.i18n.pnotr(context.deploymentUrl.toString()))
     private val linkHandler = CoderProtocolHandler(context, dialogUi, isInitialized)
-    override val environments: MutableStateFlow<LoadableState<List<RemoteProviderEnvironment>>> = MutableStateFlow(
+
+    override val environments: MutableStateFlow<LoadableState<List<CoderRemoteEnvironment>>> = MutableStateFlow(
         LoadableState.Loading
     )
 
@@ -127,7 +128,7 @@ class CoderRemoteProvider(
                 environments.update {
                     LoadableState.Value(resolvedEnvironments.toList())
                 }
-                if (isInitialized.value == false) {
+                if (!isInitialized.value) {
                     context.logger.info("Environments for ${client.url} are now initialized")
                     isInitialized.update {
                         true
@@ -137,6 +138,21 @@ class CoderRemoteProvider(
                     clear()
                     addAll(resolvedEnvironments.sortedBy { it.id })
                 }
+
+                if (WorkspaceConnectionManager.shouldEstablishWorkspaceConnections) {
+                    WorkspaceConnectionManager.allConnected().forEach { wsId ->
+                        val env = lastEnvironments.firstOrNull() { it.id == wsId }
+                        if (env != null && !env.isConnected()) {
+                            context.logger.info("Establishing lost SSH connection for workspace with id $wsId")
+                            if (!env.startSshConnection()) {
+                                context.logger.info("Can't establish lost SSH connection for workspace with id $wsId")
+                            }
+                        }
+                    }
+                    WorkspaceConnectionManager.reset()
+                }
+
+                WorkspaceConnectionManager.collectStatuses(lastEnvironments)
             } catch (_: CancellationException) {
                 context.logger.debug("${client.url} polling loop canceled")
                 break
@@ -147,6 +163,9 @@ class CoderRemoteProvider(
                     client.setupSession()
                 } else {
                     context.logger.error(ex, "workspace polling error encountered, trying to auto-login")
+                    if (ex is APIResponseException && ex.isTokenExpired) {
+                        WorkspaceConnectionManager.shouldEstablishWorkspaceConnections = true
+                    }
                     close()
                     // force auto-login
                     firstRun = true
@@ -179,6 +198,7 @@ class CoderRemoteProvider(
         // Keep the URL and token to make it easy to log back in, but set
         // rememberMe to false so we do not try to automatically log in.
         context.secrets.rememberMe = false
+        WorkspaceConnectionManager.reset()
         close()
     }
 
@@ -361,7 +381,7 @@ class CoderRemoteProvider(
         goToEnvironmentsPage()
     }
 
-    private fun MutableStateFlow<LoadableState<List<RemoteProviderEnvironment>>>.showLoadingMessage() {
+    private fun MutableStateFlow<LoadableState<List<CoderRemoteEnvironment>>>.showLoadingMessage() {
         this.update {
             LoadableState.Loading
         }

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

@@ -0,0 +1,22 @@
+package com.coder.toolbox
+
+object WorkspaceConnectionManager {
+    private val workspaceConnectionState = mutableMapOf<String, Boolean>()
+
+    var shouldEstablishWorkspaceConnections = false
+
+    fun allConnected(): Set<String> = workspaceConnectionState.filter { it.value }.map { it.key }.toSet()
+
+    fun collectStatuses(workspaces: Set<CoderRemoteEnvironment>) {
+        workspaces.forEach { register(it.id, it.isConnected()) }
+    }
+
+    private fun register(wsId: String, isConnected: Boolean) {
+        workspaceConnectionState[wsId] = isConnected
+    }
+
+    fun reset() {
+        workspaceConnectionState.clear()
+        shouldEstablishWorkspaceConnections = false
+    }
+}

src/main/kotlin/com/coder/toolbox/sdk/ex/APIResponseException.kt

@@ -10,6 +10,7 @@ class APIResponseException(action: String, url: URL, code: Int, errorResponse: A
 
     val reason = errorResponse?.detail
     val isUnauthorized = HttpURLConnection.HTTP_UNAUTHORIZED == code
+    val isTokenExpired = isUnauthorized && reason?.contains("API key expired") == true
 
     companion object {
         private fun formatToPretty(