Prevent duplicate CLI sign-ins and show the correct wizard when switching deployments

CHANGELOG.md

@@ -19,6 +19,8 @@ It follows [Keep a Changelog](https://keepachangelog.com/en/1.1.0/).
 
 - SSH connections being blocked while non-blocking startup scripts run or after startup scripts fail or time out.
 - Automatic connection not starting when a workspace agent becomes connectable after the initial workspace load.
+- Duplicate CLI sign-ins that could send the setup wizard back to the token screen after a successful connection.
+- Opening a workspace link for another deployment could leave the previous setup page on screen until Toolbox regained focus.
 
 ### Security
 

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

@@ -34,6 +34,7 @@ import com.coder.toolbox.views.CoderSetupWizardPage
 import com.coder.toolbox.views.NewEnvironmentPage
 import com.coder.toolbox.views.SuspendBiConsumer
 import com.coder.toolbox.views.state.CoderOAuthSessionContext
+import com.coder.toolbox.views.state.ConnectionOrigin
 import com.coder.toolbox.views.state.Credentials
 import com.coder.toolbox.views.state.PageRouter
 import com.coder.toolbox.views.state.PendingOAuthConnection
@@ -51,11 +52,13 @@ import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.Job
 import kotlinx.coroutines.cancelAndJoin
 import kotlinx.coroutines.channels.Channel
+import kotlinx.coroutines.currentCoroutineContext
 import kotlinx.coroutines.ensureActive
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.update
 import kotlinx.coroutines.isActive
+import kotlinx.coroutines.job
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.selects.onTimeout
 import kotlinx.coroutines.selects.select
@@ -418,6 +421,11 @@ class CoderRemoteProvider(
      * Also called as part of our own logout.
      */
     override fun close() {
+        resetConnection()
+        router.clear()
+    }
+
+    private fun resetConnection() {
         val sessionIds = lastEnvironments.currentSessionIds()
         softClose()
         client = null
@@ -428,7 +436,6 @@ class CoderRemoteProvider(
         environments.value = LoadableState.Value(emptyList())
         isInitialized.update { false }
         accountDropdownField.visibility.update { false }
-        router.clear()
         context.logger.info(sessionIds, "Coder plugin is now closed")
         context.logger.flush("provider_closed")
     }
@@ -811,7 +818,7 @@ class CoderRemoteProvider(
         linkHandler.handle(params, url, client, cli)
     }
 
-    private fun connectFromLink(params: Map<String, String>, url: URL, token: String?) {
+    private suspend fun connectFromLink(params: Map<String, String>, url: URL, token: String?) {
         val credentials = token?.let { Credentials.Token(it) } ?: Credentials.MTls
         val wizard = CoderSetupWizardPage.connectStep(
             context, settingsPage,
@@ -820,10 +827,12 @@ class CoderRemoteProvider(
             onConnect = onConnect.andThen(deferredLinkHandler(params, url)),
             onTokenRefreshed = ::onTokenRefreshed,
         )
-        // The router tears down any in-flight wizard and installs the new one before
-        // showPluginEnvironmentsPage() asks getOverrideUiPage() for the page to display.
+        // Replace any in-flight wizard before showPluginEnvironmentsPage() requests a refresh.
+        // If Toolbox calls getOverrideUiPage(), it must get the wizard for this link.
         router.navigate(wizard)
         context.popupPluginMainPage()
+        // An already visible provider can keep showing the old wizard, so show the new one explicitly.
+        context.ui.showUiPageSuspending(wizard)
     }
 
     /** Returns a [SuspendBiConsumer] that handles the given link parameters in a background coroutine. */
@@ -849,21 +858,16 @@ 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
-     * list.
-     */
+    /** Show interactive setup when needed. Saved sessions restore behind the workspace list. */
     override fun getOverrideUiPage(): UiPage? {
         // Show the setup wizard if one is already scheduled.
-        router.activePage?.let { return it }
+        router.activePage?.let { return it.pageForDisplay() }
 
         // Let the default workspace UI render if the HTTP client is initialized.
         if (client != null) return null
 
         // Otherwise, schedule our own setup wizard.
-        return router.getOrCreate { buildSetupWizard() }
+        return router.getOrCreate { buildSetupWizard() }.pageForDisplay()
     }
 
     /**
@@ -887,6 +891,7 @@ class CoderRemoteProvider(
                     credentials = credentials,
                     onConnect = onConnect,
                     onTokenRefreshed = ::onTokenRefreshed,
+                    origin = ConnectionOrigin.SAVED_SESSION,
                 )
             } catch (ex: Exception) {
                 context.logger.logAndShowError(
@@ -915,9 +920,13 @@ class CoderRemoteProvider(
     fun canAutoLogin(): Boolean = autoSetupCredentials(context.deploymentUrl) != null
 
     private val onConnect: SuspendBiConsumer<CoderRestClient, CoderCLIManager> = SuspendBiConsumer { client, cli ->
+        router.finishSignIn(currentCoroutineContext().job) { startSession(client, cli) }
+    }
+
+    private fun startSession(client: CoderRestClient, cli: CoderCLIManager) {
         // Store the URL and token for use next time.
         storeSession(client)
-        close()
+        resetConnection()
         this.client = client
         this.cli = cli
         lastEnvironments.forEach { it.updateClientAndCli(client, cli) }

src/main/kotlin/com/coder/toolbox/sdk/CoderRestClient.kt

@@ -95,11 +95,7 @@ open class CoderRestClient(
 
         val interceptors = buildList {
             if (context.settingsStore.requiresTokenAuth) {
-                val oauthOrApiToken = oauthContext?.tokenResponse?.accessToken ?: token
-                if (oauthOrApiToken.isNullOrBlank()) {
-                    throw IllegalStateException("OAuth or API token is required for $url deployment")
-                }
-                add(Interceptors.tokenAuth(oauthOrApiToken))
+                add(Interceptors.tokenAuth(requireSessionToken()))
             }
             add((Interceptors.userAgent(pluginVersion)))
             add(Interceptors.externalHeaders(context, url))
@@ -123,6 +119,12 @@ open class CoderRestClient(
                 .build().create(CoderV2RestFacade::class.java)
     }
 
+    internal fun requireSessionToken(): String {
+        val sessionToken = oauthContext?.tokenResponse?.accessToken ?: token
+        check(!sessionToken.isNullOrBlank()) { "OAuth or API token is required for $url deployment" }
+        return sessionToken
+    }
+
     /**
      * Load information about the current user and the build version.
      *

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

@@ -4,12 +4,14 @@ import com.coder.toolbox.CoderToolboxContext
 import com.coder.toolbox.cli.CoderCLIManager
 import com.coder.toolbox.sdk.CoderRestClient
 import com.coder.toolbox.views.state.CoderOAuthSessionContext
+import com.coder.toolbox.views.state.ConnectionOrigin
 import com.coder.toolbox.views.state.Credentials
 import com.coder.toolbox.views.state.PendingOAuthConnection
 import com.coder.toolbox.views.state.WizardModel
 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.CoroutineStart
 import kotlinx.coroutines.Job
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.update
@@ -19,7 +21,6 @@ import java.net.URL
 class CoderSetupWizardPage private constructor(
     private val context: CoderToolboxContext,
     private val settingsPage: CoderSettingsPage,
-    private var autoLogin: Boolean = false,
     onConnect: SuspendBiConsumer<CoderRestClient, CoderCLIManager>,
     onTokenRefreshed: (suspend (url: URL, oauthSessionCtx: CoderOAuthSessionContext) -> Unit)? = null
 ) : CoderPage(MutableStateFlow(context.i18n.ptrl("Setting up Coder")), false) {
@@ -38,36 +39,51 @@ class CoderSetupWizardPage private constructor(
         onTokenRefreshed = onTokenRefreshed
     )
     private var stateCollectJob: Job? = null
+    @Volatile
+    private var disposed = false
 
     /**
      * Fields for this page, displayed in order.
      */
     override val fields: MutableStateFlow<List<UiField>> = MutableStateFlow(emptyList())
     override val actionButtons: MutableStateFlow<List<RunnableActionDescription>> = MutableStateFlow(emptyList())
 
+    fun pageForDisplay(): CoderSetupWizardPage? = when (model.currentStep()) {
+        WizardStep.CONNECT -> if (model.connectionOrigin == ConnectionOrigin.SAVED_SESSION) null else this
+        WizardStep.SIGNED_IN -> null
+        else -> this
+    }
+
+    internal fun startSavedSessionRestore() {
+        if (model.connectionOrigin == ConnectionOrigin.SAVED_SESSION) connectStep.startConnection()
+    }
+
     override fun beforeShow() {
-        stateCollectJob?.cancel()
-        stateCollectJob = context.cs.launch {
+        prepareStepObserver()?.start()
+    }
+
+    @Synchronized
+    private fun prepareStepObserver(): Job? {
+        if (disposed || stateCollectJob?.isCompleted == false) return null
+        val job = context.cs.launch(start = CoroutineStart.LAZY) {
             model.step.collect { step ->
                 context.logger.info("Wizard step changed to $step")
-                displaySteps()
+                displayStep(step)
             }
         }
+        stateCollectJob = job
+        return job
     }
 
-    private fun displaySteps() {
-        when (model.currentStep()) {
+    private fun displayStep(step: WizardStep) {
+        when (step) {
             WizardStep.URL_REQUEST -> {
                 fields.update {
                     listOf(deploymentUrlStep.panel)
                 }
                 actionButtons.update {
                     listOf(
-                        Action(context, "Next", closesPage = false, actionBlock = {
-                            if (deploymentUrlStep.onNext()) {
-                                displaySteps()
-                            }
-                        }),
+                        Action(context, "Next", closesPage = false) { deploymentUrlStep.onNext() },
                         settingsAction
                     )
                 }
@@ -80,16 +96,9 @@ class CoderSetupWizardPage private constructor(
                 }
                 actionButtons.update {
                     listOf(
-                        Action(context, "Connect", closesPage = false, actionBlock = {
-                            if (tokenStep.onNext()) {
-                                displaySteps()
-                            }
-                        }),
+                        Action(context, "Connect", closesPage = false) { tokenStep.onNext() },
                         settingsAction,
-                        Action(context, "Back", closesPage = false, actionBlock = {
-                            tokenStep.onBack()
-                            displaySteps()
-                        })
+                        Action(context, "Back", closesPage = false) { tokenStep.onBack() }
                     )
                 }
                 tokenStep.onVisible()
@@ -102,14 +111,15 @@ class CoderSetupWizardPage private constructor(
                 actionButtons.update {
                     listOf(
                         settingsAction,
-                        Action(context, "Back", closesPage = false, actionBlock = {
-                            connectStep.onBack()
-                            displaySteps()
-                        })
+                        Action(context, "Back", closesPage = false) { connectStep.onBack() }
                     )
                 }
                 connectStep.onVisible()
             }
+
+            // Toolbox can bring this wizard back after Settings or during startup.
+            // Ask for the workspace list when the completed wizard is actually shown.
+            WizardStep.SIGNED_IN -> context.envPageManager.showPluginEnvironmentsPage(false)
         }
     }
 
@@ -120,16 +130,18 @@ class CoderSetupWizardPage private constructor(
     }
 
     private fun navigateBackFromConnect() {
-        if (autoLogin) {
-            autoLogin = false
-            model.clearFormData()
-            model.goToFirst()
+        if (model.connectionOrigin == ConnectionOrigin.FORM) {
+            model.goTo(if (context.settingsStore.requiresTokenAuth) WizardStep.TOKEN_REQUEST else WizardStep.URL_REQUEST)
             return
         }
-        if (context.settingsStore.requiresTokenAuth) {
-            model.goToPrevious()
-        } else {
-            model.goToFirst()
+        val revealForm = model.connectionOrigin == ConnectionOrigin.SAVED_SESSION
+        model.clearFormData()
+        model.goToFirst()
+        model.connectionOrigin = ConnectionOrigin.FORM
+        if (revealForm) {
+            context.cs.launch {
+                if (!disposed) context.ui.showUiPageSuspending(this@CoderSetupWizardPage)
+            }
         }
     }
 
@@ -138,13 +150,17 @@ class CoderSetupWizardPage private constructor(
      * when the wizard is being replaced (e.g. by a deep link to a different
      * deployment) so its connect job doesn't keep running and clobber the new one.
      */
+    @Synchronized
     fun dispose() {
-        stateCollectJob?.cancel()
+        disposed = true
+        afterHide()
         connectStep.dispose()
     }
 
+    @Synchronized
     override fun afterHide() {
         stateCollectJob?.cancel()
+        stateCollectJob = null
     }
 
     companion object {
@@ -166,12 +182,13 @@ class CoderSetupWizardPage private constructor(
             credentials: Credentials,
             onConnect: SuspendBiConsumer<CoderRestClient, CoderCLIManager>,
             onTokenRefreshed: (suspend (url: URL, oauthSessionCtx: CoderOAuthSessionContext) -> Unit)? = null,
+            origin: ConnectionOrigin = ConnectionOrigin.SUPPLIED_CREDENTIALS,
         ): CoderSetupWizardPage = CoderSetupWizardPage(
             context, settingsPage,
-            autoLogin = true,
             onConnect = onConnect,
             onTokenRefreshed = onTokenRefreshed,
         ).apply {
+            model.connectionOrigin = origin
             model.url = url
             when (credentials) {
                 is Credentials.MTls -> Unit