fix: token input screen is closed after switching between Toolbox and browser

- 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

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

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

@@ -7,11 +7,11 @@ 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
@@ -67,7 +67,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())
@@ -189,7 +189,7 @@ class CoderRemoteProvider(
         if (username != null) {
             return dropDownFactory(context.i18n.pnotr(username)) {
                 logout()
-                context.ui.showUiPage(getOverrideUiPage()!!)
+                context.envPageManager.showPluginEnvironmentsPage()
             }
         }
         return null
@@ -215,6 +215,7 @@ class CoderRemoteProvider(
         environments.value = LoadableState.Value(emptyList())
         isInitialized.update { false }
         client = null
+        AuthWizardState.resetSteps()
     }
 
     override val svgIcon: SvgIcon =
@@ -306,7 +307,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,40 +318,20 @@ 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"
 
-    /**
-     * 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 ?: ""
@@ -378,22 +360,4 @@ class CoderRemoteProvider(
             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()
-        }
-    }
 }

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,42 @@ 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>? = 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
+        }
+    }
+
+}

src/main/kotlin/com/coder/toolbox/util/URLExtensions.kt

@@ -4,7 +4,7 @@ import java.net.IDN
 import java.net.URI
 import java.net.URL
 
-fun String.toURL(): URL = URL(this)
+fun String.toURL(): URL = URI.create(this).toURL()
 
 fun URL.withPath(path: String): URL = URL(
     this.protocol,

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

@@ -0,0 +1,94 @@
+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.views.state.AuthWizardState
+import com.coder.toolbox.views.state.WizardStep
+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
+
+class AuthWizardPage(
+    private val context: CoderToolboxContext,
+    private val isAutoLogin: Boolean = false,
+    onConnect: (
+        client: CoderRestClient,
+        cli: CoderCLIManager,
+    ) -> Unit,
+) : CoderPage(context, context.i18n.ptrl("Authenticate to Coder")) {
+    private val signInStep = SignInStep(context)
+    private val tokenStep = TokenStep(context)
+    private val connectStep = ConnectStep(context, this::notify, onConnect)
+
+
+    /**
+     * Fields for this page, displayed in order.
+     */
+    override val fields: MutableStateFlow<List<UiField>> = MutableStateFlow(emptyList())
+    override val actionButtons: MutableStateFlow<List<RunnableActionDescription>> = MutableStateFlow(emptyList())
+
+    override fun beforeShow() {
+        displaySteps()
+    }
+
+    private fun displaySteps() {
+        when (AuthWizardState.currentStep()) {
+            WizardStep.URL_REQUEST -> {
+                fields.update {
+                    listOf(signInStep.panel)
+                }
+                actionButtons.update {
+                    listOf(
+                        Action(context.i18n.ptrl("Sign In"), closesPage = false, actionBlock = {
+                            if (signInStep.onNext()) {
+                                displaySteps()
+                            }
+                        })
+                    )
+                }
+                signInStep.onVisible()
+            }
+
+            WizardStep.TOKEN_REQUEST -> {
+                fields.update {
+                    listOf(tokenStep.panel)
+                }
+                actionButtons.update {
+                    listOf(
+                        Action(context.i18n.ptrl("Back"), closesPage = false, actionBlock = {
+                            tokenStep.onBack()
+                            displaySteps()
+                        }),
+                        Action(context.i18n.ptrl("Connect"), closesPage = false, actionBlock = {
+                            if (tokenStep.onNext()) {
+                                displaySteps()
+                            }
+                        })
+                    )
+                }
+                tokenStep.onVisible()
+            }
+
+            WizardStep.LOGIN -> {
+                fields.update {
+                    listOf(connectStep.panel)
+                }
+                actionButtons.update {
+                    listOf(
+                        Action(context.i18n.ptrl("Back"), closesPage = false, actionBlock = {
+                            if (isAutoLogin) {
+                                AuthWizardState.resetSteps()
+                            } else {
+                                connectStep.onBack()
+                            }
+                            displaySteps()
+                        })
+                    )
+                }
+                connectStep.onVisible()
+            }
+        }
+    }
+}

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

@@ -5,10 +5,7 @@ import com.jetbrains.toolbox.api.core.ui.icons.SvgIcon
 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.UiField
 import com.jetbrains.toolbox.api.ui.components.UiPage
-import com.jetbrains.toolbox.api.ui.components.ValidationErrorField
-import java.util.function.Consumer
 
 /**
  * Base page that handles the icon, displaying error notifications, and
@@ -25,19 +22,10 @@ abstract class CoderPage(
     title: LocalizableString,
     showIcon: Boolean = true,
 ) : UiPage(title) {
-    /**
-     * An error to display on the page.
-     *
-     * The current assumption is you only have one field per page.
-     */
-    protected var errorField: ValidationErrorField? = null
 
     /** Toolbox uses this to show notifications on the page. */
     private var notifier: ((Throwable) -> Unit)? = null
 
-    /** Let Toolbox know the fields should be updated. */
-    protected var listener: Consumer<UiField?>? = null
-
     /** Stores errors until the notifier is attached. */
     private var errorBuffer: MutableList<Throwable> = mutableListOf()
 
@@ -76,14 +64,6 @@ abstract class CoderPage(
             errorBuffer.clear()
         }
     }
-
-    /**
-     * Set/unset the field error and update the form.
-     */
-    protected fun updateError(error: String?) {
-        errorField = error?.let { ValidationErrorField(context.i18n.pnotr(error)) }
-        listener?.accept(null) // Make Toolbox get the fields again.
-    }
 }
 
 /**