impl: remember the ssh connection state

CHANGELOG.md

@@ -2,6 +2,15 @@
 
 ## 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.
+
 ## 0.2.3 - 2025-05-26
 
 ### Changed

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,10 +66,18 @@ 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
     )
 
+    private val visibilityState = MutableStateFlow(
+        ProviderVisibilityState(
+            applicationVisible = false,
+            providerVisible = false
+        )
+    )
+
     /**
      * With the provided client, start polling for workspaces.  Every time a new
      * workspace is added, reconfigure SSH using the provided cli (including the
@@ -120,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
@@ -130,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
@@ -140,7 +163,12 @@ 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
                     goToEnvironmentsPage()
                     break
                 }
@@ -170,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()
     }
 
@@ -263,7 +292,11 @@ class CoderRemoteProvider(
      *        a place to put a timer ("last updated 10 seconds ago" for example)
      *        and a manual refresh button.
      */
-    override fun setVisible(visibilityState: ProviderVisibilityState) {}
+    override fun setVisible(visibility: ProviderVisibilityState) {
+        visibilityState.update {
+            visibility
+        }
+    }
 
     /**
      * Handle incoming links (like from the dashboard).
@@ -309,7 +342,7 @@ class CoderRemoteProvider(
                     if (autologin && lastDeploymentURL.isNotBlank() && (lastToken.isNotBlank() || !settings.requireTokenAuth)) {
                         try {
                             AuthWizardState.goToStep(WizardStep.LOGIN)
-                            return AuthWizardPage(context, settingsPage, true, ::onConnect)
+                            return AuthWizardPage(context, settingsPage, visibilityState, true, ::onConnect)
                         } catch (ex: Exception) {
                             errorBuffer.add(ex)
                         }
@@ -319,7 +352,7 @@ class CoderRemoteProvider(
             firstRun = false
 
             // Login flow.
-            val authWizard = AuthWizardPage(context, settingsPage, false, ::onConnect)
+            val authWizard = AuthWizardPage(context, settingsPage, visibilityState, onConnect = ::onConnect)
             // We might have navigated here due to a polling error.
             errorBuffer.forEach {
                 authWizard.notify("Error encountered", it)
@@ -348,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

@@ -8,8 +8,9 @@ import java.net.URL
 class APIResponseException(action: String, url: URL, code: Int, errorResponse: ApiErrorResponse?) :
     IOException(formatToPretty(action, url, code, errorResponse)) {
 
-
+    val reason = errorResponse?.detail
     val isUnauthorized = HttpURLConnection.HTTP_UNAUTHORIZED == code
+    val isTokenExpired = isUnauthorized && reason?.contains("API key expired") == true
 
     companion object {
         private fun formatToPretty(

src/main/kotlin/com/coder/toolbox/views/AuthWizardPage.kt

@@ -3,18 +3,23 @@ package com.coder.toolbox.views
 import com.coder.toolbox.CoderToolboxContext
 import com.coder.toolbox.cli.CoderCLIManager
 import com.coder.toolbox.sdk.CoderRestClient
+import com.coder.toolbox.sdk.ex.APIResponseException
 import com.coder.toolbox.util.toURL
 import com.coder.toolbox.views.state.AuthContext
 import com.coder.toolbox.views.state.AuthWizardState
 import com.coder.toolbox.views.state.WizardStep
+import com.jetbrains.toolbox.api.remoteDev.ProviderVisibilityState
 import com.jetbrains.toolbox.api.ui.actions.RunnableActionDescription
 import com.jetbrains.toolbox.api.ui.components.UiField
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.update
+import kotlinx.coroutines.launch
+import java.util.UUID
 
 class AuthWizardPage(
     private val context: CoderToolboxContext,
     private val settingsPage: CoderSettingsPage,
+    private val visibilityState: MutableStateFlow<ProviderVisibilityState>,
     initialAutoLogin: Boolean = false,
     onConnect: (
         client: CoderRestClient,
@@ -42,6 +47,8 @@ class AuthWizardPage(
     override val fields: MutableStateFlow<List<UiField>> = MutableStateFlow(emptyList())
     override val actionButtons: MutableStateFlow<List<RunnableActionDescription>> = MutableStateFlow(emptyList())
 
+    private val errorBuffer = mutableListOf<Throwable>()
+
     init {
         if (shouldAutoLogin.value) {
             AuthContext.url = context.secrets.lastDeploymentURL.toURL()
@@ -51,6 +58,12 @@ class AuthWizardPage(
 
     override fun beforeShow() {
         displaySteps()
+        if (errorBuffer.isNotEmpty() && visibilityState.value.applicationVisible) {
+            errorBuffer.forEach {
+                showError(it)
+            }
+            errorBuffer.clear()
+        }
     }
 
     private fun displaySteps() {
@@ -113,4 +126,34 @@ class AuthWizardPage(
             }
         }
     }
+
+    /**
+     * Show an error as a popup on this page.
+     */
+    fun notify(logPrefix: String, ex: Throwable) {
+        context.logger.error(ex, logPrefix)
+        if (!visibilityState.value.applicationVisible) {
+            context.logger.debug("Toolbox is not yet visible, scheduling error to be displayed later")
+            errorBuffer.add(ex)
+            return
+        }
+        showError(ex)
+    }
+
+    private fun showError(ex: Throwable) {
+        val textError = if (ex is APIResponseException) {
+            if (!ex.reason.isNullOrBlank()) {
+                ex.reason
+            } else ex.message
+        } else ex.message
+
+        context.cs.launch {
+            context.ui.showSnackbar(
+                UUID.randomUUID().toString(),
+                context.i18n.ptrl("Error encountered during authentication"),
+                context.i18n.pnotr(textError ?: ""),
+                context.i18n.ptrl("Dismiss")
+            )
+        }
+    }
 }

src/main/kotlin/com/coder/toolbox/views/CoderPage.kt

@@ -6,8 +6,6 @@ import com.jetbrains.toolbox.api.core.ui.icons.SvgIcon.IconType
 import com.jetbrains.toolbox.api.localization.LocalizableString
 import com.jetbrains.toolbox.api.ui.actions.RunnableActionDescription
 import com.jetbrains.toolbox.api.ui.components.UiPage
-import kotlinx.coroutines.launch
-import java.util.UUID
 
 /**
  * Base page that handles the icon, displaying error notifications, and
@@ -39,21 +37,6 @@ abstract class CoderPage(
         SvgIcon(byteArrayOf(), type = IconType.Masked)
     }
 
-    /**
-     * Show an error as a popup on this page.
-     */
-    fun notify(logPrefix: String, ex: Throwable) {
-        context.logger.error(ex, logPrefix)
-        context.cs.launch {
-            context.ui.showSnackbar(
-                UUID.randomUUID().toString(),
-                context.i18n.pnotr(logPrefix),
-                context.i18n.pnotr(ex.message ?: ""),
-                context.i18n.ptrl("Dismiss")
-            )
-        }
-    }
-
     companion object {
         fun emptyPage(ctx: CoderToolboxContext): UiPage = UiPage(ctx.i18n.pnotr(""))
     }

src/main/resources/localization/defaultMessages.po

@@ -137,4 +137,7 @@ msgid "Network Status"
 msgstr ""
 
 msgid "Create workspace"
+msgstr ""
+
+msgid "Error encountered during authentication"
 msgstr ""