fix: remember token when switching deployments (#120)

If we log in on deployment 1, then log out and login to deployment 2 and
then in the same session we try to log in back to deployment 1, the
token is no longer valid. The plugin will associate with deployment 1
the token from the second deployment.

There is an overly complicated block of code inherited from Gateway
plugin with multiple fallback sequences for both the deployment url and
token from multiple sources (secrets store, data dir config, env,
etc...). This fix simplifies the approach, we only store the url and the
token in the secrets store, the token is always associated to a
hostname. If there is no previous URL to remember (like the first time
login) we default to https://dev.coder.com/ and empty token.

Author: fioan89

CHANGELOG.md

@@ -10,6 +10,7 @@
 ### Fixed
 
 - url on the main page is now refreshed when switching between multiple deployments (via logout/login or URI handling)
+- tokens are now remembered after switching between multiple deployments
 
 ## 0.2.2 - 2025-05-21
 

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

@@ -64,7 +64,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(context.deploymentUrl?.first ?: ""))
+    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(
         LoadableState.Loading
@@ -336,6 +336,7 @@ class CoderRemoteProvider(
         // Store the URL and token for use next time.
         context.secrets.lastDeploymentURL = client.url.toString()
         context.secrets.lastToken = client.token ?: ""
+        context.secrets.storeTokenFor(client.url, context.secrets.lastToken)
         // Currently we always remember, but this could be made an option.
         context.secrets.rememberMe = true
         this.client = client

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

@@ -1,6 +1,5 @@
 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
@@ -13,6 +12,7 @@ import com.jetbrains.toolbox.api.remoteDev.states.EnvironmentStateColorPalette
 import com.jetbrains.toolbox.api.remoteDev.ui.EnvironmentUiPageManager
 import com.jetbrains.toolbox.api.ui.ToolboxUi
 import kotlinx.coroutines.CoroutineScope
+import java.net.URL
 
 data class CoderToolboxContext(
     val ui: ToolboxUi,
@@ -37,31 +37,11 @@ data class CoderToolboxContext(
      * 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()
+    val deploymentUrl: URL
+        get() {
+            if (this.secrets.lastDeploymentURL.isNotBlank()) {
+                return this.secrets.lastDeploymentURL.toURL()
             }
+            return this.settingsStore.defaultURL.toURL()
         }
-
-    /**
-     * 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/browser/BrowserUtil.kt

@@ -1,12 +1,12 @@
 package com.coder.toolbox.browser
 
+import com.coder.toolbox.util.toURL
 import com.jetbrains.toolbox.api.core.os.LocalDesktopManager
-import java.net.URI
 
 
 suspend fun LocalDesktopManager.browse(rawUrl: String, errorHandler: suspend (BrowserException) -> Unit) {
     try {
-        val url = URI.create(rawUrl).toURL()
+        val url = rawUrl.toURL()
         this.openUrl(url)
     } catch (e: Exception) {
         errorHandler(

src/main/kotlin/com/coder/toolbox/settings/ReadOnlyCoderSettings.kt

@@ -10,7 +10,7 @@ interface ReadOnlyCoderSettings {
     /**
      * The default URL to show in the connection window.
      */
-    val defaultURL: String?
+    val defaultURL: String
 
     /**
      * Used to download the Coder CLI which is necessary to proxy SSH
@@ -116,16 +116,6 @@ interface ReadOnlyCoderSettings {
      */
     val networkInfoDir: String
 
-    /**
-     * The default URL to show in the connection window.
-     */
-    fun defaultURL(): Pair<String, SettingSource>?
-
-    /**
-     * Given a deployment URL, try to find a token for it if required.
-     */
-    fun token(deploymentURL: URL): Pair<String, SettingSource>?
-
     /**
      * Where the specified deployment should put its data.
      */

src/main/kotlin/com/coder/toolbox/store/CoderSecretsStore.kt

@@ -1,6 +1,7 @@
 package com.coder.toolbox.store
 
 import com.jetbrains.toolbox.api.core.PluginSecretStore
+import java.net.URL
 
 
 /**
@@ -26,4 +27,10 @@ class CoderSecretsStore(private val store: PluginSecretStore) {
     var rememberMe: Boolean
         get() = get("remember-me").toBoolean()
         set(value) = set("remember-me", value.toString())
+
+    fun tokenFor(url: URL): String? = store[url.host]
+
+    fun storeTokenFor(url: URL, token: String) {
+        store[url.host] = token
+    }
 }

src/main/kotlin/com/coder/toolbox/store/CoderSettingsStore.kt

@@ -3,7 +3,6 @@ package com.coder.toolbox.store
 import com.coder.toolbox.settings.Environment
 import com.coder.toolbox.settings.ReadOnlyCoderSettings
 import com.coder.toolbox.settings.ReadOnlyTLSSettings
-import com.coder.toolbox.settings.SettingSource
 import com.coder.toolbox.util.Arch
 import com.coder.toolbox.util.OS
 import com.coder.toolbox.util.expand
@@ -35,7 +34,7 @@ class CoderSettingsStore(
     ) : ReadOnlyTLSSettings
 
     // Properties implementation
-    override val defaultURL: String? get() = store[DEFAULT_URL]
+    override val defaultURL: String get() = store[DEFAULT_URL] ?: "https://dev.coder.com"
     override val binarySource: String? get() = store[BINARY_SOURCE]
     override val binaryDirectory: String? get() = store[BINARY_DIRECTORY]
     override val defaultCliBinaryNameByOsAndArch: String get() = getCoderCLIForOS(getOS(), getArch())
@@ -71,48 +70,6 @@ class CoderSettingsStore(
             .normalize()
             .toString()
 
-    /**
-     * The default URL to show in the connection window.
-     */
-    override fun defaultURL(): Pair<String, SettingSource>? {
-        val envURL = env.get(CODER_URL)
-        if (!defaultURL.isNullOrEmpty()) {
-            return defaultURL!! to SettingSource.SETTINGS
-        } else if (envURL.isNotBlank()) {
-            return envURL to SettingSource.ENVIRONMENT
-        } else {
-            val (configUrl, _) = readConfig(Path.of(globalConfigDir))
-            if (!configUrl.isNullOrBlank()) {
-                return configUrl to SettingSource.CONFIG
-            }
-        }
-        return null
-    }
-
-    /**
-     * Given a deployment URL, try to find a token for it if required.
-     */
-    override fun token(deploymentURL: URL): Pair<String, SettingSource>? {
-        // No need to bother if we do not need token auth anyway.
-        if (!requireTokenAuth) {
-            return null
-        }
-        // Try the deployment's config directory.  This could exist if someone
-        // has entered a URL that they are not currently connected to, but have
-        // connected to in the past.
-        val (_, deploymentToken) = readConfig(dataDir(deploymentURL).resolve("config"))
-        if (!deploymentToken.isNullOrBlank()) {
-            return deploymentToken to SettingSource.DEPLOYMENT_CONFIG
-        }
-        // Try the global config directory, in case they previously set up the
-        // CLI with this URL.
-        val (configUrl, configToken) = readConfig(Path.of(globalConfigDir))
-        if (configUrl == deploymentURL.toString() && !configToken.isNullOrBlank()) {
-            return configToken to SettingSource.CONFIG
-        }
-        return null
-    }
-
     /**
      * Where the specified deployment should put its data.
      */

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

@@ -3,6 +3,8 @@ 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.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.ui.actions.RunnableActionDescription
@@ -23,17 +25,30 @@ class AuthWizardPage(
     private val settingsAction = Action(context.i18n.ptrl("Settings"), actionBlock = {
         context.ui.showUiPage(settingsPage)
     })
+
     private val signInStep = SignInStep(context, this::notify)
     private val tokenStep = TokenStep(context)
-    private val connectStep = ConnectStep(context, shouldAutoLogin, this::notify, this::displaySteps, onConnect)
-
+    private val connectStep = ConnectStep(
+        context,
+        shouldAutoLogin,
+        this::notify,
+        this::displaySteps,
+        onConnect
+    )
 
     /**
      * Fields for this page, displayed in order.
      */
     override val fields: MutableStateFlow<List<UiField>> = MutableStateFlow(emptyList())
     override val actionButtons: MutableStateFlow<List<RunnableActionDescription>> = MutableStateFlow(emptyList())
 
+    init {
+        if (shouldAutoLogin.value) {
+            AuthContext.url = context.secrets.lastDeploymentURL.toURL()
+            AuthContext.token = context.secrets.lastToken
+        }
+    }
+
     override fun beforeShow() {
         displaySteps()
     }

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

@@ -5,7 +5,7 @@ import com.coder.toolbox.cli.CoderCLIManager
 import com.coder.toolbox.cli.ensureCLI
 import com.coder.toolbox.plugin.PluginManager
 import com.coder.toolbox.sdk.CoderRestClient
-import com.coder.toolbox.util.toURL
+import com.coder.toolbox.views.state.AuthContext
 import com.coder.toolbox.views.state.AuthWizardState
 import com.jetbrains.toolbox.api.localization.LocalizableString
 import com.jetbrains.toolbox.api.ui.components.LabelField
@@ -50,34 +50,38 @@ class ConnectStep(
             context.i18n.pnotr("")
         }
 
-        val url = context.deploymentUrl?.first?.toURL()
-        statusField.textState.update { context.i18n.pnotr("Connecting to ${url?.host}...") }
+        if (AuthContext.isNotReadyForAuth()) {
+            errorField.textState.update {
+                context.i18n.pnotr("URL and token were not properly configured. Please go back and provide a proper URL and token!")
+            }
+            return
+        }
+
+        statusField.textState.update { context.i18n.pnotr("Connecting to ${AuthContext.url!!.host}...") }
         connect()
     }
 
     /**
      * Try connecting to Coder with the provided URL and token.
      */
     private fun connect() {
-        val url = context.deploymentUrl?.first?.toURL()
-        val token = context.getToken(context.deploymentUrl?.first)?.first
-        if (url == null) {
+        if (!AuthContext.hasUrl()) {
             errorField.textState.update { context.i18n.ptrl("URL is required") }
             return
         }
 
-        if (token.isNullOrBlank()) {
+        if (!AuthContext.hasToken()) {
             errorField.textState.update { context.i18n.ptrl("Token is required") }
             return
         }
         signInJob?.cancel()
         signInJob = context.cs.launch {
             try {
-                statusField.textState.update { (context.i18n.ptrl("Authenticating to ${url.host}...")) }
+                statusField.textState.update { (context.i18n.ptrl("Authenticating to ${AuthContext.url!!.host}...")) }
                 val client = CoderRestClient(
                     context,
-                    url,
-                    token,
+                    AuthContext.url!!,
+                    AuthContext.token!!,
                     PluginManager.pluginInfo.version,
                 )
                 // allows interleaving with the back/cancel action
@@ -92,19 +96,20 @@ class ConnectStep(
                     yield()
                     cli.login(client.token)
                 }
-                statusField.textState.update { (context.i18n.ptrl("Successfully configured ${url.host}...")) }
+                statusField.textState.update { (context.i18n.ptrl("Successfully configured ${AuthContext.url!!.host}...")) }
                 // allows interleaving with the back/cancel action
                 yield()
-                onConnect(client, cli)
+                AuthContext.reset()
                 AuthWizardState.resetSteps()
+                onConnect(client, cli)
             } catch (ex: CancellationException) {
                 if (ex.message != USER_HIT_THE_BACK_BUTTON) {
-                    notify("Connection to ${url.host} was configured", ex)
+                    notify("Connection to ${AuthContext.url!!.host} was configured", ex)
                     onBack()
                     refreshWizard()
                 }
             } catch (ex: Exception) {
-                notify("Failed to configure ${url.host}", ex)
+                notify("Failed to configure ${AuthContext.url!!.host}", ex)
                 onBack()
                 refreshWizard()
             }
@@ -120,6 +125,7 @@ class ConnectStep(
             signInJob?.cancel(CancellationException(USER_HIT_THE_BACK_BUTTON))
         } finally {
             if (shouldAutoLogin.value) {
+                AuthContext.reset()
                 AuthWizardState.resetSteps()
                 context.secrets.rememberMe = false
             } else {