From 0ec198587a0f952f2fa4ac198c48a39fd5e153c5 Mon Sep 17 00:00:00 2001 From: Faur Ioan-Aurel Date: Wed, 11 Nov 2026 21:18:11 +0200 Subject: [PATCH 1/7] Prevent duplicate CLI sign-ins during setup and session restore Toolbox can show the connection step more than once during startup. A second CLI login could fail after the first succeeded and send the user back to token entry. The wizard now starts one connection attempt at a time. A retry waits for the cancelled attempt to finish, and only the active attempt can publish a session or report a failure. Saved sessions restore behind the workspace list. If restoration fails, the wizard opens the sign-in form so the user can reconnect. --- CHANGELOG.md | 1 + .../com/coder/toolbox/CoderRemoteProvider.kt | 26 +- .../toolbox/views/CoderSetupWizardPage.kt | 76 ++++-- .../com/coder/toolbox/views/ConnectStep.kt | 197 ++++++++------- .../coder/toolbox/views/DeploymentUrlStep.kt | 6 +- .../com/coder/toolbox/views/TokenStep.kt | 6 +- .../coder/toolbox/views/state/PageRouter.kt | 33 ++- .../coder/toolbox/views/state/WizardModel.kt | 16 +- .../coder/toolbox/CoderRemoteProviderTest.kt | 81 ++++-- .../toolbox/views/CoderSetupWizardPageTest.kt | 234 ++++++++++++++++++ .../toolbox/views/state/PageRouterTest.kt | 54 ++++ 11 files changed, 545 insertions(+), 185 deletions(-) create mode 100644 src/test/kotlin/com/coder/toolbox/views/CoderSetupWizardPageTest.kt diff --git a/CHANGELOG.md b/CHANGELOG.md index a0950256..58aeb6d7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -19,6 +19,7 @@ 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. ### Security diff --git a/src/main/kotlin/com/coder/toolbox/CoderRemoteProvider.kt b/src/main/kotlin/com/coder/toolbox/CoderRemoteProvider.kt index e7610993..8680e23c 100644 --- a/src/main/kotlin/com/coder/toolbox/CoderRemoteProvider.kt +++ b/src/main/kotlin/com/coder/toolbox/CoderRemoteProvider.kt @@ -51,11 +51,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 +420,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 +435,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") } @@ -849,21 +855,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 +888,7 @@ class CoderRemoteProvider( credentials = credentials, onConnect = onConnect, onTokenRefreshed = ::onTokenRefreshed, + restoreSavedSession = true, ) } catch (ex: Exception) { context.logger.logAndShowError( @@ -915,9 +917,13 @@ class CoderRemoteProvider( fun canAutoLogin(): Boolean = autoSetupCredentials(context.deploymentUrl) != null private val onConnect: SuspendBiConsumer = 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) } diff --git a/src/main/kotlin/com/coder/toolbox/views/CoderSetupWizardPage.kt b/src/main/kotlin/com/coder/toolbox/views/CoderSetupWizardPage.kt index 2481bd8e..a6dc515f 100644 --- a/src/main/kotlin/com/coder/toolbox/views/CoderSetupWizardPage.kt +++ b/src/main/kotlin/com/coder/toolbox/views/CoderSetupWizardPage.kt @@ -10,6 +10,7 @@ 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 @@ -38,6 +39,8 @@ class CoderSetupWizardPage private constructor( onTokenRefreshed = onTokenRefreshed ) private var stateCollectJob: Job? = null + @Volatile + private var disposed = false /** * Fields for this page, displayed in order. @@ -45,29 +48,43 @@ class CoderSetupWizardPage private constructor( override val fields: MutableStateFlow> = MutableStateFlow(emptyList()) override val actionButtons: MutableStateFlow> = MutableStateFlow(emptyList()) + fun pageForDisplay(): CoderSetupWizardPage? = when (model.currentStep()) { + // Restore saved credentials behind Toolbox's loading workspace list. + // A successful restore then needs no page navigation during window startup. + WizardStep.RESTORE_SESSION -> { + connectStep.onVisible() + null + } + WizardStep.SIGNED_IN -> null + else -> this + } + 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,36 +97,30 @@ 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() } - WizardStep.CONNECT -> { + WizardStep.CONNECT, WizardStep.RESTORE_SESSION -> { fields.update { listOf(connectStep.panel) } 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,14 +131,20 @@ class CoderSetupWizardPage private constructor( } private fun navigateBackFromConnect() { + val restoringSession = model.currentStep() == WizardStep.RESTORE_SESSION if (autoLogin) { autoLogin = false model.clearFormData() model.goToFirst() + if (restoringSession) { + context.cs.launch { + if (!disposed) context.ui.showUiPageSuspending(this@CoderSetupWizardPage) + } + } return } if (context.settingsStore.requiresTokenAuth) { - model.goToPrevious() + model.goTo(WizardStep.TOKEN_REQUEST) } else { model.goToFirst() } @@ -138,13 +155,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,6 +187,7 @@ class CoderSetupWizardPage private constructor( credentials: Credentials, onConnect: SuspendBiConsumer, onTokenRefreshed: (suspend (url: URL, oauthSessionCtx: CoderOAuthSessionContext) -> Unit)? = null, + restoreSavedSession: Boolean = false, ): CoderSetupWizardPage = CoderSetupWizardPage( context, settingsPage, autoLogin = true, @@ -178,7 +200,7 @@ class CoderSetupWizardPage private constructor( is Credentials.Token -> model.token = credentials.value is Credentials.OAuth -> model.oauthSession = credentials.session } - model.goTo(WizardStep.CONNECT) + model.goTo(if (restoreSavedSession) WizardStep.RESTORE_SESSION else WizardStep.CONNECT) } } } diff --git a/src/main/kotlin/com/coder/toolbox/views/ConnectStep.kt b/src/main/kotlin/com/coder/toolbox/views/ConnectStep.kt index 4adc782e..cbce7def 100644 --- a/src/main/kotlin/com/coder/toolbox/views/ConnectStep.kt +++ b/src/main/kotlin/com/coder/toolbox/views/ConnectStep.kt @@ -6,17 +6,25 @@ import com.coder.toolbox.cli.ensureCLI import com.coder.toolbox.oauth.OAuth2Client import com.coder.toolbox.plugin.PluginManager import com.coder.toolbox.sdk.CoderRestClient +import com.coder.toolbox.util.origin import com.coder.toolbox.views.state.CoderOAuthSessionContext import com.coder.toolbox.views.state.WizardModel +import com.coder.toolbox.views.state.WizardStep.CONNECT +import com.coder.toolbox.views.state.WizardStep.RESTORE_SESSION import com.jetbrains.toolbox.api.ui.components.LabelField import com.jetbrains.toolbox.api.ui.components.RowGroup import com.jetbrains.toolbox.api.ui.components.ValidationErrorField import kotlinx.coroutines.CancellationException import kotlinx.coroutines.CoroutineName -import kotlinx.coroutines.CoroutineScope +import kotlinx.coroutines.CoroutineStart import kotlinx.coroutines.Job +import kotlinx.coroutines.currentCoroutineContext +import kotlinx.coroutines.ensureActive import kotlinx.coroutines.flow.update +import kotlinx.coroutines.job import kotlinx.coroutines.launch +import kotlinx.coroutines.sync.Mutex +import kotlinx.coroutines.sync.withLock import kotlinx.coroutines.yield import java.net.URL @@ -34,6 +42,8 @@ class ConnectStep( private val onTokenRefreshed: (suspend (url: URL, oauthSessionCtx: CoderOAuthSessionContext) -> Unit)? = null ) : WizardStep { private var signInJob: Job? = null + private var disposed = false + private val connectionMutex = Mutex() private val statusField = LabelField(context.i18n.pnotr("")) private val errorField = ValidationErrorField(context.i18n.pnotr("")) @@ -44,124 +54,93 @@ class ConnectStep( ) override fun onVisible() { - errorField.textState.update { - context.i18n.pnotr("") - } + prepareConnection()?.start() + } + @Synchronized + private fun prepareConnection(): Job? { + if (disposed || signInJob != null || model.currentStep() !in setOf(CONNECT, RESTORE_SESSION)) return null + + errorField.textState.update { context.i18n.pnotr("") } if (context.settingsStore.requiresTokenAuth && model.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 + return null } - - // Don't launch another connection attempt if one is already in progress. - if (signInJob?.isActive == true) { - context.logger.info("Connection already in progress, skipping duplicate") - return - } - - statusField.textState.update { context.i18n.pnotr("Connecting to ${model.url?.host ?: "unknown host"}...") } - connect() - } - - /** - * Try connecting to Coder with the provided URL and token. - */ - private fun connect() { val url = model.url if (url == null) { errorField.textState.update { context.i18n.ptrl("URL is required") } - return + return null } - if (context.settingsStore.requiresTokenAuth && !model.hasToken() && !model.hasOAuthSession()) { - errorField.textState.update { context.i18n.ptrl("Token is required") } - return + statusField.textState.update { context.i18n.pnotr("Connecting to ${url.origin() ?: url.host}...") } + // Record the job before dispatching it. Toolbox can show the same step again during startup. + val job = context.cs.launch(CoroutineName("Async Http and CLI Setup"), start = CoroutineStart.LAZY) { + // A retry waits for a cancelled CLI process to finish cleaning up. + connectionMutex.withLock { connect(url) } } + signInJob = job + return job + } - // Capture the host name early for error reporting - val hostName = url.host - - // Cancel previous job regardless of the new mode - signInJob?.cancel() - - // 1. Extract the logic into a reusable suspend lambda - val connectionLogic: suspend CoroutineScope.() -> Unit = { - try { - var oauthSession: CoderOAuthSessionContext? = null - if (context.settingsStore.requiresTokenAuth && context.settingsStore.preferOAuth2IfAvailable && model.hasOAuthSession()) { - refreshOAuthToken() - oauthSession = model.oauthSession!!.copy() - } + private suspend fun connect(url: URL) { + var pendingClient: CoderRestClient? = null + try { + var oauthSession: CoderOAuthSessionContext? = null + if (context.settingsStore.requiresTokenAuth && context.settingsStore.preferOAuth2IfAvailable && model.hasOAuthSession()) { + refreshOAuthToken() + oauthSession = model.oauthSession!!.copy() + } - val apiToken = if (context.settingsStore.requiresTokenAuth) model.token else null - - context.logger.info("Setting up the HTTP client...") - val client = CoderRestClient( - context, - url, - apiToken, - oauthSession, - PluginManager.pluginInfo.version, - onTokenRefreshed - ) - // allows interleaving with the back/cancel action - yield() - client.initializeSession() - logAndReportProgress("Checking Coder CLI...") - val cli = ensureCLI( - context, client.url, - client.buildVersion - ) { progress -> - statusField.textState.update { (context.i18n.pnotr(progress)) } - } - // We only need to log in if we are using token-based auth. - if (context.settingsStore.requiresTokenAuth) { - logAndReportProgress("Configuring Coder CLI...") - // allows interleaving with the back/cancel action - yield() - if (oauthSession != null) { - cli.login(oauthSession.tokenResponse!!.accessToken) - } else { - cli.login(apiToken!!) - } - } - logAndReportProgress("Successfully configured ${hostName}...") + val apiToken = if (context.settingsStore.requiresTokenAuth) model.token else null + + context.logger.info("Setting up the HTTP client...") + val client = CoderRestClient( + context, + url, + apiToken, + oauthSession, + PluginManager.pluginInfo.version, + onTokenRefreshed + ) + pendingClient = client + // allows interleaving with the back/cancel action + yield() + client.initializeSession() + logAndReportProgress("Checking Coder CLI...") + val cli = ensureCLI( + context, client.url, + client.buildVersion + ) { progress -> + statusField.textState.update { (context.i18n.pnotr(progress)) } + } + // We only need to log in if we are using token-based auth. + if (context.settingsStore.requiresTokenAuth) { + logAndReportProgress("Configuring Coder CLI...") // allows interleaving with the back/cancel action yield() - context.logger.info("Connection setup done, initializing the workspace poller...") - onConnect.accept(client, cli) - // Only invoke onTokenRefreshed when we actually have an OAuth session - oauthSession?.let { session -> - onTokenRefreshed?.invoke(client.url, session) + if (oauthSession != null) { + cli.login(oauthSession.tokenResponse!!.accessToken) + } else { + cli.login(apiToken!!) } - // The provider's onConnect ran close() which clears the router; combined - // with client now being non-null this drops the wizard from getOverrideUiPage. - context.envPageManager.showPluginEnvironmentsPage(false) - } catch (ex: CancellationException) { - // Back-button cancellation already navigates in onBack(), while - // dispose() must cancel without navigating. Treat these control-flow - // cancellations separately so we do not run navigateBack() twice. - if (ex.message != USER_HIT_THE_BACK_BUTTON && ex.message != WIZARD_WAS_DISPOSED) { - context.logger.logAndShowError( - "Error encountered while setting up Coder", - "Failed to configure $hostName. ${ex.message}", - ex - ) - navigateBack() - } - } catch (ex: Exception) { - context.logger.logAndShowError( - "Error encountered while setting up Coder", - "Failed to configure $hostName. ${ex.message}", - ex - ) - navigateBack() } + logAndReportProgress("Successfully configured ${url.host}...") + // allows interleaving with the back/cancel action + yield() + context.logger.info("Connection setup done, initializing the workspace poller...") + onConnect.accept(client, cli) + pendingClient = null // The provider now owns this client. + // Only invoke onTokenRefreshed when we actually have an OAuth session + oauthSession?.let { session -> + onTokenRefreshed?.invoke(client.url, session) + } + } catch (ex: Exception) { + reportConnectionFailure(currentCoroutineContext().job, url, ex) + } finally { + pendingClient?.close() } - - signInJob = context.cs.launch(CoroutineName("Async Http and CLI Setup"), block = connectionLogic) } private suspend fun refreshOAuthToken() { @@ -179,15 +158,31 @@ class ConnectStep( statusField.textState.update { context.i18n.pnotr(msg) } } + @Synchronized + private fun reportConnectionFailure(job: Job, url: URL, ex: Exception) { + job.ensureActive() + if (signInJob !== job) return + signInJob = null + context.logger.logAndShowError( + "Error encountered while setting up Coder", + "Failed to configure ${url.origin() ?: url.host}. ${ex.message}", + ex, + ) + navigateBack() + } + override suspend fun onNext(): Boolean { return false } + @Synchronized override fun onBack() { + if (disposed || model.currentStep() != CONNECT) return try { context.logger.info("Back button was pressed, cancelling in-progress connection setup...") signInJob?.cancel(CancellationException(USER_HIT_THE_BACK_BUTTON)) } finally { + signInJob = null navigateBack() } } @@ -197,7 +192,9 @@ class ConnectStep( * itself is being torn down by an external trigger (e.g. a deep link to a * different deployment). */ + @Synchronized fun dispose() { + disposed = true signInJob?.cancel(CancellationException(WIZARD_WAS_DISPOSED)) signInJob = null } diff --git a/src/main/kotlin/com/coder/toolbox/views/DeploymentUrlStep.kt b/src/main/kotlin/com/coder/toolbox/views/DeploymentUrlStep.kt index 72325344..ceac44b3 100644 --- a/src/main/kotlin/com/coder/toolbox/views/DeploymentUrlStep.kt +++ b/src/main/kotlin/com/coder/toolbox/views/DeploymentUrlStep.kt @@ -12,6 +12,8 @@ import com.coder.toolbox.util.toURL import com.coder.toolbox.util.validateStrictWebUrl import com.coder.toolbox.views.state.CoderOAuthSessionContext import com.coder.toolbox.views.state.WizardModel +import com.coder.toolbox.views.state.WizardStep.CONNECT +import com.coder.toolbox.views.state.WizardStep.TOKEN_REQUEST import com.jetbrains.toolbox.api.localization.LocalizableString import com.jetbrains.toolbox.api.ui.components.CheckboxField import com.jetbrains.toolbox.api.ui.components.FieldModifier @@ -88,7 +90,7 @@ class DeploymentUrlStep( } if (context.settingsStore.requiresMTlsAuth) { - model.goToLast() + model.goTo(CONNECT) return true } if (context.settingsStore.requiresTokenAuth && context.settingsStore.preferOAuth2IfAvailable) { @@ -102,7 +104,7 @@ class DeploymentUrlStep( } } // if all else fails try the good old API token auth - model.goToNext() + model.goTo(TOKEN_REQUEST) return true } diff --git a/src/main/kotlin/com/coder/toolbox/views/TokenStep.kt b/src/main/kotlin/com/coder/toolbox/views/TokenStep.kt index 902f14fd..42b290e6 100644 --- a/src/main/kotlin/com/coder/toolbox/views/TokenStep.kt +++ b/src/main/kotlin/com/coder/toolbox/views/TokenStep.kt @@ -3,6 +3,8 @@ package com.coder.toolbox.views import com.coder.toolbox.CoderToolboxContext import com.coder.toolbox.util.withPath import com.coder.toolbox.views.state.WizardModel +import com.coder.toolbox.views.state.WizardStep.CONNECT +import com.coder.toolbox.views.state.WizardStep.URL_REQUEST import com.jetbrains.toolbox.api.localization.LocalizableString import com.jetbrains.toolbox.api.ui.components.FieldModifier import com.jetbrains.toolbox.api.ui.components.LinkField @@ -56,12 +58,12 @@ class TokenStep( } model.token = token - model.goToNext() + model.goTo(CONNECT) return true } override fun onBack() { - model.goToPrevious() + model.goTo(URL_REQUEST) } private fun reportError(message: LocalizableString?) { diff --git a/src/main/kotlin/com/coder/toolbox/views/state/PageRouter.kt b/src/main/kotlin/com/coder/toolbox/views/state/PageRouter.kt index a78b27cb..0ab622dc 100644 --- a/src/main/kotlin/com/coder/toolbox/views/state/PageRouter.kt +++ b/src/main/kotlin/com/coder/toolbox/views/state/PageRouter.kt @@ -1,7 +1,8 @@ package com.coder.toolbox.views.state -import com.coder.toolbox.views.CoderPage import com.coder.toolbox.views.CoderSetupWizardPage +import kotlinx.coroutines.Job +import kotlinx.coroutines.ensureActive import kotlinx.coroutines.flow.MutableStateFlow /** @@ -20,24 +21,22 @@ sealed interface PageRoute { class PageRouter { private val route = MutableStateFlow(PageRoute.None) - private val activeWizard: CoderSetupWizardPage? + val activePage: CoderSetupWizardPage? get() = (route.value as? PageRoute.Wizard)?.page - val activePage: CoderPage? - get() = activeWizard - val hasActiveWizard: Boolean - get() = activeWizard != null + get() = activePage != null val pendingOAuthConnection: PendingOAuthConnection? - get() = activeWizard?.pendingOAuthConnection() + get() = activePage?.pendingOAuthConnection() /** * Returns the page already on this route, or builds a new one and * registers it. */ - fun getOrCreate(build: () -> CoderSetupWizardPage): CoderPage { - (route.value as? PageRoute.Wizard)?.let { return it.page } + @Synchronized + fun getOrCreate(build: () -> CoderSetupWizardPage): CoderSetupWizardPage { + activePage?.let { return it } return build().also { route.value = PageRoute.Wizard(it) } } @@ -47,13 +46,25 @@ class PageRouter { * (e.g. a deep link to a different deployment) needs to forcibly install * a new wizard. */ + @Synchronized fun navigate(page: CoderSetupWizardPage) { - activeWizard?.dispose() + activePage?.dispose() route.value = PageRoute.Wizard(page) } + @Synchronized fun clear() { - activeWizard?.dispose() + activePage?.dispose() + route.value = PageRoute.None + } + + @Synchronized + fun finishSignIn(connection: Job, startSession: () -> Unit) { + // Replacing a wizard cancels its connection. Only the winner can publish a session. + connection.ensureActive() + startSession() + val wizard = activePage ?: return route.value = PageRoute.None + wizard.model.goTo(WizardStep.SIGNED_IN) } } diff --git a/src/main/kotlin/com/coder/toolbox/views/state/WizardModel.kt b/src/main/kotlin/com/coder/toolbox/views/state/WizardModel.kt index 0104a5f6..b8285cc8 100644 --- a/src/main/kotlin/com/coder/toolbox/views/state/WizardModel.kt +++ b/src/main/kotlin/com/coder/toolbox/views/state/WizardModel.kt @@ -25,24 +25,10 @@ class WizardModel { _step.value = step } - fun goToNext() { - val entries = WizardStep.entries - _step.value = entries[(_step.value.ordinal + 1) % entries.size] - } - - fun goToPrevious() { - val entries = WizardStep.entries - _step.value = entries[(_step.value.ordinal - 1 + entries.size) % entries.size] - } - fun goToFirst() { _step.value = WizardStep.URL_REQUEST } - fun goToLast() { - _step.value = WizardStep.CONNECT - } - fun hasUrl(): Boolean = url != null fun hasToken(): Boolean = !token.isNullOrBlank() fun hasOAuthSession(): Boolean = oauthSession?.tokenResponse?.accessToken != null @@ -57,7 +43,7 @@ class WizardModel { } enum class WizardStep { - URL_REQUEST, TOKEN_REQUEST, CONNECT; + URL_REQUEST, TOKEN_REQUEST, CONNECT, RESTORE_SESSION, SIGNED_IN; } sealed interface Credentials { diff --git a/src/test/kotlin/com/coder/toolbox/CoderRemoteProviderTest.kt b/src/test/kotlin/com/coder/toolbox/CoderRemoteProviderTest.kt index bbac7437..0a061b3b 100644 --- a/src/test/kotlin/com/coder/toolbox/CoderRemoteProviderTest.kt +++ b/src/test/kotlin/com/coder/toolbox/CoderRemoteProviderTest.kt @@ -642,6 +642,11 @@ class CoderRemoteProviderTest { } } + private fun setupWizard(provider: CoderRemoteProvider): CoderSetupWizardPage { + val router = provider.javaClass.getDeclaredField("router").apply { isAccessible = true }.get(provider) as PageRouter + return requireNotNull(router.activePage) + } + @Test fun `links to the active origin continue without confirmation and preserve its URL`() = runTest { val secrets = mockContext.secrets @@ -839,6 +844,7 @@ class CoderRemoteProviderTest { @Test fun `declining an account change preserves an in-flight auto-connect wizard`() = runTest { + every { mockContext.cs } returns backgroundScope val url = URI("https://coder.example.com").toURL() every { mockContext.deploymentUrl } returns url every { mockContext.settingsStore.requiresMTlsAuth } returns false @@ -850,12 +856,14 @@ class CoderRemoteProviderTest { val provider = spyk(remoteProvider, recordPrivateCalls = true) coEvery { provider["userForToken"](url, "link-token") } returns DataGen.user().copy(username = "other-user") coEvery { provider["userForToken"](url, "saved-token") } returns DataGen.user().copy(username = "saved-user") - val existingWizard = provider.getOverrideUiPage() as CoderSetupWizardPage + assertNull(provider.getOverrideUiPage()) + val existingWizard = setupWizard(provider) provider.handleUri(workspaceLink("https://coder.example.com")) - assertSame(existingWizard, provider.getOverrideUiPage()) - assertEquals(WizardStep.CONNECT, existingWizard.model.currentStep()) + assertNull(provider.getOverrideUiPage()) + assertSame(existingWizard, setupWizard(provider)) + assertEquals(WizardStep.RESTORE_SESSION, existingWizard.model.currentStep()) assertEquals("saved-token", existingWizard.model.token) verify(exactly = 0) { mockContext.popupPluginMainPage() } } @@ -1029,16 +1037,45 @@ class CoderRemoteProviderTest { every { mockContext.settingsStore } returns settings every { mockContext.cs } returns backgroundScope every { mockClient.token } returns "saved-token" - val callback = remoteProvider.javaClass.getDeclaredField("onConnect").apply { isAccessible = true } - .get(remoteProvider) as SuspendBiConsumer - callback.accept(mockClient, mockCli) + completeSignIn() verify { secrets.storeApiTokenFor(url, "saved-token") } verify { secrets.trustDeployment(url) } assertEquals("https://coder.example.com", settings.lastDeploymentURL) } + @Test + @OptIn(ExperimentalCoroutinesApi::class) + fun `sign-in behind Settings opens the workspace list when the user returns to setup`() = runTest { + every { mockContext.cs } returns backgroundScope + coEvery { mockClient.workspaces(any()) } returns emptyList() + val wizard = remoteProvider.getOverrideUiPage() as CoderSetupWizardPage + val pages = mockContext.envPageManager + wizard.beforeShow() + runCurrent() + wizard.afterHide() + + completeSignIn() + runCurrent() + verify(exactly = 0) { pages.showPluginEnvironmentsPage(false) } + assertNull(remoteProvider.getOverrideUiPage()) + + wizard.beforeShow() + wizard.beforeShow() + runCurrent() + verify(exactly = 1) { pages.showPluginEnvironmentsPage(false) } + wizard.afterHide() + remoteProvider.close() + } + + @Suppress("UNCHECKED_CAST") + private suspend fun completeSignIn() { + val callback = remoteProvider.javaClass.getDeclaredField("onConnect").apply { isAccessible = true } + .get(remoteProvider) as SuspendBiConsumer + callback.accept(mockClient, mockCli) + } + @Test fun `declining a fresh link preserves the existing setup wizard`() = runTest { every { mockContext.settingsStore.requiresMTlsAuth } returns false @@ -1138,7 +1175,8 @@ class CoderRemoteProviderTest { } @Test - fun `given mTLS is required when auto setup has stored credentials then mTLS takes precedence`() { + fun `given mTLS is required when auto setup has stored credentials then mTLS takes precedence`() = runTest { + every { mockContext.cs } returns backgroundScope // given val url = URI("https://coder.example.com").toURL() every { mockContext.deploymentUrl } returns url @@ -1150,16 +1188,18 @@ class CoderRemoteProviderTest { val provider = CoderRemoteProvider(mockContext) // when - val overridePage = provider.getOverrideUiPage() as CoderSetupWizardPage + assertNull(provider.getOverrideUiPage()) + val overridePage = setupWizard(provider) // then - assertEquals(WizardStep.CONNECT, overridePage.model.currentStep()) + assertEquals(WizardStep.RESTORE_SESSION, overridePage.model.currentStep()) assertNull(overridePage.model.token) assertNull(overridePage.model.oauthSession) } @Test - fun `given OAuth is preferred when auto setup has token and OAuth session then OAuth is used`() { + fun `given OAuth is preferred when auto setup has token and OAuth session then OAuth is used`() = runTest { + every { mockContext.cs } returns backgroundScope // given val url = URI("https://coder.example.com").toURL() every { mockContext.deploymentUrl } returns url @@ -1171,16 +1211,18 @@ class CoderRemoteProviderTest { val provider = CoderRemoteProvider(mockContext) // when - val overridePage = provider.getOverrideUiPage() as CoderSetupWizardPage + assertNull(provider.getOverrideUiPage()) + val overridePage = setupWizard(provider) // then - assertEquals(WizardStep.CONNECT, overridePage.model.currentStep()) + assertEquals(WizardStep.RESTORE_SESSION, overridePage.model.currentStep()) assertNull(overridePage.model.token) assertNotNull(overridePage.model.oauthSession) } @Test - fun `given OAuth is not preferred when auto setup has token and OAuth session then token is used`() { + fun `given OAuth is not preferred when auto setup has token and OAuth session then token is used`() = runTest { + every { mockContext.cs } returns backgroundScope // given val url = URI("https://coder.example.com").toURL() every { mockContext.deploymentUrl } returns url @@ -1192,16 +1234,18 @@ class CoderRemoteProviderTest { val provider = CoderRemoteProvider(mockContext) // when - val overridePage = provider.getOverrideUiPage() as CoderSetupWizardPage + assertNull(provider.getOverrideUiPage()) + val overridePage = setupWizard(provider) // then - assertEquals(WizardStep.CONNECT, overridePage.model.currentStep()) + assertEquals(WizardStep.RESTORE_SESSION, overridePage.model.currentStep()) assertEquals("token", overridePage.model.token) assertNull(overridePage.model.oauthSession) } @Test - fun `given OAuth is not preferred when auto setup has API token then token is used`() { + fun `given OAuth is not preferred when auto setup has API token then token is used`() = runTest { + every { mockContext.cs } returns backgroundScope // given val url = URI("https://coder.example.com").toURL() every { mockContext.deploymentUrl } returns url @@ -1212,10 +1256,11 @@ class CoderRemoteProviderTest { val provider = CoderRemoteProvider(mockContext) // when - val overridePage = provider.getOverrideUiPage() as CoderSetupWizardPage + assertNull(provider.getOverrideUiPage()) + val overridePage = setupWizard(provider) // then - assertEquals(WizardStep.CONNECT, overridePage.model.currentStep()) + assertEquals(WizardStep.RESTORE_SESSION, overridePage.model.currentStep()) assertEquals("api-token", overridePage.model.token) assertNull(overridePage.model.oauthSession) } diff --git a/src/test/kotlin/com/coder/toolbox/views/CoderSetupWizardPageTest.kt b/src/test/kotlin/com/coder/toolbox/views/CoderSetupWizardPageTest.kt new file mode 100644 index 00000000..668378ce --- /dev/null +++ b/src/test/kotlin/com/coder/toolbox/views/CoderSetupWizardPageTest.kt @@ -0,0 +1,234 @@ +package com.coder.toolbox.views + +import com.coder.toolbox.CoderToolboxContext +import com.coder.toolbox.cli.CoderCLIManager +import com.coder.toolbox.cli.ensureCLI +import com.coder.toolbox.sdk.CoderRestClient +import com.coder.toolbox.sdk.DataGen +import com.coder.toolbox.settings.Environment +import com.coder.toolbox.store.CoderSettingsStore +import com.coder.toolbox.util.pluginTestSettingsStore +import com.coder.toolbox.views.state.WizardModel +import com.coder.toolbox.views.state.Credentials +import com.coder.toolbox.views.state.WizardStep +import io.mockk.coEvery +import io.mockk.coVerify +import io.mockk.every +import io.mockk.justRun +import io.mockk.mockk +import io.mockk.mockkConstructor +import io.mockk.mockkStatic +import io.mockk.unmockkConstructor +import io.mockk.unmockkStatic +import io.mockk.verify +import kotlinx.coroutines.CancellationException +import kotlinx.coroutines.CompletableDeferred +import kotlinx.coroutines.CoroutineDispatcher +import kotlinx.coroutines.CoroutineScope +import kotlinx.coroutines.ExperimentalCoroutinesApi +import kotlinx.coroutines.NonCancellable +import kotlinx.coroutines.awaitCancellation +import kotlinx.coroutines.test.StandardTestDispatcher +import kotlinx.coroutines.test.runCurrent +import kotlinx.coroutines.test.runTest +import kotlinx.coroutines.withContext +import java.net.Proxy +import java.net.URI +import kotlin.coroutines.CoroutineContext +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertNull + +@OptIn(ExperimentalCoroutinesApi::class) +class CoderSetupWizardPageTest { + private val context = mockk(relaxed = true) + private val url = URI("https://coder.example.test").toURL() + private val model = WizardModel().apply { + url = this@CoderSetupWizardPageTest.url + token = "test-token" + goTo(WizardStep.CONNECT) + } + private val onConnect = mockk>(relaxed = true) + + @Test + fun `saved credentials restore once behind the workspace list without page navigation`() = runTest { + every { context.cs } returns backgroundScope + val page = CoderSetupWizardPage.connectStep( + context, mockk(relaxed = true), url, Credentials.Token("test-token"), onConnect, restoreSavedSession = true, + ) + val pages = context.envPageManager + val ui = context.ui + + withMockCli { cli -> + assertNull(page.pageForDisplay()) + assertNull(page.pageForDisplay()) + runCurrent() + + coVerify(exactly = 1) { cli.login("test-token") } + coVerify(exactly = 1) { onConnect.accept(any(), cli) } + verify(exactly = 0) { pages.showPluginEnvironmentsPage(any()) } + coVerify(exactly = 0) { ui.showUiPageSuspending(any()) } + } + } + + @Test + fun `a failed restore presents the sign-in form and permits a new connection`() = runTest { + every { context.cs } returns backgroundScope + val page = CoderSetupWizardPage.connectStep( + context, mockk(relaxed = true), url, Credentials.Token("test-token"), onConnect, restoreSavedSession = true, + ) + val ui = context.ui + + withMockCli { cli -> + coEvery { cli.login("test-token") } throws IllegalStateException("Keychain is locked") andThen "Signed in" + assertNull(page.pageForDisplay()) + runCurrent() + + assertEquals(WizardStep.URL_REQUEST, page.model.currentStep()) + coVerify(exactly = 1) { ui.showUiPageSuspending(page) } + coVerify(exactly = 0) { onConnect.accept(any(), any()) } + + page.model.url = url + page.model.token = "test-token" + page.model.goTo(WizardStep.CONNECT) + page.beforeShow() + runCurrent() + coVerify(exactly = 2) { cli.login("test-token") } + coVerify(exactly = 1) { onConnect.accept(any(), cli) } + page.dispose() + } + } + + @Test + fun `submitting a token displays the connection step once`() = runTest { + every { context.cs } returns backgroundScope + every { context.secrets.apiTokenFor(url) } returns "test-token" + mockkConstructor(ConnectStep::class) + try { + justRun { anyConstructed().onVisible() } + val page = CoderSetupWizardPage.deploymentUrlStep(context, mockk(), mockk()) + page.model.url = url + page.model.goTo(WizardStep.TOKEN_REQUEST) + page.beforeShow() + runCurrent() + + page.actionButtons.value.first().run() + runCurrent() + + assertEquals(WizardStep.CONNECT, page.model.currentStep()) + verify(exactly = 1) { anyConstructed().onVisible() } + page.dispose() + page.beforeShow() + runCurrent() + verify(exactly = 1) { anyConstructed().onVisible() } + } finally { + unmockkConstructor(ConnectStep::class) + } + } + + @Test + fun `a second visibility callback during startup logs in to the CLI once`() = runTest { + lateinit var step: ConnectStep + val scheduler = StandardTestDispatcher(testScheduler) + val dispatcher = object : CoroutineDispatcher() { + private var firstDispatch = true + + override fun dispatch(context: CoroutineContext, block: Runnable) { + if (firstDispatch) { + firstDispatch = false + step.onVisible() + } + scheduler.dispatch(context, block) + } + } + every { context.cs } returns CoroutineScope(backgroundScope.coroutineContext + dispatcher) + step = ConnectStep(context, model, { model.goTo(WizardStep.TOKEN_REQUEST) }, onConnect) + withMockCli { cli -> + step.onVisible() + runCurrent() + + coVerify(exactly = 1) { cli.login("test-token") } + coVerify(exactly = 1) { onConnect.accept(any(), cli) } + } + } + + private fun withMockCli(test: (CoderCLIManager) -> Unit) { + every { context.settingsStore } returns CoderSettingsStore( + pluginTestSettingsStore(), Environment(), mockk(relaxed = true), + ) + every { context.proxySettings.getProxy() } returns Proxy.NO_PROXY + every { context.proxySettings.getProxyAuth() } returns null + val cli = mockk(relaxed = true) + mockkConstructor(CoderRestClient::class) + mockkStatic(::ensureCLI) + try { + coEvery { anyConstructed().initializeSession() } returns DataGen.user() + every { anyConstructed().buildVersion } returns "2.39.0" + coEvery { ensureCLI(context, url, "2.39.0", any()) } returns cli + test(cli) + } finally { + unmockkStatic(::ensureCLI) + unmockkConstructor(CoderRestClient::class) + } + } + + @Test + fun `a retry waits for a cancelled login and ignores its late error`() = runTest { + every { context.cs } returns backgroundScope + val step = ConnectStep(context, model, { model.goTo(WizardStep.TOKEN_REQUEST) }, onConnect) + val allowCleanup = CompletableDeferred() + withMockCli { cli -> + coEvery { cli.login("test-token") } coAnswers { + try { + awaitCancellation() + } catch (ex: CancellationException) { + withContext(NonCancellable) { allowCleanup.await() } + throw IllegalStateException("Cancelled CLI process exited", ex) + } + } andThen "Signed in" + try { + step.onVisible() + runCurrent() + step.onBack() + assertEquals(WizardStep.TOKEN_REQUEST, model.currentStep()) + model.goTo(WizardStep.CONNECT) + step.onVisible() + runCurrent() + coVerify(exactly = 1) { cli.login("test-token") } + + allowCleanup.complete(Unit) + runCurrent() + + assertEquals(WizardStep.CONNECT, model.currentStep()) + coVerify(exactly = 2) { cli.login("test-token") } + coVerify(exactly = 1) { onConnect.accept(any(), cli) } + } finally { + allowCleanup.complete(Unit) + step.dispose() + } + } + } + + @Test + fun `a genuine CLI failure returns to token entry and allows a new sign-in`() = runTest { + every { context.cs } returns backgroundScope + val step = ConnectStep(context, model, { model.goTo(WizardStep.TOKEN_REQUEST) }, onConnect) + withMockCli { cli -> + coEvery { cli.login("test-token") } throws IllegalStateException("Keychain is locked") andThen "Signed in" + step.onVisible() + runCurrent() + assertEquals(WizardStep.TOKEN_REQUEST, model.currentStep()) + coVerify(exactly = 0) { onConnect.accept(any(), any()) } + + model.goTo(WizardStep.CONNECT) + step.onVisible() + runCurrent() + step.dispose() + step.onVisible() + runCurrent() + + coVerify(exactly = 2) { cli.login("test-token") } + coVerify(exactly = 1) { onConnect.accept(any(), cli) } + } + } +} diff --git a/src/test/kotlin/com/coder/toolbox/views/state/PageRouterTest.kt b/src/test/kotlin/com/coder/toolbox/views/state/PageRouterTest.kt index c0fd8dbc..ec93e121 100644 --- a/src/test/kotlin/com/coder/toolbox/views/state/PageRouterTest.kt +++ b/src/test/kotlin/com/coder/toolbox/views/state/PageRouterTest.kt @@ -4,10 +4,19 @@ import com.coder.toolbox.oauth.TokenEndpointAuthMethod import com.coder.toolbox.views.CoderSetupWizardPage import io.mockk.every import io.mockk.mockk +import io.mockk.verify +import kotlinx.coroutines.CancellationException +import kotlinx.coroutines.Job import java.net.URI +import java.util.concurrent.CountDownLatch +import java.util.concurrent.Executors +import java.util.concurrent.TimeUnit import kotlin.test.Test import kotlin.test.assertEquals +import kotlin.test.assertFailsWith import kotlin.test.assertNull +import kotlin.test.assertSame +import kotlin.test.assertTrue class PageRouterTest { @@ -42,6 +51,51 @@ class PageRouterTest { assertEquals(pendingConnection, pendingOAuthConnection) } + @Test + fun `a workspace link replaces a wizard being built for auto-connect`() { + val router = PageRouter() + val automatic = mockk(relaxed = true) + val fromLink = mockk() + val building = CountDownLatch(1) + val switching = CountDownLatch(1) + val workers = Executors.newFixedThreadPool(2) + try { + val autoConnect = workers.submit { + router.getOrCreate { + building.countDown() + check(switching.await(5, TimeUnit.SECONDS)) + automatic + } + } + assertTrue(building.await(5, TimeUnit.SECONDS)) + val openLink = workers.submit { + switching.countDown() + router.navigate(fromLink) + } + autoConnect.get(5, TimeUnit.SECONDS) + openLink.get(5, TimeUnit.SECONDS) + + assertSame(fromLink, router.activePage) + verify(exactly = 1) { automatic.dispose() } + } finally { + workers.shutdownNow() + } + } + + @Test + fun `a cancelled sign-in cannot publish a session or dismiss its replacement`() { + val router = PageRouter() + val replacement = mockk() + router.navigate(replacement) + val cancelledConnection = Job().apply { cancel() } + + assertFailsWith { + router.finishSignIn(cancelledConnection) { error("The cancelled session was started") } + } + + assertSame(replacement, router.activePage) + } + private fun oauthSession(): CoderOAuthSessionContext = CoderOAuthSessionContext( clientId = "client-id", clientSecret = "client-secret", From c43e049424088818fdeb2dd06d75b484a54fada8 Mon Sep 17 00:00:00 2001 From: Faur Ioan-Aurel Date: Thu, 12 Nov 2026 00:35:44 +0200 Subject: [PATCH 2/7] Show the replacement setup wizard when opening a workspace link A workspace link for another deployment could leave the previous setup page visible until Toolbox regained focus. The provider now displays the replacement wizard as soon as it handles the link, so setup can continue while Toolbox stays in focus. --- CHANGELOG.md | 1 + .../com/coder/toolbox/CoderRemoteProvider.kt | 8 +- .../coder/toolbox/CoderRemoteProviderTest.kt | 174 ++++++++++++++---- .../toolbox/util/InMemoryPluginSecretStore.kt | 17 ++ 4 files changed, 166 insertions(+), 34 deletions(-) create mode 100644 src/test/kotlin/com/coder/toolbox/util/InMemoryPluginSecretStore.kt diff --git a/CHANGELOG.md b/CHANGELOG.md index 58aeb6d7..56c81b2b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -20,6 +20,7 @@ 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 diff --git a/src/main/kotlin/com/coder/toolbox/CoderRemoteProvider.kt b/src/main/kotlin/com/coder/toolbox/CoderRemoteProvider.kt index 8680e23c..4f7ab6fc 100644 --- a/src/main/kotlin/com/coder/toolbox/CoderRemoteProvider.kt +++ b/src/main/kotlin/com/coder/toolbox/CoderRemoteProvider.kt @@ -817,7 +817,7 @@ class CoderRemoteProvider( linkHandler.handle(params, url, client, cli) } - private fun connectFromLink(params: Map, url: URL, token: String?) { + private suspend fun connectFromLink(params: Map, url: URL, token: String?) { val credentials = token?.let { Credentials.Token(it) } ?: Credentials.MTls val wizard = CoderSetupWizardPage.connectStep( context, settingsPage, @@ -826,10 +826,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. */ diff --git a/src/test/kotlin/com/coder/toolbox/CoderRemoteProviderTest.kt b/src/test/kotlin/com/coder/toolbox/CoderRemoteProviderTest.kt index 0a061b3b..6b8d5b9f 100644 --- a/src/test/kotlin/com/coder/toolbox/CoderRemoteProviderTest.kt +++ b/src/test/kotlin/com/coder/toolbox/CoderRemoteProviderTest.kt @@ -1,6 +1,7 @@ package com.coder.toolbox import com.coder.toolbox.cli.CoderCLIManager +import com.coder.toolbox.cli.ensureCLI import com.coder.toolbox.diagnostics.CoderLogger import com.coder.toolbox.oauth.TokenEndpointAuthMethod import com.coder.toolbox.sdk.CoderRestClient @@ -16,8 +17,10 @@ import com.coder.toolbox.sdk.v2.models.WorkspaceResource import com.coder.toolbox.sdk.v2.models.WorkspaceStatus import com.coder.toolbox.session.SessionId import com.coder.toolbox.settings.Environment +import com.coder.toolbox.store.CoderSecretsStore import com.coder.toolbox.store.CoderSettingsStore import com.coder.toolbox.util.CoderProtocolHandler +import com.coder.toolbox.util.InMemoryPluginSecretStore import com.coder.toolbox.util.pluginTestSettingsStore import com.coder.toolbox.views.CoderSetupWizardPage import com.coder.toolbox.views.SuspendBiConsumer @@ -35,16 +38,20 @@ import io.mockk.coEvery import io.mockk.coVerify import io.mockk.every import io.mockk.mockk +import io.mockk.mockkConstructor +import io.mockk.mockkStatic import io.mockk.slot import io.mockk.spyk +import io.mockk.unmockkConstructor +import io.mockk.unmockkStatic import io.mockk.verify -import io.mockk.verifyOrder import kotlinx.coroutines.CancellationException import kotlinx.coroutines.CompletableDeferred import kotlinx.coroutines.CoroutineScope import kotlinx.coroutines.Dispatchers import kotlinx.coroutines.ExperimentalCoroutinesApi import kotlinx.coroutines.async +import kotlinx.coroutines.awaitCancellation import kotlinx.coroutines.test.StandardTestDispatcher import kotlinx.coroutines.test.runCurrent import kotlinx.coroutines.test.runTest @@ -138,6 +145,48 @@ class CoderRemoteProviderTest { assertFalse(remoteProvider.getAccountDropDown().visibility.value) } + @CsvSource( + "https://coder.example.com, false, true", + "https://other.example.com, false, true", + "https://other.example.com, true, false", + ) + @ParameterizedTest(name = "After logout, link to {0}, B remembered={1}, deployment prompt={2}") + fun `logout forgets A while each subsequent link uses its own deployment history`( + target: String, + rememberB: Boolean, + expectPrompt: Boolean, + ) = runTest { + val a = mockClient.url + val b = URI("https://other.example.com").toURL() + val secrets = CoderSecretsStore(InMemoryPluginSecretStore()) + every { mockContext.secrets } returns secrets + every { mockContext.settingsStore.requiresTokenAuth } returns true + every { mockContext.settingsStore.requiresMTlsAuth } returns false + every { mockContext.settingsStore.preferOAuth2IfAvailable } returns false + secrets.storeApiTokenFor(a, "saved-a") + secrets.trustDeployment(a) + if (rememberB) { + secrets.storeApiTokenFor(b, "saved-b") + secrets.trustDeployment(b) + } + val provider = spyk(remoteProvider, recordPrivateCalls = true) + setPrivateField(provider, "client", mockClient) + setPrivateField(provider, "cli", mockCli) + val user = DataGen.user() + coEvery { provider["userForToken"](b, "saved-b") } returns user + coEvery { provider["userForToken"](b, "link-token") } returns user + coEvery { mockUi.showYesNoPopup(any(), any(), any(), any()) } returns false + + provider.logout() + provider.handleUri(workspaceLink(target)) + + assertNull(secrets.apiTokenFor(a)) + assertFalse(secrets.isUrlTrusted(a)) + assertEquals(rememberB, secrets.isUrlTrusted(b)) + coVerify(exactly = if (expectPrompt) 1 else 0) { mockUi.showYesNoPopup(any(), any(), any(), any()) } + coVerify(exactly = if (expectPrompt) 0 else 1) { mockUi.showUiPageSuspending(any()) } + } + @Test fun `provider shutdown preserves credentials`() { setPrivateField(remoteProvider, "client", mockClient) @@ -611,42 +660,105 @@ class CoderRemoteProviderTest { ] ) @ParameterizedTest - fun `accepted deployment switch replaces the wizard before Toolbox requests its page`(target: String) = runTest { - every { mockContext.settingsStore.requiresMTlsAuth } returns false - every { mockContext.settingsStore.requiresTokenAuth } returns true - coEvery { mockUi.showYesNoPopup(any(), any(), any(), any()) } returns true - setPrivateField(remoteProvider, "client", mockClient) - setPrivateField(remoteProvider, "cli", mockCli) - assertNull(remoteProvider.getOverrideUiPage()) - - val previousWizard = mockk(relaxed = true) - val router = PageRouter() - router.navigate(previousWizard) - setPrivateField(remoteProvider, "router", router) - var displayedWizard: CoderSetupWizardPage? = null - every { mockContext.popupPluginMainPage() } answers { - displayedWizard = remoteProvider.getOverrideUiPage() as? CoderSetupWizardPage - } - - remoteProvider.handleUri(workspaceLink(target)) - - val wizard = assertNotNull(displayedWizard) - assertEquals(WizardStep.CONNECT, wizard.model.currentStep()) - assertEquals("https://coder2.example.com", wizard.model.url.toString()) - assertEquals("link-token", wizard.model.token) - assertSame(wizard, remoteProvider.getOverrideUiPage()) - coVerify(exactly = 1) { mockUi.showYesNoPopup(any(), any(), any(), any()) } - verifyOrder { - previousWizard.dispose() - mockContext.popupPluginMainPage() + fun `accepted cold start link presents the new deployment while the old session restores`(target: String) = + runTest { + every { mockContext.cs } returns backgroundScope + val previousUrl = URI("https://coder.example.com").toURL() + every { mockContext.deploymentUrl } returns previousUrl + every { mockContext.secrets.apiTokenFor(previousUrl) } returns "saved-token" + every { mockContext.settingsStore.requiresMTlsAuth } returns false + every { mockContext.settingsStore.requiresTokenAuth } returns true + coEvery { mockUi.showYesNoPopup(any(), any(), any(), any()) } returns true + assertNull(remoteProvider.getOverrideUiPage()) + val previousWizard = setupWizard(remoteProvider) + assertEquals(previousUrl, previousWizard.model.url) + assertEquals(WizardStep.RESTORE_SESSION, previousWizard.model.currentStep()) + val displayedWizard = slot() + coEvery { mockUi.showUiPageSuspending(capture(displayedWizard)) } returns Unit + + remoteProvider.handleUri(workspaceLink(target)) + + // Selecting an already visible provider can leave Toolbox on the old page. + // No second visibility callback or request for the override page is simulated. + coVerify(exactly = 1) { mockUi.showUiPageSuspending(any()) } + val wizard = displayedWizard.captured + assertEquals(WizardStep.CONNECT, wizard.model.currentStep()) + assertEquals("https://coder2.example.com", wizard.model.url.toString()) + assertEquals("link-token", wizard.model.token) + assertSame(wizard, remoteProvider.getOverrideUiPage()) + coVerify(exactly = 1) { mockUi.showYesNoPopup(any(), any(), any(), any()) } } - } private fun setupWizard(provider: CoderRemoteProvider): CoderSetupWizardPage { - val router = provider.javaClass.getDeclaredField("router").apply { isAccessible = true }.get(provider) as PageRouter + val router = + provider.javaClass.getDeclaredField("router").apply { isAccessible = true }.get(provider) as PageRouter return requireNotNull(router.activePage) } + @Test + @OptIn(ExperimentalCoroutinesApi::class) + fun `a cold link signs in to B and hands off its workspace without a focus change`() = runTest { + val a = URI("https://a.example.com").toURL() + val b = URI("https://b.example.com").toURL() + val secrets = mockContext.secrets + val environmentsPage = mockContext.envPageManager + every { mockContext.settingsStore } returns CoderSettingsStore( + pluginTestSettingsStore(), Environment(), underlyingLogger, + ) + every { mockContext.cs } returns backgroundScope + every { mockContext.proxySettings.getProxy() } returns Proxy.NO_PROXY + every { mockContext.proxySettings.getProxyAuth() } returns null + every { mockContext.deploymentUrl } returns a + every { secrets.apiTokenFor(a) } returns "saved-a" + val handler = mockk(relaxed = true) + setPrivateField(remoteProvider, "linkHandler", handler) + val user = DataGen.user() + val aStarted = CompletableDeferred() + val aCancelled = CompletableDeferred() + val shown = slot() + mockkConstructor(CoderRestClient::class) + mockkStatic(::ensureCLI) + try { + coEvery { anyConstructed().initializeSession() } coAnswers { + aStarted.complete(Unit) + try { + awaitCancellation() + } finally { + aCancelled.complete(Unit) + } + } andThen user + every { anyConstructed().me } returns user + every { anyConstructed().buildVersion } returns "2.39.0" + coEvery { anyConstructed().workspaces(any()) } returns emptyList() + coEvery { ensureCLI(mockContext, b, "2.39.0", any()) } returns mockCli + coEvery { mockUi.showYesNoPopup(any(), any(), any(), any()) } returns true + coEvery { mockUi.showUiPageSuspending(capture(shown)) } coAnswers { + shown.captured.beforeShow() + } + every { environmentsPage.showPluginEnvironmentsPage(false) } answers { + assertNull(remoteProvider.getOverrideUiPage()) + shown.captured.afterHide() + } + assertNull(remoteProvider.getOverrideUiPage()) + runCurrent() + assertTrue(aStarted.isCompleted) + remoteProvider.handleUri(workspaceLink(b.toString())) + runCurrent() + + assertTrue(aCancelled.isCompleted) + verify(exactly = 1) { anyConstructed().close() } + coVerify(exactly = 1) { mockCli.login("link-token") } + verify(exactly = 1) { secrets.storeApiTokenFor(b, "link-token") } + coVerify(exactly = 1) { handler.handle(any(), b, any(), mockCli) } + verify(exactly = 1) { environmentsPage.showPluginEnvironmentsPage(false) } + assertNull(remoteProvider.getOverrideUiPage()) + } finally { + remoteProvider.close() + unmockkStatic(::ensureCLI) + unmockkConstructor(CoderRestClient::class) + } + } + @Test fun `links to the active origin continue without confirmation and preserve its URL`() = runTest { val secrets = mockContext.secrets diff --git a/src/test/kotlin/com/coder/toolbox/util/InMemoryPluginSecretStore.kt b/src/test/kotlin/com/coder/toolbox/util/InMemoryPluginSecretStore.kt new file mode 100644 index 00000000..f90259ea --- /dev/null +++ b/src/test/kotlin/com/coder/toolbox/util/InMemoryPluginSecretStore.kt @@ -0,0 +1,17 @@ +package com.coder.toolbox.util + +import com.jetbrains.toolbox.api.core.PluginSecretStore + +internal class InMemoryPluginSecretStore : PluginSecretStore { + private val values = mutableMapOf() + + override fun get(key: String): String? = values[key] + + override fun set(key: String, value: String) { + values[key] = value + } + + override fun clear(key: String) { + values.remove(key) + } +} From 2581ce9bad053cab5a331123a7b4c6d1f9cd431d Mon Sep 17 00:00:00 2001 From: Faur Ioan-Aurel Date: Fri, 13 Nov 2026 18:14:15 +0200 Subject: [PATCH 3/7] Keep the workspace list open after an OAuth save error An OAuth credential update can fail after the provider accepts the session. Treating that error as a failed sign-in sends the user back to setup even though the workspace connection is active. The wizard now logs the error and keeps the signed-in session open. A regression test covers a credential store that rejects the update. --- .../com/coder/toolbox/views/ConnectStep.kt | 5 ++ .../toolbox/views/CoderSetupWizardPageTest.kt | 47 ++++++++++++++++++- 2 files changed, 51 insertions(+), 1 deletion(-) diff --git a/src/main/kotlin/com/coder/toolbox/views/ConnectStep.kt b/src/main/kotlin/com/coder/toolbox/views/ConnectStep.kt index cbce7def..5963d081 100644 --- a/src/main/kotlin/com/coder/toolbox/views/ConnectStep.kt +++ b/src/main/kotlin/com/coder/toolbox/views/ConnectStep.kt @@ -11,6 +11,7 @@ import com.coder.toolbox.views.state.CoderOAuthSessionContext import com.coder.toolbox.views.state.WizardModel import com.coder.toolbox.views.state.WizardStep.CONNECT import com.coder.toolbox.views.state.WizardStep.RESTORE_SESSION +import com.coder.toolbox.views.state.WizardStep.SIGNED_IN import com.jetbrains.toolbox.api.ui.components.LabelField import com.jetbrains.toolbox.api.ui.components.RowGroup import com.jetbrains.toolbox.api.ui.components.ValidationErrorField @@ -162,6 +163,10 @@ class ConnectStep( private fun reportConnectionFailure(job: Job, url: URL, ex: Exception) { job.ensureActive() if (signInJob !== job) return + if (model.currentStep() == SIGNED_IN) { + context.logger.warn(ex, "Failed to update session credentials after sign-in") + return + } signInJob = null context.logger.logAndShowError( "Error encountered while setting up Coder", diff --git a/src/test/kotlin/com/coder/toolbox/views/CoderSetupWizardPageTest.kt b/src/test/kotlin/com/coder/toolbox/views/CoderSetupWizardPageTest.kt index 668378ce..4fb02c7a 100644 --- a/src/test/kotlin/com/coder/toolbox/views/CoderSetupWizardPageTest.kt +++ b/src/test/kotlin/com/coder/toolbox/views/CoderSetupWizardPageTest.kt @@ -3,13 +3,17 @@ package com.coder.toolbox.views import com.coder.toolbox.CoderToolboxContext import com.coder.toolbox.cli.CoderCLIManager import com.coder.toolbox.cli.ensureCLI +import com.coder.toolbox.oauth.OAuthTokenResponse +import com.coder.toolbox.oauth.TokenEndpointAuthMethod import com.coder.toolbox.sdk.CoderRestClient import com.coder.toolbox.sdk.DataGen import com.coder.toolbox.settings.Environment import com.coder.toolbox.store.CoderSettingsStore import com.coder.toolbox.util.pluginTestSettingsStore -import com.coder.toolbox.views.state.WizardModel +import com.coder.toolbox.views.state.CoderOAuthSessionContext import com.coder.toolbox.views.state.Credentials +import com.coder.toolbox.views.state.PageRouter +import com.coder.toolbox.views.state.WizardModel import com.coder.toolbox.views.state.WizardStep import io.mockk.coEvery import io.mockk.coVerify @@ -28,6 +32,8 @@ import kotlinx.coroutines.CoroutineScope import kotlinx.coroutines.ExperimentalCoroutinesApi import kotlinx.coroutines.NonCancellable import kotlinx.coroutines.awaitCancellation +import kotlinx.coroutines.currentCoroutineContext +import kotlinx.coroutines.job import kotlinx.coroutines.test.StandardTestDispatcher import kotlinx.coroutines.test.runCurrent import kotlinx.coroutines.test.runTest @@ -99,6 +105,45 @@ class CoderSetupWizardPageTest { } } + @Test + fun `an OAuth credential update failure keeps the signed-in wizard on the workspace list`() = runTest { + every { context.cs } returns backgroundScope + val session = CoderOAuthSessionContext( + clientId = "client-id", + clientSecret = "client-secret", + tokenCodeVerifier = "verifier", + state = "state", + tokenEndpoint = "$url/oauth/token", + tokenResponse = OAuthTokenResponse("access-token", "Bearer", null, "refresh-token", null), + tokenAuthMethod = TokenEndpointAuthMethod.CLIENT_SECRET_BASIC, + ) + val router = PageRouter() + val failure = IllegalStateException("Keychain is locked") + val page = router.getOrCreate { + CoderSetupWizardPage.connectStep( + context, mockk(), url, Credentials.OAuth(session), + onConnect = SuspendBiConsumer { _, _ -> router.finishSignIn(currentCoroutineContext().job) {} }, + onTokenRefreshed = { _, _ -> throw failure }, + ) + } + val pages = context.envPageManager + val logger = context.logger + val log = logger.delegate + + withMockCli { + context.settingsStore.updatePreferAuthViaOAuth2(true) + page.beforeShow() + runCurrent() + + assertEquals(WizardStep.SIGNED_IN, page.model.currentStep()) + verify(exactly = 1) { pages.showPluginEnvironmentsPage(false) } + verify(exactly = 1) { log.warn(failure, any()) } + verify(exactly = 0) { logger.showInfoPopup(any(), any()) } + verify(exactly = 0) { anyConstructed().close() } + page.dispose() + } + } + @Test fun `submitting a token displays the connection step once`() = runTest { every { context.cs } returns backgroundScope From de8c881f1122d4a8458fa84c73015fdbac8a3c5d Mon Sep 17 00:00:00 2001 From: Faur Ioan-Aurel Date: Fri, 13 Nov 2026 18:36:37 +0200 Subject: [PATCH 4/7] Log in to the CLI once when restoring an OAuth session Setup logs the CLI in with the current OAuth token. Calling the token refresh handler again after sign-in repeats that login and can leave the CLI and the saved session out of sync. The wizard now saves the OAuth session directly after setup. The refresh handler still updates credentials when a running session receives a new token. A test checks the login count and saved token. --- .../com/coder/toolbox/views/ConnectStep.kt | 4 +- .../coder/toolbox/CoderRemoteProviderTest.kt | 42 +++++++++++++++++++ .../toolbox/views/CoderSetupWizardPageTest.kt | 2 +- 3 files changed, 45 insertions(+), 3 deletions(-) diff --git a/src/main/kotlin/com/coder/toolbox/views/ConnectStep.kt b/src/main/kotlin/com/coder/toolbox/views/ConnectStep.kt index 5963d081..678a6a43 100644 --- a/src/main/kotlin/com/coder/toolbox/views/ConnectStep.kt +++ b/src/main/kotlin/com/coder/toolbox/views/ConnectStep.kt @@ -133,9 +133,9 @@ class ConnectStep( context.logger.info("Connection setup done, initializing the workspace poller...") onConnect.accept(client, cli) pendingClient = null // The provider now owns this client. - // Only invoke onTokenRefreshed when we actually have an OAuth session + // The CLI already uses this token. Save the OAuth session for the next sign-in. oauthSession?.let { session -> - onTokenRefreshed?.invoke(client.url, session) + context.secrets.storeOAuthFor(client.url.toString(), session) } } catch (ex: Exception) { reportConnectionFailure(currentCoroutineContext().job, url, ex) diff --git a/src/test/kotlin/com/coder/toolbox/CoderRemoteProviderTest.kt b/src/test/kotlin/com/coder/toolbox/CoderRemoteProviderTest.kt index 6b8d5b9f..d3ecdbf9 100644 --- a/src/test/kotlin/com/coder/toolbox/CoderRemoteProviderTest.kt +++ b/src/test/kotlin/com/coder/toolbox/CoderRemoteProviderTest.kt @@ -3,6 +3,7 @@ package com.coder.toolbox import com.coder.toolbox.cli.CoderCLIManager import com.coder.toolbox.cli.ensureCLI import com.coder.toolbox.diagnostics.CoderLogger +import com.coder.toolbox.oauth.OAuth2Client import com.coder.toolbox.oauth.TokenEndpointAuthMethod import com.coder.toolbox.sdk.CoderRestClient import com.coder.toolbox.sdk.DataGen @@ -209,6 +210,47 @@ class CoderRemoteProviderTest { coVerify(exactly = 0) { mockCli.logout() } } + @Test + @OptIn(ExperimentalCoroutinesApi::class) + fun `restoring an OAuth session logs in to the CLI once and saves the rotated token`() = runTest { + val url = mockCli.deploymentURL + val stored = storedOAuthSession() + val session = stored.toSessionContext() + val refreshed = session.tokenResponse!!.copy(accessToken = "new-access-token", refreshToken = "new-refresh-token") + val secrets = mockContext.secrets + every { secrets.oauthSessionFor(url.toString()) } returns stored + every { mockContext.deploymentUrl } returns url + every { mockContext.cs } returns backgroundScope + every { mockContext.settingsStore } returns CoderSettingsStore( + pluginTestSettingsStore(), Environment(), underlyingLogger, + ).apply { updatePreferAuthViaOAuth2(true) } + every { mockContext.proxySettings.getProxy() } returns Proxy.NO_PROXY + every { mockContext.proxySettings.getProxyAuth() } returns null + + mockkConstructor(OAuth2Client::class, CoderRestClient::class) + mockkStatic(::ensureCLI) + try { + coEvery { anyConstructed().refreshToken(session) } returns refreshed + val user = DataGen.user() + coEvery { anyConstructed().initializeSession() } returns user + every { anyConstructed().me } returns user + every { anyConstructed().buildVersion } returns "2.39.0" + coEvery { anyConstructed().templates() } returns emptyList() + coEvery { anyConstructed().workspaces(any()) } returns emptyList() + coEvery { ensureCLI(mockContext, url, "2.39.0", any()) } returns mockCli + + assertNull(remoteProvider.getOverrideUiPage()) + runCurrent() + + coVerify(exactly = 1) { mockCli.login(refreshed.accessToken) } + verify(exactly = 1) { secrets.storeOAuthFor(url.toString(), session.copy(tokenResponse = refreshed)) } + } finally { + remoteProvider.close() + unmockkStatic(::ensureCLI) + unmockkConstructor(OAuth2Client::class, CoderRestClient::class) + } + } + @Test fun `OAuth refresh updates the CLI session for the active origin`() = runTest { setPrivateField(remoteProvider, "client", mockClient) diff --git a/src/test/kotlin/com/coder/toolbox/views/CoderSetupWizardPageTest.kt b/src/test/kotlin/com/coder/toolbox/views/CoderSetupWizardPageTest.kt index 4fb02c7a..decd2445 100644 --- a/src/test/kotlin/com/coder/toolbox/views/CoderSetupWizardPageTest.kt +++ b/src/test/kotlin/com/coder/toolbox/views/CoderSetupWizardPageTest.kt @@ -119,11 +119,11 @@ class CoderSetupWizardPageTest { ) val router = PageRouter() val failure = IllegalStateException("Keychain is locked") + every { context.secrets.storeOAuthFor(url.toString(), session) } throws failure val page = router.getOrCreate { CoderSetupWizardPage.connectStep( context, mockk(), url, Credentials.OAuth(session), onConnect = SuspendBiConsumer { _, _ -> router.finishSignIn(currentCoroutineContext().job) {} }, - onTokenRefreshed = { _, _ -> throw failure }, ) } val pages = context.envPageManager From 996a40ae896f8bc7f0fa75bf8241685c2f651dd8 Mon Sep 17 00:00:00 2001 From: Faur Ioan-Aurel Date: Fri, 13 Nov 2026 18:59:46 +0200 Subject: [PATCH 5/7] Clarify setup and session restore tests The provider tests use short deployment names and repeat the setup for several credential choices. This makes it hard to see which deployment receives a link and which credentials the wizard should use. The tests name the source and target deployments and cover the restore choices in a table. Their assertions identify when the CLI signs in and which workspace page the provider opens. --- .../coder/toolbox/CoderRemoteProviderTest.kt | 297 ++++++++---------- .../toolbox/views/CoderSetupWizardPageTest.kt | 8 +- 2 files changed, 135 insertions(+), 170 deletions(-) diff --git a/src/test/kotlin/com/coder/toolbox/CoderRemoteProviderTest.kt b/src/test/kotlin/com/coder/toolbox/CoderRemoteProviderTest.kt index d3ecdbf9..9f3a7045 100644 --- a/src/test/kotlin/com/coder/toolbox/CoderRemoteProviderTest.kt +++ b/src/test/kotlin/com/coder/toolbox/CoderRemoteProviderTest.kt @@ -79,6 +79,9 @@ import kotlin.test.assertNotNull import kotlin.test.assertNull import kotlin.test.assertSame +private const val PRODUCTION_URL = "https://coder.example.com" +private const val STAGING_URL = "https://staging.example.com" + class CoderRemoteProviderTest { private lateinit var mockClient: CoderRestClient @@ -108,9 +111,9 @@ class CoderRemoteProviderTest { every { mockContext.settingsStore } returns settingsStore every { mockContext.logger } returns mockLogger every { mockContext.secrets.isUrlTrusted(any()) } returns false - every { mockClient.url } returns URI("https://coder.example.com").toURL() + every { mockClient.url } returns URI(PRODUCTION_URL).toURL() every { mockClient.me } returns DataGen.user() - every { mockCli.deploymentURL } returns URI("https://coder.example.com").toURL() + every { mockCli.deploymentURL } returns URI(PRODUCTION_URL).toURL() remoteProvider = CoderRemoteProvider(mockContext) } @@ -147,43 +150,52 @@ class CoderRemoteProviderTest { } @CsvSource( - "https://coder.example.com, false, true", - "https://other.example.com, false, true", - "https://other.example.com, true, false", + "logged out deployment, link target, other remembered deployment, expect prompt", + "$PRODUCTION_URL, $PRODUCTION_URL, none, true", + "$STAGING_URL, $STAGING_URL, $PRODUCTION_URL, true", + "$PRODUCTION_URL, $STAGING_URL, none, true", + "$STAGING_URL, $PRODUCTION_URL, $PRODUCTION_URL, false", + nullValues = ["none"], + useHeadersInDisplayName = true, ) - @ParameterizedTest(name = "After logout, link to {0}, B remembered={1}, deployment prompt={2}") - fun `logout forgets A while each subsequent link uses its own deployment history`( - target: String, - rememberB: Boolean, + @ParameterizedTest(name = "After logout, {arguments}") + fun `logout clears the active deployment and each link uses its target's trust history`( + loggedOutDeployment: URL, + linkTarget: URL, + otherRememberedDeployment: URL?, expectPrompt: Boolean, ) = runTest { - val a = mockClient.url - val b = URI("https://other.example.com").toURL() + every { mockClient.url } returns loggedOutDeployment + every { mockCli.deploymentURL } returns loggedOutDeployment val secrets = CoderSecretsStore(InMemoryPluginSecretStore()) every { mockContext.secrets } returns secrets every { mockContext.settingsStore.requiresTokenAuth } returns true every { mockContext.settingsStore.requiresMTlsAuth } returns false every { mockContext.settingsStore.preferOAuth2IfAvailable } returns false - secrets.storeApiTokenFor(a, "saved-a") - secrets.trustDeployment(a) - if (rememberB) { - secrets.storeApiTokenFor(b, "saved-b") - secrets.trustDeployment(b) - } + secrets.storeApiTokenFor(loggedOutDeployment, "logged-out-token") + secrets.trustDeployment(loggedOutDeployment) val provider = spyk(remoteProvider, recordPrivateCalls = true) setPrivateField(provider, "client", mockClient) setPrivateField(provider, "cli", mockCli) - val user = DataGen.user() - coEvery { provider["userForToken"](b, "saved-b") } returns user - coEvery { provider["userForToken"](b, "link-token") } returns user + if (otherRememberedDeployment != null) { + secrets.storeApiTokenFor(otherRememberedDeployment, "remembered-token") + secrets.trustDeployment(otherRememberedDeployment) + val rememberedUser = DataGen.user() + coEvery { provider["userForToken"](otherRememberedDeployment, "remembered-token") } returns rememberedUser + coEvery { provider["userForToken"](otherRememberedDeployment, "link-token") } returns rememberedUser + } coEvery { mockUi.showYesNoPopup(any(), any(), any(), any()) } returns false provider.logout() - provider.handleUri(workspaceLink(target)) - - assertNull(secrets.apiTokenFor(a)) - assertFalse(secrets.isUrlTrusted(a)) - assertEquals(rememberB, secrets.isUrlTrusted(b)) + provider.handleUri(workspaceLink(linkTarget.toString())) + + assertNull(secrets.apiTokenFor(loggedOutDeployment)) + assertFalse(secrets.isUrlTrusted(loggedOutDeployment)) + assertEquals(!expectPrompt, secrets.isUrlTrusted(linkTarget)) + if (otherRememberedDeployment != null) { + assertTrue(secrets.isUrlTrusted(otherRememberedDeployment)) + assertEquals("remembered-token", secrets.apiTokenFor(otherRememberedDeployment)) + } coVerify(exactly = if (expectPrompt) 1 else 0) { mockUi.showYesNoPopup(any(), any(), any(), any()) } coVerify(exactly = if (expectPrompt) 0 else 1) { mockUi.showUiPageSuspending(any()) } } @@ -695,41 +707,43 @@ class CoderRemoteProviderTest { verify(exactly = 0) { mockContext.popupPluginMainPage() } } - @ValueSource( - strings = [ - "https://coder2.example.com/", - "https://coder2.example.com/dashboard?view=recent#workspaces", - ] + @CsvSource( + "restoring deployment, link target, expected origin", + "$PRODUCTION_URL, $STAGING_URL/, $STAGING_URL", + "$STAGING_URL, $PRODUCTION_URL/dashboard?view=recent#workspaces, $PRODUCTION_URL", + "$PRODUCTION_URL, $STAGING_URL:8443/dashboard, $STAGING_URL:8443", + useHeadersInDisplayName = true, ) - @ParameterizedTest - fun `accepted cold start link presents the new deployment while the old session restores`(target: String) = - runTest { - every { mockContext.cs } returns backgroundScope - val previousUrl = URI("https://coder.example.com").toURL() - every { mockContext.deploymentUrl } returns previousUrl - every { mockContext.secrets.apiTokenFor(previousUrl) } returns "saved-token" - every { mockContext.settingsStore.requiresMTlsAuth } returns false - every { mockContext.settingsStore.requiresTokenAuth } returns true - coEvery { mockUi.showYesNoPopup(any(), any(), any(), any()) } returns true - assertNull(remoteProvider.getOverrideUiPage()) - val previousWizard = setupWizard(remoteProvider) - assertEquals(previousUrl, previousWizard.model.url) - assertEquals(WizardStep.RESTORE_SESSION, previousWizard.model.currentStep()) - val displayedWizard = slot() - coEvery { mockUi.showUiPageSuspending(capture(displayedWizard)) } returns Unit - - remoteProvider.handleUri(workspaceLink(target)) - - // Selecting an already visible provider can leave Toolbox on the old page. - // No second visibility callback or request for the override page is simulated. - coVerify(exactly = 1) { mockUi.showUiPageSuspending(any()) } - val wizard = displayedWizard.captured - assertEquals(WizardStep.CONNECT, wizard.model.currentStep()) - assertEquals("https://coder2.example.com", wizard.model.url.toString()) - assertEquals("link-token", wizard.model.token) - assertSame(wizard, remoteProvider.getOverrideUiPage()) - coVerify(exactly = 1) { mockUi.showYesNoPopup(any(), any(), any(), any()) } - } + @ParameterizedTest(name = "Cold-start link, {arguments}") + fun `accepted cold start link presents the new deployment while the old session restores`( + restoringDeployment: URL, + linkTarget: String, + expectedOrigin: String, + ) = runTest { + every { mockContext.cs } returns backgroundScope + every { mockContext.deploymentUrl } returns restoringDeployment + every { mockContext.secrets.apiTokenFor(restoringDeployment) } returns "saved-token" + every { mockContext.settingsStore.requiresMTlsAuth } returns false + every { mockContext.settingsStore.requiresTokenAuth } returns true + coEvery { mockUi.showYesNoPopup(any(), any(), any(), any()) } returns true + assertNull(remoteProvider.getOverrideUiPage()) + val restoringWizard = setupWizard(remoteProvider) + assertEquals(restoringDeployment, restoringWizard.model.url) + assertEquals(WizardStep.RESTORE_SESSION, restoringWizard.model.currentStep()) + val displayedWizard = slot() + coEvery { mockUi.showUiPageSuspending(capture(displayedWizard)) } returns Unit + + remoteProvider.handleUri(workspaceLink(linkTarget)) + + // Toolbox has not requested the override page again since the link arrived. + coVerify(exactly = 1) { mockUi.showUiPageSuspending(any()) } + val replacementWizard = displayedWizard.captured + assertEquals(WizardStep.CONNECT, replacementWizard.model.currentStep()) + assertEquals(expectedOrigin, replacementWizard.model.url.toString()) + assertEquals("link-token", replacementWizard.model.token) + assertSame(replacementWizard, remoteProvider.getOverrideUiPage()) + coVerify(exactly = 1) { mockUi.showYesNoPopup(any(), any(), any(), any()) } + } private fun setupWizard(provider: CoderRemoteProvider): CoderSetupWizardPage { val router = @@ -739,9 +753,9 @@ class CoderRemoteProviderTest { @Test @OptIn(ExperimentalCoroutinesApi::class) - fun `a cold link signs in to B and hands off its workspace without a focus change`() = runTest { - val a = URI("https://a.example.com").toURL() - val b = URI("https://b.example.com").toURL() + fun `a cold link signs in to the linked deployment and hands off its workspace without a focus change`() = runTest { + val restoringDeployment = URI(PRODUCTION_URL).toURL() + val linkedDeployment = URI(STAGING_URL).toURL() val secrets = mockContext.secrets val environmentsPage = mockContext.envPageManager every { mockContext.settingsStore } returns CoderSettingsStore( @@ -750,48 +764,48 @@ class CoderRemoteProviderTest { every { mockContext.cs } returns backgroundScope every { mockContext.proxySettings.getProxy() } returns Proxy.NO_PROXY every { mockContext.proxySettings.getProxyAuth() } returns null - every { mockContext.deploymentUrl } returns a - every { secrets.apiTokenFor(a) } returns "saved-a" - val handler = mockk(relaxed = true) - setPrivateField(remoteProvider, "linkHandler", handler) + every { mockContext.deploymentUrl } returns restoringDeployment + every { secrets.apiTokenFor(restoringDeployment) } returns "saved-token" + val linkHandler = mockk(relaxed = true) + setPrivateField(remoteProvider, "linkHandler", linkHandler) val user = DataGen.user() - val aStarted = CompletableDeferred() - val aCancelled = CompletableDeferred() - val shown = slot() + val restoreStarted = CompletableDeferred() + val restoreCancelled = CompletableDeferred() + val displayedWizard = slot() mockkConstructor(CoderRestClient::class) mockkStatic(::ensureCLI) try { coEvery { anyConstructed().initializeSession() } coAnswers { - aStarted.complete(Unit) + restoreStarted.complete(Unit) try { awaitCancellation() } finally { - aCancelled.complete(Unit) + restoreCancelled.complete(Unit) } } andThen user every { anyConstructed().me } returns user every { anyConstructed().buildVersion } returns "2.39.0" coEvery { anyConstructed().workspaces(any()) } returns emptyList() - coEvery { ensureCLI(mockContext, b, "2.39.0", any()) } returns mockCli + coEvery { ensureCLI(mockContext, linkedDeployment, "2.39.0", any()) } returns mockCli coEvery { mockUi.showYesNoPopup(any(), any(), any(), any()) } returns true - coEvery { mockUi.showUiPageSuspending(capture(shown)) } coAnswers { - shown.captured.beforeShow() + coEvery { mockUi.showUiPageSuspending(capture(displayedWizard)) } coAnswers { + displayedWizard.captured.beforeShow() } every { environmentsPage.showPluginEnvironmentsPage(false) } answers { assertNull(remoteProvider.getOverrideUiPage()) - shown.captured.afterHide() + displayedWizard.captured.afterHide() } assertNull(remoteProvider.getOverrideUiPage()) runCurrent() - assertTrue(aStarted.isCompleted) - remoteProvider.handleUri(workspaceLink(b.toString())) + assertTrue(restoreStarted.isCompleted) + remoteProvider.handleUri(workspaceLink(linkedDeployment.toString())) runCurrent() - assertTrue(aCancelled.isCompleted) + assertTrue(restoreCancelled.isCompleted) verify(exactly = 1) { anyConstructed().close() } coVerify(exactly = 1) { mockCli.login("link-token") } - verify(exactly = 1) { secrets.storeApiTokenFor(b, "link-token") } - coVerify(exactly = 1) { handler.handle(any(), b, any(), mockCli) } + verify(exactly = 1) { secrets.storeApiTokenFor(linkedDeployment, "link-token") } + coVerify(exactly = 1) { linkHandler.handle(match { it["workspace"] == "dev" }, linkedDeployment, any(), mockCli) } verify(exactly = 1) { environmentsPage.showPluginEnvironmentsPage(false) } assertNull(remoteProvider.getOverrideUiPage()) } finally { @@ -1201,7 +1215,7 @@ class CoderRemoteProviderTest { @Test @OptIn(ExperimentalCoroutinesApi::class) - fun `sign-in behind Settings opens the workspace list when the user returns to setup`() = runTest { + fun `returning from Settings after sign-in opens the workspace list once even if beforeShow runs twice`() = runTest { every { mockContext.cs } returns backgroundScope coEvery { mockClient.workspaces(any()) } returns emptyList() val wizard = remoteProvider.getOverrideUiPage() as CoderSetupWizardPage @@ -1215,8 +1229,8 @@ class CoderRemoteProviderTest { verify(exactly = 0) { pages.showPluginEnvironmentsPage(false) } assertNull(remoteProvider.getOverrideUiPage()) - wizard.beforeShow() - wizard.beforeShow() + // Both visibility callbacks arrive before the step observer runs. + repeat(2) { wizard.beforeShow() } runCurrent() verify(exactly = 1) { pages.showPluginEnvironmentsPage(false) } wizard.afterHide() @@ -1328,95 +1342,42 @@ class CoderRemoteProviderTest { assertFalse(accountDropdown.visibility.value) } - @Test - fun `given mTLS is required when auto setup has stored credentials then mTLS takes precedence`() = runTest { - every { mockContext.cs } returns backgroundScope - // given - val url = URI("https://coder.example.com").toURL() - every { mockContext.deploymentUrl } returns url - every { mockContext.settingsStore.requiresMTlsAuth } returns true - every { mockContext.settingsStore.requiresTokenAuth } returns false - every { mockContext.settingsStore.preferOAuth2IfAvailable } returns true - every { mockContext.secrets.apiTokenFor(url) } returns "token" - every { mockContext.secrets.oauthSessionFor(url.toString()) } returns storedOAuthSession() - val provider = CoderRemoteProvider(mockContext) - - // when - assertNull(provider.getOverrideUiPage()) - val overridePage = setupWizard(provider) - - // then - assertEquals(WizardStep.RESTORE_SESSION, overridePage.model.currentStep()) - assertNull(overridePage.model.token) - assertNull(overridePage.model.oauthSession) - } - - @Test - fun `given OAuth is preferred when auto setup has token and OAuth session then OAuth is used`() = runTest { - every { mockContext.cs } returns backgroundScope - // given - val url = URI("https://coder.example.com").toURL() - every { mockContext.deploymentUrl } returns url - every { mockContext.settingsStore.requiresMTlsAuth } returns false - every { mockContext.settingsStore.requiresTokenAuth } returns true - every { mockContext.settingsStore.preferOAuth2IfAvailable } returns true - every { mockContext.secrets.apiTokenFor(url) } returns "token" - every { mockContext.secrets.oauthSessionFor(url.toString()) } returns storedOAuthSession() - val provider = CoderRemoteProvider(mockContext) - - // when - assertNull(provider.getOverrideUiPage()) - val overridePage = setupWizard(provider) - - // then - assertEquals(WizardStep.RESTORE_SESSION, overridePage.model.currentStep()) - assertNull(overridePage.model.token) - assertNotNull(overridePage.model.oauthSession) - } - - @Test - fun `given OAuth is not preferred when auto setup has token and OAuth session then token is used`() = runTest { - every { mockContext.cs } returns backgroundScope - // given - val url = URI("https://coder.example.com").toURL() - every { mockContext.deploymentUrl } returns url - every { mockContext.settingsStore.requiresMTlsAuth } returns false - every { mockContext.settingsStore.requiresTokenAuth } returns true - every { mockContext.settingsStore.preferOAuth2IfAvailable } returns false - every { mockContext.secrets.apiTokenFor(url) } returns "token" - every { mockContext.secrets.oauthSessionFor(url.toString()) } returns storedOAuthSession() - val provider = CoderRemoteProvider(mockContext) - - // when - assertNull(provider.getOverrideUiPage()) - val overridePage = setupWizard(provider) - - // then - assertEquals(WizardStep.RESTORE_SESSION, overridePage.model.currentStep()) - assertEquals("token", overridePage.model.token) - assertNull(overridePage.model.oauthSession) - } - - @Test - fun `given OAuth is not preferred when auto setup has API token then token is used`() = runTest { + @CsvSource( + "require mTLS, prefer OAuth, saved API token, saved OAuth session, expected API token, expect OAuth", + "true, true, api-token, true, none, false", + "false, true, api-token, true, none, true", + "false, false, api-token, true, api-token, false", + "false, false, api-token, false, api-token, false", + "false, true, api-token, false, api-token, false", + "false, true, none, true, none, true", + nullValues = ["none"], + useHeadersInDisplayName = true, + ) + @ParameterizedTest(name = "Restore credentials, {arguments}") + fun `session restore selects credentials using the deployment settings`( + requireMtls: Boolean, + preferOAuth: Boolean, + savedApiToken: String?, + hasSavedOAuthSession: Boolean, + expectedApiToken: String?, + expectOAuth: Boolean, + ) = runTest { every { mockContext.cs } returns backgroundScope - // given - val url = URI("https://coder.example.com").toURL() + val url = mockClient.url + val savedOAuthSession = if (hasSavedOAuthSession) storedOAuthSession() else null every { mockContext.deploymentUrl } returns url - every { mockContext.settingsStore.requiresMTlsAuth } returns false - every { mockContext.settingsStore.requiresTokenAuth } returns true - every { mockContext.settingsStore.preferOAuth2IfAvailable } returns false - every { mockContext.secrets.apiTokenFor(url) } returns "api-token" - val provider = CoderRemoteProvider(mockContext) + every { mockContext.settingsStore.requiresMTlsAuth } returns requireMtls + every { mockContext.settingsStore.requiresTokenAuth } returns !requireMtls + every { mockContext.settingsStore.preferOAuth2IfAvailable } returns preferOAuth + every { mockContext.secrets.apiTokenFor(url) } returns savedApiToken + every { mockContext.secrets.oauthSessionFor(url.toString()) } returns savedOAuthSession - // when - assertNull(provider.getOverrideUiPage()) - val overridePage = setupWizard(provider) + assertNull(remoteProvider.getOverrideUiPage()) + val model = setupWizard(remoteProvider).model - // then - assertEquals(WizardStep.RESTORE_SESSION, overridePage.model.currentStep()) - assertEquals("api-token", overridePage.model.token) - assertNull(overridePage.model.oauthSession) + assertEquals(WizardStep.RESTORE_SESSION, model.currentStep()) + assertEquals(expectedApiToken, model.token) + assertEquals(if (expectOAuth) requireNotNull(savedOAuthSession).toSessionContext() else null, model.oauthSession) } @Test diff --git a/src/test/kotlin/com/coder/toolbox/views/CoderSetupWizardPageTest.kt b/src/test/kotlin/com/coder/toolbox/views/CoderSetupWizardPageTest.kt index decd2445..9ff1947c 100644 --- a/src/test/kotlin/com/coder/toolbox/views/CoderSetupWizardPageTest.kt +++ b/src/test/kotlin/com/coder/toolbox/views/CoderSetupWizardPageTest.kt @@ -38,6 +38,8 @@ import kotlinx.coroutines.test.StandardTestDispatcher import kotlinx.coroutines.test.runCurrent import kotlinx.coroutines.test.runTest import kotlinx.coroutines.withContext +import org.junit.jupiter.params.ParameterizedTest +import org.junit.jupiter.params.provider.EnumSource import java.net.Proxy import java.net.URI import kotlin.coroutines.CoroutineContext @@ -171,8 +173,10 @@ class CoderSetupWizardPageTest { } } - @Test - fun `a second visibility callback during startup logs in to the CLI once`() = runTest { + @EnumSource(WizardStep::class, names = ["CONNECT", "RESTORE_SESSION"]) + @ParameterizedTest(name = "Repeated visibility during {0}") + fun `a second visibility callback during startup logs in to the CLI once`(connectionStep: WizardStep) = runTest { + model.goTo(connectionStep) lateinit var step: ConnectStep val scheduler = StandardTestDispatcher(testScheduler) val dispatcher = object : CoroutineDispatcher() { From ca7f86ca600583ede8c9d61014f0d79dcc346784 Mon Sep 17 00:00:00 2001 From: Faur Ioan-Aurel Date: Fri, 13 Nov 2026 19:23:08 +0200 Subject: [PATCH 6/7] Validate the session token before CLI login Setup assumes a token exists when it hands credentials to the CLI. If the user switches from certificate to token authentication during setup, a missing token can produce an error with no explanation. The HTTP client and CLI now use the same token check. The wizard reports a missing token and skips CLI login for certificate sessions. Tests cover both paths. --- .../com/coder/toolbox/sdk/CoderRestClient.kt | 12 ++--- .../com/coder/toolbox/views/ConnectStep.kt | 31 ++++++------- .../toolbox/views/CoderSetupWizardPageTest.kt | 44 +++++++++++++++++++ 3 files changed, 64 insertions(+), 23 deletions(-) diff --git a/src/main/kotlin/com/coder/toolbox/sdk/CoderRestClient.kt b/src/main/kotlin/com/coder/toolbox/sdk/CoderRestClient.kt index e1843ecc..f0bde3f8 100644 --- a/src/main/kotlin/com/coder/toolbox/sdk/CoderRestClient.kt +++ b/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. * diff --git a/src/main/kotlin/com/coder/toolbox/views/ConnectStep.kt b/src/main/kotlin/com/coder/toolbox/views/ConnectStep.kt index 678a6a43..c01d3060 100644 --- a/src/main/kotlin/com/coder/toolbox/views/ConnectStep.kt +++ b/src/main/kotlin/com/coder/toolbox/views/ConnectStep.kt @@ -88,13 +88,11 @@ class ConnectStep( private suspend fun connect(url: URL) { var pendingClient: CoderRestClient? = null try { - var oauthSession: CoderOAuthSessionContext? = null - if (context.settingsStore.requiresTokenAuth && context.settingsStore.preferOAuth2IfAvailable && model.hasOAuthSession()) { - refreshOAuthToken() - oauthSession = model.oauthSession!!.copy() - } - - val apiToken = if (context.settingsStore.requiresTokenAuth) model.token else null + val requiresTokenAuth = context.settingsStore.requiresTokenAuth + val oauthSession = if (requiresTokenAuth && context.settingsStore.preferOAuth2IfAvailable) { + oauthSessionForConnection() + } else null + val apiToken = if (requiresTokenAuth) model.token else null context.logger.info("Setting up the HTTP client...") val client = CoderRestClient( @@ -116,16 +114,11 @@ class ConnectStep( ) { progress -> statusField.textState.update { (context.i18n.pnotr(progress)) } } - // We only need to log in if we are using token-based auth. if (context.settingsStore.requiresTokenAuth) { logAndReportProgress("Configuring Coder CLI...") // allows interleaving with the back/cancel action yield() - if (oauthSession != null) { - cli.login(oauthSession.tokenResponse!!.accessToken) - } else { - cli.login(apiToken!!) - } + cli.login(client.requireSessionToken()) } logAndReportProgress("Successfully configured ${url.host}...") // allows interleaving with the back/cancel action @@ -144,14 +137,16 @@ class ConnectStep( } } - private suspend fun refreshOAuthToken() { - val session = model.oauthSession ?: return - if (!session.tokenResponse?.accessToken.isNullOrBlank()) return + private suspend fun oauthSessionForConnection(): CoderOAuthSessionContext? { + val session = model.oauthSession ?: return null + val response = session.tokenResponse ?: return null + if (response.accessToken.isNotBlank()) return session logAndReportProgress("Refreshing OAuth token...") - val tokenResponse = OAuth2Client(context).refreshToken(session) + val refreshed = session.copy(tokenResponse = OAuth2Client(context).refreshToken(session)) context.logger.info("Successfully refreshed access token") - model.oauthSession = session.copy(tokenResponse = tokenResponse) + model.oauthSession = refreshed + return refreshed } private fun logAndReportProgress(msg: String) { diff --git a/src/test/kotlin/com/coder/toolbox/views/CoderSetupWizardPageTest.kt b/src/test/kotlin/com/coder/toolbox/views/CoderSetupWizardPageTest.kt index 9ff1947c..0b5ce3e4 100644 --- a/src/test/kotlin/com/coder/toolbox/views/CoderSetupWizardPageTest.kt +++ b/src/test/kotlin/com/coder/toolbox/views/CoderSetupWizardPageTest.kt @@ -22,6 +22,7 @@ import io.mockk.justRun import io.mockk.mockk import io.mockk.mockkConstructor import io.mockk.mockkStatic +import io.mockk.spyk import io.mockk.unmockkConstructor import io.mockk.unmockkStatic import io.mockk.verify @@ -201,6 +202,49 @@ class CoderSetupWizardPageTest { } } + @Test + fun `certificate sign-in connects without a CLI login`() = runTest { + every { context.cs } returns backgroundScope + model.token = null + val step = ConnectStep(context, model, { model.goTo(WizardStep.URL_REQUEST) }, onConnect) + withMockCli { cli -> + val settings = spyk(context.settingsStore) + every { settings.requiresTokenAuth } returns false + every { context.settingsStore } returns settings + + step.onVisible() + runCurrent() + + coVerify(exactly = 0) { cli.login(any()) } + coVerify(exactly = 1) { onConnect.accept(any(), cli) } + } + } + + @Test + fun `switching from certificates to token authentication during setup reports the missing token`() = runTest { + every { context.cs } returns backgroundScope + model.token = null + val step = ConnectStep(context, model, { model.goTo(WizardStep.TOKEN_REQUEST) }, onConnect) + val logger = context.logger + withMockCli { cli -> + val settings = spyk(context.settingsStore) + var requiresToken = false + every { settings.requiresTokenAuth } answers { requiresToken } + every { context.settingsStore } returns settings + coEvery { anyConstructed().initializeSession() } coAnswers { + requiresToken = true + DataGen.user() + } + step.onVisible() + runCurrent() + + assertEquals(WizardStep.TOKEN_REQUEST, model.currentStep()) + verify { logger.showInfoPopup(any(), match { it.contains("OAuth or API token is required") }) } + coVerify(exactly = 0) { cli.login(any()) } + coVerify(exactly = 0) { onConnect.accept(any(), any()) } + } + } + private fun withMockCli(test: (CoderCLIManager) -> Unit) { every { context.settingsStore } returns CoderSettingsStore( pluginTestSettingsStore(), Environment(), mockk(relaxed = true), From ac38953ac5cdfd06dcb3b337a5d8dff55e1a9398 Mon Sep 17 00:00:00 2001 From: Faur Ioan-Aurel Date: Fri, 13 Nov 2026 19:52:12 +0200 Subject: [PATCH 7/7] Start saved-session restore when the router opens the wizard Looking up a page to display can start a CLI sign-in. The wizard also uses a separate step and a flag to track how the connection started. The router now starts a saved-session restore when it registers the wizard. Looking up the page only selects what Toolbox displays. The wizard records the origin of the connection so a failed restore can show the sign-in form and repeated callbacks reuse the active job. --- .../com/coder/toolbox/CoderRemoteProvider.kt | 3 +- .../toolbox/views/CoderSetupWizardPage.kt | 45 +++++++++---------- .../com/coder/toolbox/views/ConnectStep.kt | 10 +++-- .../coder/toolbox/views/state/PageRouter.kt | 28 +++++------- .../coder/toolbox/views/state/WizardModel.kt | 8 +++- .../coder/toolbox/CoderRemoteProviderTest.kt | 12 ++--- .../toolbox/views/CoderSetupWizardPageTest.kt | 38 ++++++++++++---- .../toolbox/views/state/PageRouterTest.kt | 4 +- 8 files changed, 82 insertions(+), 66 deletions(-) diff --git a/src/main/kotlin/com/coder/toolbox/CoderRemoteProvider.kt b/src/main/kotlin/com/coder/toolbox/CoderRemoteProvider.kt index 4f7ab6fc..9cad4054 100644 --- a/src/main/kotlin/com/coder/toolbox/CoderRemoteProvider.kt +++ b/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 @@ -890,7 +891,7 @@ class CoderRemoteProvider( credentials = credentials, onConnect = onConnect, onTokenRefreshed = ::onTokenRefreshed, - restoreSavedSession = true, + origin = ConnectionOrigin.SAVED_SESSION, ) } catch (ex: Exception) { context.logger.logAndShowError( diff --git a/src/main/kotlin/com/coder/toolbox/views/CoderSetupWizardPage.kt b/src/main/kotlin/com/coder/toolbox/views/CoderSetupWizardPage.kt index a6dc515f..7b742298 100644 --- a/src/main/kotlin/com/coder/toolbox/views/CoderSetupWizardPage.kt +++ b/src/main/kotlin/com/coder/toolbox/views/CoderSetupWizardPage.kt @@ -4,6 +4,7 @@ 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 @@ -20,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, onTokenRefreshed: (suspend (url: URL, oauthSessionCtx: CoderOAuthSessionContext) -> Unit)? = null ) : CoderPage(MutableStateFlow(context.i18n.ptrl("Setting up Coder")), false) { @@ -49,16 +49,15 @@ class CoderSetupWizardPage private constructor( override val actionButtons: MutableStateFlow> = MutableStateFlow(emptyList()) fun pageForDisplay(): CoderSetupWizardPage? = when (model.currentStep()) { - // Restore saved credentials behind Toolbox's loading workspace list. - // A successful restore then needs no page navigation during window startup. - WizardStep.RESTORE_SESSION -> { - connectStep.onVisible() - null - } + 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() { prepareStepObserver()?.start() } @@ -105,7 +104,7 @@ class CoderSetupWizardPage private constructor( tokenStep.onVisible() } - WizardStep.CONNECT, WizardStep.RESTORE_SESSION -> { + WizardStep.CONNECT -> { fields.update { listOf(connectStep.panel) } @@ -131,22 +130,18 @@ class CoderSetupWizardPage private constructor( } private fun navigateBackFromConnect() { - val restoringSession = model.currentStep() == WizardStep.RESTORE_SESSION - if (autoLogin) { - autoLogin = false - model.clearFormData() - model.goToFirst() - if (restoringSession) { - context.cs.launch { - if (!disposed) context.ui.showUiPageSuspending(this@CoderSetupWizardPage) - } - } + if (model.connectionOrigin == ConnectionOrigin.FORM) { + model.goTo(if (context.settingsStore.requiresTokenAuth) WizardStep.TOKEN_REQUEST else WizardStep.URL_REQUEST) return } - if (context.settingsStore.requiresTokenAuth) { - model.goTo(WizardStep.TOKEN_REQUEST) - } 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) + } } } @@ -187,20 +182,20 @@ class CoderSetupWizardPage private constructor( credentials: Credentials, onConnect: SuspendBiConsumer, onTokenRefreshed: (suspend (url: URL, oauthSessionCtx: CoderOAuthSessionContext) -> Unit)? = null, - restoreSavedSession: Boolean = false, + 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 is Credentials.Token -> model.token = credentials.value is Credentials.OAuth -> model.oauthSession = credentials.session } - model.goTo(if (restoreSavedSession) WizardStep.RESTORE_SESSION else WizardStep.CONNECT) + model.goTo(WizardStep.CONNECT) } } } diff --git a/src/main/kotlin/com/coder/toolbox/views/ConnectStep.kt b/src/main/kotlin/com/coder/toolbox/views/ConnectStep.kt index c01d3060..14cad55c 100644 --- a/src/main/kotlin/com/coder/toolbox/views/ConnectStep.kt +++ b/src/main/kotlin/com/coder/toolbox/views/ConnectStep.kt @@ -8,9 +8,9 @@ import com.coder.toolbox.plugin.PluginManager import com.coder.toolbox.sdk.CoderRestClient import com.coder.toolbox.util.origin import com.coder.toolbox.views.state.CoderOAuthSessionContext +import com.coder.toolbox.views.state.ConnectionOrigin import com.coder.toolbox.views.state.WizardModel import com.coder.toolbox.views.state.WizardStep.CONNECT -import com.coder.toolbox.views.state.WizardStep.RESTORE_SESSION import com.coder.toolbox.views.state.WizardStep.SIGNED_IN import com.jetbrains.toolbox.api.ui.components.LabelField import com.jetbrains.toolbox.api.ui.components.RowGroup @@ -54,13 +54,15 @@ class ConnectStep( RowGroup.RowField(errorField) ) - override fun onVisible() { + override fun onVisible() = startConnection() + + fun startConnection() { prepareConnection()?.start() } @Synchronized private fun prepareConnection(): Job? { - if (disposed || signInJob != null || model.currentStep() !in setOf(CONNECT, RESTORE_SESSION)) return null + if (disposed || signInJob != null || model.currentStep() != CONNECT) return null errorField.textState.update { context.i18n.pnotr("") } if (context.settingsStore.requiresTokenAuth && model.isNotReadyForAuth()) { @@ -177,7 +179,7 @@ class ConnectStep( @Synchronized override fun onBack() { - if (disposed || model.currentStep() != CONNECT) return + if (disposed || model.currentStep() != CONNECT || model.connectionOrigin == ConnectionOrigin.SAVED_SESSION) return try { context.logger.info("Back button was pressed, cancelling in-progress connection setup...") signInJob?.cancel(CancellationException(USER_HIT_THE_BACK_BUTTON)) diff --git a/src/main/kotlin/com/coder/toolbox/views/state/PageRouter.kt b/src/main/kotlin/com/coder/toolbox/views/state/PageRouter.kt index 0ab622dc..e428ba35 100644 --- a/src/main/kotlin/com/coder/toolbox/views/state/PageRouter.kt +++ b/src/main/kotlin/com/coder/toolbox/views/state/PageRouter.kt @@ -3,26 +3,16 @@ package com.coder.toolbox.views.state import com.coder.toolbox.views.CoderSetupWizardPage import kotlinx.coroutines.Job import kotlinx.coroutines.ensureActive -import kotlinx.coroutines.flow.MutableStateFlow /** - * The page that should currently be rendered in place of the environment list. - */ -sealed interface PageRoute { - object None : PageRoute - data class Wizard(val page: CoderSetupWizardPage) : PageRoute -} - -/** - * Holds the active [PageRoute]. The same page instance is returned across + * Holds the active wizard. The same page instance is returned across * Toolbox visibility cycles so in-flight work (e.g. an ongoing connect) is * preserved instead of being thrown away every time the window reopens. */ class PageRouter { - private val route = MutableStateFlow(PageRoute.None) - - val activePage: CoderSetupWizardPage? - get() = (route.value as? PageRoute.Wizard)?.page + @Volatile + var activePage: CoderSetupWizardPage? = null + private set val hasActiveWizard: Boolean get() = activePage != null @@ -38,7 +28,7 @@ class PageRouter { fun getOrCreate(build: () -> CoderSetupWizardPage): CoderSetupWizardPage { activePage?.let { return it } - return build().also { route.value = PageRoute.Wizard(it) } + return build().also(::navigate) } /** @@ -49,13 +39,15 @@ class PageRouter { @Synchronized fun navigate(page: CoderSetupWizardPage) { activePage?.dispose() - route.value = PageRoute.Wizard(page) + activePage = page + // A restore can finish immediately. Publish its wizard before starting the job. + page.startSavedSessionRestore() } @Synchronized fun clear() { activePage?.dispose() - route.value = PageRoute.None + activePage = null } @Synchronized @@ -64,7 +56,7 @@ class PageRouter { connection.ensureActive() startSession() val wizard = activePage ?: return - route.value = PageRoute.None + activePage = null wizard.model.goTo(WizardStep.SIGNED_IN) } } diff --git a/src/main/kotlin/com/coder/toolbox/views/state/WizardModel.kt b/src/main/kotlin/com/coder/toolbox/views/state/WizardModel.kt index b8285cc8..0251c13e 100644 --- a/src/main/kotlin/com/coder/toolbox/views/state/WizardModel.kt +++ b/src/main/kotlin/com/coder/toolbox/views/state/WizardModel.kt @@ -18,6 +18,8 @@ class WizardModel { var url: URL? = null var token: String? = null var oauthSession: CoderOAuthSessionContext? = null + @Volatile + var connectionOrigin = ConnectionOrigin.FORM fun currentStep(): WizardStep = _step.value @@ -43,7 +45,11 @@ class WizardModel { } enum class WizardStep { - URL_REQUEST, TOKEN_REQUEST, CONNECT, RESTORE_SESSION, SIGNED_IN; + URL_REQUEST, TOKEN_REQUEST, CONNECT, SIGNED_IN; +} + +enum class ConnectionOrigin { + FORM, SUPPLIED_CREDENTIALS, SAVED_SESSION } sealed interface Credentials { diff --git a/src/test/kotlin/com/coder/toolbox/CoderRemoteProviderTest.kt b/src/test/kotlin/com/coder/toolbox/CoderRemoteProviderTest.kt index 9f3a7045..9725f768 100644 --- a/src/test/kotlin/com/coder/toolbox/CoderRemoteProviderTest.kt +++ b/src/test/kotlin/com/coder/toolbox/CoderRemoteProviderTest.kt @@ -729,7 +729,7 @@ class CoderRemoteProviderTest { assertNull(remoteProvider.getOverrideUiPage()) val restoringWizard = setupWizard(remoteProvider) assertEquals(restoringDeployment, restoringWizard.model.url) - assertEquals(WizardStep.RESTORE_SESSION, restoringWizard.model.currentStep()) + assertEquals(WizardStep.CONNECT, restoringWizard.model.currentStep()) val displayedWizard = slot() coEvery { mockUi.showUiPageSuspending(capture(displayedWizard)) } returns Unit @@ -792,10 +792,10 @@ class CoderRemoteProviderTest { displayedWizard.captured.beforeShow() } every { environmentsPage.showPluginEnvironmentsPage(false) } answers { - assertNull(remoteProvider.getOverrideUiPage()) + assertNull(remoteProvider.getOverrideUiPage(), "Finishing link sign-in should show the workspace list") displayedWizard.captured.afterHide() } - assertNull(remoteProvider.getOverrideUiPage()) + assertNull(remoteProvider.getOverrideUiPage(), "The first lookup should restore the saved session in the background") runCurrent() assertTrue(restoreStarted.isCompleted) remoteProvider.handleUri(workspaceLink(linkedDeployment.toString())) @@ -807,7 +807,7 @@ class CoderRemoteProviderTest { verify(exactly = 1) { secrets.storeApiTokenFor(linkedDeployment, "link-token") } coVerify(exactly = 1) { linkHandler.handle(match { it["workspace"] == "dev" }, linkedDeployment, any(), mockCli) } verify(exactly = 1) { environmentsPage.showPluginEnvironmentsPage(false) } - assertNull(remoteProvider.getOverrideUiPage()) + assertNull(remoteProvider.getOverrideUiPage(), "The workspace list should still be visible after the handoff") } finally { remoteProvider.close() unmockkStatic(::ensureCLI) @@ -1031,7 +1031,7 @@ class CoderRemoteProviderTest { assertNull(provider.getOverrideUiPage()) assertSame(existingWizard, setupWizard(provider)) - assertEquals(WizardStep.RESTORE_SESSION, existingWizard.model.currentStep()) + assertEquals(WizardStep.CONNECT, existingWizard.model.currentStep()) assertEquals("saved-token", existingWizard.model.token) verify(exactly = 0) { mockContext.popupPluginMainPage() } } @@ -1375,7 +1375,7 @@ class CoderRemoteProviderTest { assertNull(remoteProvider.getOverrideUiPage()) val model = setupWizard(remoteProvider).model - assertEquals(WizardStep.RESTORE_SESSION, model.currentStep()) + assertEquals(WizardStep.CONNECT, model.currentStep()) assertEquals(expectedApiToken, model.token) assertEquals(if (expectOAuth) requireNotNull(savedOAuthSession).toSessionContext() else null, model.oauthSession) } diff --git a/src/test/kotlin/com/coder/toolbox/views/CoderSetupWizardPageTest.kt b/src/test/kotlin/com/coder/toolbox/views/CoderSetupWizardPageTest.kt index 0b5ce3e4..5cf3adc2 100644 --- a/src/test/kotlin/com/coder/toolbox/views/CoderSetupWizardPageTest.kt +++ b/src/test/kotlin/com/coder/toolbox/views/CoderSetupWizardPageTest.kt @@ -11,6 +11,7 @@ import com.coder.toolbox.settings.Environment import com.coder.toolbox.store.CoderSettingsStore import com.coder.toolbox.util.pluginTestSettingsStore 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.WizardModel @@ -47,6 +48,7 @@ import kotlin.coroutines.CoroutineContext import kotlin.test.Test import kotlin.test.assertEquals import kotlin.test.assertNull +import kotlin.test.assertSame @OptIn(ExperimentalCoroutinesApi::class) class CoderSetupWizardPageTest { @@ -60,17 +62,25 @@ class CoderSetupWizardPageTest { private val onConnect = mockk>(relaxed = true) @Test - fun `saved credentials restore once behind the workspace list without page navigation`() = runTest { + fun `routing starts saved session restore and repeated page lookups do not log in again`() = runTest { every { context.cs } returns backgroundScope val page = CoderSetupWizardPage.connectStep( - context, mockk(relaxed = true), url, Credentials.Token("test-token"), onConnect, restoreSavedSession = true, + context, mockk(relaxed = true), url, Credentials.Token("test-token"), onConnect, origin = ConnectionOrigin.SAVED_SESSION, ) val pages = context.envPageManager val ui = context.ui withMockCli { cli -> - assertNull(page.pageForDisplay()) - assertNull(page.pageForDisplay()) + // Toolbox can ask for the page repeatedly. A lookup alone must not start a job. + repeat(2) { assertNull(page.pageForDisplay()) } + runCurrent() + coVerify(exactly = 0) { cli.login(any()) } + + val router = PageRouter() + router.getOrCreate { page } + runCurrent() + coVerify(exactly = 1) { cli.login("test-token") } + repeat(2) { assertNull(router.getOrCreate { error("The routed page should be reused") }.pageForDisplay()) } runCurrent() coVerify(exactly = 1) { cli.login("test-token") } @@ -84,16 +94,19 @@ class CoderSetupWizardPageTest { fun `a failed restore presents the sign-in form and permits a new connection`() = runTest { every { context.cs } returns backgroundScope val page = CoderSetupWizardPage.connectStep( - context, mockk(relaxed = true), url, Credentials.Token("test-token"), onConnect, restoreSavedSession = true, + context, mockk(relaxed = true), url, Credentials.Token("test-token"), onConnect, origin = ConnectionOrigin.SAVED_SESSION, ) val ui = context.ui withMockCli { cli -> coEvery { cli.login("test-token") } throws IllegalStateException("Keychain is locked") andThen "Signed in" + PageRouter().navigate(page) assertNull(page.pageForDisplay()) runCurrent() assertEquals(WizardStep.URL_REQUEST, page.model.currentStep()) + assertEquals(ConnectionOrigin.FORM, page.model.connectionOrigin) + assertSame(page, page.pageForDisplay()) coVerify(exactly = 1) { ui.showUiPageSuspending(page) } coVerify(exactly = 0) { onConnect.accept(any(), any()) } @@ -165,7 +178,9 @@ class CoderSetupWizardPageTest { assertEquals(WizardStep.CONNECT, page.model.currentStep()) verify(exactly = 1) { anyConstructed().onVisible() } + page.dispose() + // Toolbox delivers a late callback after the page observer has been disposed. page.beforeShow() runCurrent() verify(exactly = 1) { anyConstructed().onVisible() } @@ -174,10 +189,10 @@ class CoderSetupWizardPageTest { } } - @EnumSource(WizardStep::class, names = ["CONNECT", "RESTORE_SESSION"]) - @ParameterizedTest(name = "Repeated visibility during {0}") - fun `a second visibility callback during startup logs in to the CLI once`(connectionStep: WizardStep) = runTest { - model.goTo(connectionStep) + @EnumSource(ConnectionOrigin::class) + @ParameterizedTest(name = "Repeated visibility during sign-in from {0}") + fun `a second visibility callback during startup logs in to the CLI once`(origin: ConnectionOrigin) = runTest { + model.connectionOrigin = origin lateinit var step: ConnectStep val scheduler = StandardTestDispatcher(testScheduler) val dispatcher = object : CoroutineDispatcher() { @@ -186,6 +201,7 @@ class CoderSetupWizardPageTest { override fun dispatch(context: CoroutineContext, block: Runnable) { if (firstDispatch) { firstDispatch = false + // Deliver the second callback while the first is still dispatching its sign-in job. step.onVisible() } scheduler.dispatch(context, block) @@ -284,6 +300,8 @@ class CoderSetupWizardPageTest { runCurrent() step.onBack() assertEquals(WizardStep.TOKEN_REQUEST, model.currentStep()) + + // Retry while the cancelled CLI login is still waiting for cleanup. model.goTo(WizardStep.CONNECT) step.onVisible() runCurrent() @@ -316,7 +334,9 @@ class CoderSetupWizardPageTest { model.goTo(WizardStep.CONNECT) step.onVisible() runCurrent() + step.dispose() + // A late visibility callback after disposal must leave the login count at two. step.onVisible() runCurrent() diff --git a/src/test/kotlin/com/coder/toolbox/views/state/PageRouterTest.kt b/src/test/kotlin/com/coder/toolbox/views/state/PageRouterTest.kt index ec93e121..5c2fefde 100644 --- a/src/test/kotlin/com/coder/toolbox/views/state/PageRouterTest.kt +++ b/src/test/kotlin/com/coder/toolbox/views/state/PageRouterTest.kt @@ -55,7 +55,7 @@ class PageRouterTest { fun `a workspace link replaces a wizard being built for auto-connect`() { val router = PageRouter() val automatic = mockk(relaxed = true) - val fromLink = mockk() + val fromLink = mockk(relaxUnitFun = true) val building = CountDownLatch(1) val switching = CountDownLatch(1) val workers = Executors.newFixedThreadPool(2) @@ -85,7 +85,7 @@ class PageRouterTest { @Test fun `a cancelled sign-in cannot publish a session or dismiss its replacement`() { val router = PageRouter() - val replacement = mockk() + val replacement = mockk(relaxUnitFun = true) router.navigate(replacement) val cancelledConnection = Job().apply { cancel() }