fix: remember token when switching deployments

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

@@ -7,6 +7,10 @@
 - improved workspace status reporting (icon and colors) when it is failed, stopping, deleting, stopped or when we are
   establishing the SSH connection.
 
+### Fixed
+
+- tokens are now remembered after switching between multiple deployments
+
 ## 0.2.2 - 2025-05-21
 
 ### Added

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

@@ -335,6 +335,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/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/views/AuthWizardPage.kt

@@ -3,6 +3,7 @@ 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.AuthContext
 import com.coder.toolbox.views.state.AuthWizardState
 import com.coder.toolbox.views.state.WizardStep
 import com.jetbrains.toolbox.api.ui.actions.RunnableActionDescription
@@ -23,9 +24,17 @@ 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 authContext: AuthContext = AuthContext()
+    private val signInStep = SignInStep(context, authContext, this::notify)
+    private val tokenStep = TokenStep(context, authContext)
+    private val connectStep = ConnectStep(
+        context,
+        authContext,
+        shouldAutoLogin,
+        this::notify,
+        this::displaySteps, onConnect
+    )
 
 
     /**

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
@@ -25,6 +25,7 @@ private const val USER_HIT_THE_BACK_BUTTON = "User hit the back button"
  */
 class ConnectStep(
     private val context: CoderToolboxContext,
+    private val authContext: AuthContext,
     private val shouldAutoLogin: StateFlow<Boolean>,
     private val notify: (String, Throwable) -> Unit,
     private val refreshWizard: () -> Unit,
@@ -49,18 +50,18 @@ class ConnectStep(
         errorField.textState.update {
             context.i18n.pnotr("")
         }
+        if (authContext.isNotReadyForAuth()) return
 
-        val url = context.deploymentUrl?.first?.toURL()
-        statusField.textState.update { context.i18n.pnotr("Connecting to ${url?.host}...") }
+        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
+        val url = authContext.url
+        val token = authContext.token
         if (url == null) {
             errorField.textState.update { context.i18n.ptrl("URL is required") }
             return
@@ -96,6 +97,8 @@ class ConnectStep(
                 // allows interleaving with the back/cancel action
                 yield()
                 onConnect(client, cli)
+
+                authContext.reset()
                 AuthWizardState.resetSteps()
             } catch (ex: CancellationException) {
                 if (ex.message != USER_HIT_THE_BACK_BUTTON) {

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

@@ -2,31 +2,34 @@ package com.coder.toolbox.views
 
 import com.coder.toolbox.CoderToolboxContext
 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
 import com.jetbrains.toolbox.api.ui.components.RowGroup
 import com.jetbrains.toolbox.api.ui.components.TextField
 import com.jetbrains.toolbox.api.ui.components.TextType
 import com.jetbrains.toolbox.api.ui.components.ValidationErrorField
 import kotlinx.coroutines.flow.update
 import java.net.MalformedURLException
+import java.net.URI
 
 /**
  * A page with a field for providing the Coder deployment URL.
  *
  * Populates with the provided URL, at which point the user can accept or
  * enter their own.
  */
-class SignInStep(private val context: CoderToolboxContext, private val notify: (String, Throwable) -> Unit) :
+class SignInStep(
+    private val context: CoderToolboxContext,
+    private val authContext: AuthContext,
+    private val notify: (String, Throwable) -> Unit
+) :
     WizardStep {
     private val urlField = TextField(context.i18n.ptrl("Deployment URL"), "", TextType.General)
-    private val descriptionField = LabelField(context.i18n.pnotr(""))
     private val errorField = ValidationErrorField(context.i18n.pnotr(""))
 
     override val panel: RowGroup = RowGroup(
         RowGroup.RowField(urlField),
-        RowGroup.RowField(descriptionField),
         RowGroup.RowField(errorField)
     )
 
@@ -37,11 +40,7 @@ class SignInStep(private val context: CoderToolboxContext, private val notify: (
             context.i18n.pnotr("")
         }
         urlField.textState.update {
-            context.deploymentUrl?.first ?: ""
-        }
-
-        descriptionField.textState.update {
-            context.i18n.pnotr(context.deploymentUrl?.second?.description("URL") ?: "")
+            context.secrets.lastDeploymentURL
         }
     }
 
@@ -62,7 +61,7 @@ class SignInStep(private val context: CoderToolboxContext, private val notify: (
             notify("URL is invalid", e)
             return false
         }
-        context.secrets.lastDeploymentURL = url
+        authContext.url = URI.create(url).toURL()
         AuthWizardState.goToNextStep()
         return true
     }

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

@@ -3,9 +3,9 @@ package com.coder.toolbox.views
 import com.coder.toolbox.CoderToolboxContext
 import com.coder.toolbox.util.toURL
 import com.coder.toolbox.util.withPath
+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
 import com.jetbrains.toolbox.api.ui.components.LinkField
 import com.jetbrains.toolbox.api.ui.components.RowGroup
 import com.jetbrains.toolbox.api.ui.components.TextField
@@ -20,15 +20,16 @@ import kotlinx.coroutines.flow.update
  * Populate with the provided token, at which point the user can accept or
  * enter their own.
  */
-class TokenStep(private val context: CoderToolboxContext) : WizardStep {
+class TokenStep(
+    private val context: CoderToolboxContext,
+    private val authContext: AuthContext
+) : WizardStep {
     private val tokenField = TextField(context.i18n.ptrl("Token"), "", TextType.Password)
-    private val descriptionField = LabelField(context.i18n.pnotr(""))
     private val linkField = LinkField(context.i18n.ptrl("Get a token"), "")
     private val errorField = ValidationErrorField(context.i18n.pnotr(""))
 
     override val panel: RowGroup = RowGroup(
         RowGroup.RowField(tokenField),
-        RowGroup.RowField(descriptionField),
         RowGroup.RowField(linkField),
         RowGroup.RowField(errorField)
     )
@@ -39,13 +40,11 @@ class TokenStep(private val context: CoderToolboxContext) : WizardStep {
             context.i18n.pnotr("")
         }
         tokenField.textState.update {
-            context.getToken(context.deploymentUrl?.first)?.first ?: ""
-        }
-        descriptionField.textState.update {
-            context.i18n.pnotr(
-                context.getToken(context.deploymentUrl?.first)?.second?.description("token")
-                    ?: "No existing token for ${context.deploymentUrl} found."
-            )
+            if (authContext.hasUrl()) {
+                context.secrets.tokenFor(authContext.url!!) ?: ""
+            } else {
+                ""
+            }
         }
         (linkField.urlState as MutableStateFlow).update {
             context.deploymentUrl?.first?.toURL()?.withPath("/login?redirect=%2Fcli-auth")?.toString() ?: ""
@@ -59,7 +58,7 @@ class TokenStep(private val context: CoderToolboxContext) : WizardStep {
             return false
         }
 
-        context.secrets.lastToken = token
+        authContext.token = token
         AuthWizardState.goToNextStep()
         return true
     }

src/main/kotlin/com/coder/toolbox/views/state/AuthContext.kt

@@ -0,0 +1,17 @@
+package com.coder.toolbox.views.state
+
+import java.net.URL
+
+data class AuthContext(
+    var url: URL? = null,
+    var token: String? = null
+) {
+    fun hasUrl(): Boolean = url != null
+
+    fun isNotReadyForAuth(): Boolean = !(hasUrl() && token != null)
+
+    fun reset() {
+        url = null
+        token = null
+    }
+}