fix: token input screen is closed after switching between Toolbox and browser (#72)

- rough draft to fix UI state management in the authentication flow
which today has 3 pages. If user closes
Toolbox in any of these three pages (for example to go and copy the
token from a browser), then when it comes back in Toolbox does not
remember which was the last visible UiPage.
- until JetBrains improves Toolbox state management, we can work around
the problem by having only one
UiPage with three "steps" in it, similar to a wizard. With this approach
we can have complete control over
  the state of the page.

- to be noted that I've also looked over two other approaches. The first
idea was to manage the stat
ourselves, but that didn’t work out as Toolbox doesn’t clearly tell us
when the user clicks the Back
button vs. when they close the window. So we can’t reliably figure out
which page to show when it reopens.
- another option was changing the auth flow entirely and adding custom
redirect URLs for Toolbox plugins.
But that would only work with certain Coder versions, which might not be
ideal.

- resolves #45

Author: fioan89

Diff for: CHANGELOG.md

@@ -2,6 +2,10 @@
 
 ## Unreleased
 
+### Fixed
+
+- Toolbox remembers the authentication page that was last visible on the screen 
+
 ## 0.1.2 - 2025-04-04
 
 ### Fixed

Diff for: gradle.properties

@@ -1,3 +1,3 @@
-version=0.1.2
+version=0.1.3
 group=com.coder.toolbox
 name=coder-toolbox

Diff for: src/main/kotlin/com/coder/toolbox/CoderRemoteEnvironment.kt

@@ -74,26 +74,35 @@ class CoderRemoteEnvironment(
         if (wsRawStatus.canStart()) {
             if (workspace.outdated) {
                 actions.add(Action(context.i18n.ptrl("Update and start")) {
-                    val build = client.updateWorkspace(workspace)
-                    update(workspace.copy(latestBuild = build), agent)
+                    context.cs.launch {
+                        val build = client.updateWorkspace(workspace)
+                        update(workspace.copy(latestBuild = build), agent)
+                    }
                 })
             } else {
                 actions.add(Action(context.i18n.ptrl("Start")) {
-                    val build = client.startWorkspace(workspace)
-                    update(workspace.copy(latestBuild = build), agent)
+                    context.cs.launch {
+                        val build = client.startWorkspace(workspace)
+                        update(workspace.copy(latestBuild = build), agent)
+
+                    }
                 })
             }
         }
         if (wsRawStatus.canStop()) {
             if (workspace.outdated) {
                 actions.add(Action(context.i18n.ptrl("Update and restart")) {
-                    val build = client.updateWorkspace(workspace)
-                    update(workspace.copy(latestBuild = build), agent)
+                    context.cs.launch {
+                        val build = client.updateWorkspace(workspace)
+                        update(workspace.copy(latestBuild = build), agent)
+                    }
                 })
             } else {
                 actions.add(Action(context.i18n.ptrl("Stop")) {
-                    val build = client.stopWorkspace(workspace)
-                    update(workspace.copy(latestBuild = build), agent)
+                    context.cs.launch {
+                        val build = client.stopWorkspace(workspace)
+                        update(workspace.copy(latestBuild = build), agent)
+                    }
                 })
             }
         }

Diff for: src/main/kotlin/com/coder/toolbox/CoderRemoteProvider.kt

@@ -3,15 +3,14 @@ package com.coder.toolbox
 import com.coder.toolbox.cli.CoderCLIManager
 import com.coder.toolbox.sdk.CoderRestClient
 import com.coder.toolbox.sdk.v2.models.WorkspaceStatus
-import com.coder.toolbox.settings.SettingSource
 import com.coder.toolbox.util.CoderProtocolHandler
 import com.coder.toolbox.util.DialogUi
 import com.coder.toolbox.views.Action
+import com.coder.toolbox.views.AuthWizardPage
 import com.coder.toolbox.views.CoderSettingsPage
-import com.coder.toolbox.views.ConnectPage
 import com.coder.toolbox.views.NewEnvironmentPage
-import com.coder.toolbox.views.SignInPage
-import com.coder.toolbox.views.TokenPage
+import com.coder.toolbox.views.state.AuthWizardState
+import com.coder.toolbox.views.state.WizardStep
 import com.jetbrains.toolbox.api.core.ui.icons.SvgIcon
 import com.jetbrains.toolbox.api.core.ui.icons.SvgIcon.IconType
 import com.jetbrains.toolbox.api.core.util.LoadableState
@@ -32,7 +31,6 @@ import kotlinx.coroutines.selects.onTimeout
 import kotlinx.coroutines.selects.select
 import java.net.SocketTimeoutException
 import java.net.URI
-import java.net.URL
 import kotlin.coroutines.cancellation.CancellationException
 import kotlin.time.Duration.Companion.seconds
 import kotlin.time.TimeSource
@@ -67,7 +65,7 @@ class CoderRemoteProvider(
     // On the first load, automatically log in if we can.
     private var firstRun = true
     private val isInitialized: MutableStateFlow<Boolean> = MutableStateFlow(false)
-    private var coderHeaderPage = NewEnvironmentPage(context, context.i18n.pnotr(getDeploymentURL()?.first ?: ""))
+    private var coderHeaderPage = NewEnvironmentPage(context, context.i18n.pnotr(context.deploymentUrl?.first ?: ""))
     private val linkHandler = CoderProtocolHandler(context, dialogUi, isInitialized)
     override val environments: MutableStateFlow<LoadableState<List<RemoteProviderEnvironment>>> = MutableStateFlow(
         LoadableState.Value(emptyList())
@@ -177,7 +175,7 @@ class CoderRemoteProvider(
     private fun logout() {
         // 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"
+        context.secrets.rememberMe = false
         close()
     }
 
@@ -189,7 +187,7 @@ class CoderRemoteProvider(
         if (username != null) {
             return dropDownFactory(context.i18n.pnotr(username)) {
                 logout()
-                context.ui.showUiPage(getOverrideUiPage()!!)
+                context.envPageManager.showPluginEnvironmentsPage()
             }
         }
         return null
@@ -215,6 +213,7 @@ class CoderRemoteProvider(
         environments.value = LoadableState.Value(emptyList())
         isInitialized.update { false }
         client = null
+        AuthWizardState.resetSteps()
     }
 
     override val svgIcon: SvgIcon =
@@ -293,7 +292,7 @@ class CoderRemoteProvider(
     /**
      * Return the sign-in page if we do not have a valid client.
 
-     * Otherwise return null, which causes Toolbox to display the environment
+     * Otherwise, return null, which causes Toolbox to display the environment
      * list.
      */
     override fun getOverrideUiPage(): UiPage? {
@@ -306,7 +305,8 @@ class CoderRemoteProvider(
                 context.secrets.lastDeploymentURL.let { lastDeploymentURL ->
                     if (autologin && lastDeploymentURL.isNotBlank() && (lastToken.isNotBlank() || !settings.requireTokenAuth)) {
                         try {
-                            return createConnectPage(URL(lastDeploymentURL), lastToken)
+                            AuthWizardState.goToStep(WizardStep.LOGIN)
+                            return AuthWizardPage(context, true, ::onConnect)
                         } catch (ex: Exception) {
                             autologinEx = ex
                         }
@@ -316,84 +316,29 @@ class CoderRemoteProvider(
             firstRun = false
 
             // Login flow.
-            val signInPage =
-                SignInPage(context, getDeploymentURL()) { deploymentURL ->
-                    context.ui.showUiPage(
-                        TokenPage(
-                            context,
-                            deploymentURL,
-                            getToken(deploymentURL)
-                        ) { selectedToken ->
-                            context.ui.showUiPage(createConnectPage(deploymentURL, selectedToken))
-                        },
-                    )
-                }
-
+            val authWizard = AuthWizardPage(context, false, ::onConnect)
             // We might have tried and failed to automatically log in.
-            autologinEx?.let { signInPage.notify("Error logging in", it) }
+            autologinEx?.let { authWizard.notify("Error logging in", it) }
             // We might have navigated here due to a polling error.
-            pollError?.let { signInPage.notify("Error fetching workspaces", it) }
+            pollError?.let { authWizard.notify("Error fetching workspaces", it) }
 
-            return signInPage
+            return authWizard
         }
         return null
     }
 
-    private fun shouldDoAutoLogin(): Boolean = firstRun && context.secrets.rememberMe == "true"
+    private fun shouldDoAutoLogin(): Boolean = firstRun && context.secrets.rememberMe == true
 
-    /**
-     * Create a connect page that starts polling and resets the UI on success.
-     */
-    private fun createConnectPage(deploymentURL: URL, token: String?): ConnectPage = ConnectPage(
-        context,
-        deploymentURL,
-        token,
-        ::goToEnvironmentsPage,
-    ) { client, cli ->
+    private fun onConnect(client: CoderRestClient, cli: CoderCLIManager) {
         // Store the URL and token for use next time.
         context.secrets.lastDeploymentURL = client.url.toString()
         context.secrets.lastToken = client.token ?: ""
         // Currently we always remember, but this could be made an option.
-        context.secrets.rememberMe = "true"
+        context.secrets.rememberMe = true
         this.client = client
         pollError = null
         pollJob?.cancel()
         pollJob = poll(client, cli)
         goToEnvironmentsPage()
     }
-
-    /**
-     * Try to find a token.
-     *
-     * Order of preference:
-     *
-     * 1. Last used token, if it was for this deployment.
-     * 2. Token on disk for this deployment.
-     * 3. Global token for Coder, if it matches the deployment.
-     */
-    private fun getToken(deploymentURL: URL): Pair<String, SettingSource>? = context.secrets.lastToken.let {
-        if (it.isNotBlank() && context.secrets.lastDeploymentURL == deploymentURL.toString()) {
-            it to SettingSource.LAST_USED
-        } else {
-            settings.token(deploymentURL)
-        }
-    }
-
-    /**
-     * Try to find a URL.
-     *
-     * In order of preference:
-     *
-     * 1. Last used URL.
-     * 2. URL in settings.
-     * 3. CODER_URL.
-     * 4. URL in global cli config.
-     */
-    private fun getDeploymentURL(): Pair<String, SettingSource>? = context.secrets.lastDeploymentURL.let {
-        if (it.isNotBlank()) {
-            it to SettingSource.LAST_USED
-        } else {
-            context.settingsStore.defaultURL()
-        }
-    }
 }

Diff for: src/main/kotlin/com/coder/toolbox/CoderToolboxContext.kt

@@ -1,7 +1,9 @@
 package com.coder.toolbox
 
+import com.coder.toolbox.settings.SettingSource
 import com.coder.toolbox.store.CoderSecretsStore
 import com.coder.toolbox.store.CoderSettingsStore
+import com.coder.toolbox.util.toURL
 import com.jetbrains.toolbox.api.core.diagnostics.Logger
 import com.jetbrains.toolbox.api.localization.LocalizableStringFactory
 import com.jetbrains.toolbox.api.remoteDev.connection.ClientHelper
@@ -20,4 +22,43 @@ data class CoderToolboxContext(
     val i18n: LocalizableStringFactory,
     val settingsStore: CoderSettingsStore,
     val secrets: CoderSecretsStore
-)
+) {
+    /**
+     * Try to find a URL.
+     *
+     * In order of preference:
+     *
+     * 1. Last used URL.
+     * 2. URL in settings.
+     * 3. CODER_URL.
+     * 4. URL in global cli config.
+     */
+    val deploymentUrl: Pair<String, SettingSource>?
+        get() = this.secrets.lastDeploymentURL.let {
+            if (it.isNotBlank()) {
+                it to SettingSource.LAST_USED
+            } else {
+                this.settingsStore.defaultURL()
+            }
+        }
+
+    /**
+     * Try to find a token.
+     *
+     * Order of preference:
+     *
+     * 1. Last used token, if it was for this deployment.
+     * 2. Token on disk for this deployment.
+     * 3. Global token for Coder, if it matches the deployment.
+     */
+    fun getToken(deploymentURL: String?): Pair<String, SettingSource>? = this.secrets.lastToken.let {
+        if (it.isNotBlank() && this.secrets.lastDeploymentURL == deploymentURL) {
+            it to SettingSource.LAST_USED
+        } else {
+            if (deploymentURL != null) {
+                this.settingsStore.token(deploymentURL.toURL())
+            } else null
+        }
+    }
+
+}