impl: connect page router with URI handler

URI handler now schedules pages to be displayed via pager router
and at the end trigger the window to show up which instead it
will call getOverrideUiPage.

Author: fioan89

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

@@ -28,8 +28,7 @@ import com.coder.toolbox.views.CoderSettingsPage
 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.CoderSetupWizardContext
-import com.coder.toolbox.views.state.CoderSetupWizardState
+import com.coder.toolbox.views.state.PageRouter
 import com.coder.toolbox.views.state.WizardStep
 import com.jetbrains.toolbox.api.core.ui.icons.SvgIcon
 import com.jetbrains.toolbox.api.core.ui.icons.SvgIcon.IconType
@@ -113,6 +112,8 @@ class CoderRemoteProvider(
 
     private val errorBuffer = mutableListOf<Throwable>()
 
+    private val router = PageRouter()
+
     /**
      * With the provided client, start polling for workspaces.  Every time a new
      * workspace is added, reconfigure SSH using the provided cli (including the
@@ -269,6 +270,7 @@ class CoderRemoteProvider(
         lastEnvironments.clear()
         environments.value = LoadableState.Value(emptyList())
         isInitialized.update { false }
+        router.clear()
         context.logger.info("Coder plugin is now closed")
     }
 
@@ -345,9 +347,6 @@ class CoderRemoteProvider(
      */
     override suspend fun handleUri(uri: URI) {
         try {
-            // Obtain focus. This switches to the main plugin screen, even
-            // if last opened provider was not Coder
-            context.envPageManager.showPluginEnvironmentsPage()
             if (uri.toString().startsWith("jetbrains://gateway/com.coder.toolbox/auth")) {
                 handleOAuthUri(uri)
                 return
@@ -372,25 +371,25 @@ class CoderRemoteProvider(
                 linkHandler.handle(params, newUrl, this.client!!, this.cli!!)
                 coderHeaderPage.isBusy.update { false }
             } else {
-                // Different URL - we need a new connection.
-                // Chain the link handling after onConnect so it runs once the connection is established.
-                CoderSetupWizardContext.apply {
-                    url = newUrl
-                    token = newToken
-                }
-                CoderSetupWizardState.goToStep(WizardStep.CONNECT)
-                context.ui.showUiPage(
-                    CoderCliSetupWizardPage(
-                        context, settingsPage, visibilityState,
-                        initialAutoSetup = true,
-                        jumpToMainPageOnError = true,
-                        onConnect = onConnect.andThen(deferredLinkHandler(params, newUrl))
-                            .andThen { _, _ ->
-                                coderHeaderPage.isBusy.update { false }
-                            },
-                        onTokenRefreshed = ::onTokenRefreshed
-                    )
+                // Different URL - we need a new connection. Tear down any
+                // in-flight wizard, install a fresh one on the router, and let
+                // showPluginEnvironmentsPage() pull it through getOverrideUiPage.
+                router.activeWizard?.dispose()
+                val wizard = CoderCliSetupWizardPage(
+                    context, settingsPage, visibilityState,
+                    initialAutoSetup = true,
+                    jumpToMainPageOnError = true,
+                    onConnect = onConnect.andThen(deferredLinkHandler(params, newUrl))
+                        .andThen { _, _ ->
+                            coderHeaderPage.isBusy.update { false }
+                        },
+                    onTokenRefreshed = ::onTokenRefreshed
                 )
+                wizard.model.url = newUrl
+                wizard.model.token = newToken
+                wizard.model.goTo(WizardStep.CONNECT)
+                router.replaceWith(wizard)
+                context.envPageManager.showPluginEnvironmentsPage()
             }
         } catch (ex: Exception) {
             val textError = if (ex is APIResponseException) {
@@ -403,7 +402,6 @@ class CoderRemoteProvider(
                 textError ?: ""
             )
             coderHeaderPage.isBusy.update { false }
-            context.envPageManager.showPluginEnvironmentsPage()
         } finally {
             firstRun = false
         }
@@ -424,7 +422,15 @@ class CoderRemoteProvider(
             )
         }
 
-        params["state"]?.takeIf { it == CoderSetupWizardContext.oauthSession?.state }
+        val activeWizard = router.activeWizard ?: return context.logAndShowError(
+            FAILED_TO_HANDLE_OAUTH2_TITLE,
+            "OAuth2 callback arrived but the setup wizard is no longer active"
+        )
+        val activeOAuthSession = activeWizard.model.oauthSession ?: return context.logAndShowError(
+            FAILED_TO_HANDLE_OAUTH2_TITLE,
+            "OAuth2 callback arrived but no OAuth session was started"
+        )
+        params["state"]?.takeIf { it == activeOAuthSession.state }
             ?: return context.logAndShowError(
                 FAILED_TO_HANDLE_OAUTH2_TITLE,
                 "Server responded back with an invalid state that does not match the initial authorization state sent to the server"
@@ -442,18 +448,20 @@ class CoderRemoteProvider(
             )
             return
         }
-        exchangeOAuthCodeForToken(code, CoderSetupWizardContext.oauthSession!!)
+        exchangeOAuthCodeForToken(code, activeOAuthSession, activeWizard)
     }
 
-    private suspend fun exchangeOAuthCodeForToken(code: String, oauthSessionContext: CoderOAuthSessionContext) {
+    private suspend fun exchangeOAuthCodeForToken(
+        code: String,
+        oauthSessionContext: CoderOAuthSessionContext,
+        wizard: CoderCliSetupWizardPage,
+    ) {
         try {
             context.logger.info("Handling OAuth callback...")
 
             val tokenResponse = OAuth2Client(context).exchangeCode(oauthSessionContext, code)
-            CoderSetupWizardContext.oauthSession = oauthSessionContext.copy(tokenResponse = tokenResponse)
-
-            CoderSetupWizardState.goToStep(WizardStep.CONNECT)
-
+            wizard.model.oauthSession = oauthSessionContext.copy(tokenResponse = tokenResponse)
+            wizard.model.goTo(WizardStep.CONNECT)
         } catch (e: Exception) {
             context.logAndShowError("OAuth Error", "Exception during token exchange: ${e.message}", e)
         }
@@ -520,67 +528,70 @@ class CoderRemoteProvider(
      * list.
      */
     override fun getOverrideUiPage(): UiPage? {
-        // Show the setup page if we have not configured the client yet.
-        if (client == null) {
-            // When coming back to the application, initializeSession immediately.
-            if (shouldDoAutoSetup()) {
-                try {
-                    val storedOAuthSession = context.secrets.oauthSessionFor(context.deploymentUrl.toString())
-                    CoderSetupWizardContext.apply {
-                        url = context.deploymentUrl
-                        token = context.secrets.apiTokenFor(context.deploymentUrl)
-                        if (storedOAuthSession != null) {
-                            oauthSession = CoderOAuthSessionContext(
-                                clientId = storedOAuthSession.clientId,
-                                clientSecret = storedOAuthSession.clientSecret,
-                                tokenCodeVerifier = "",
-                                state = "",
-                                tokenEndpoint = storedOAuthSession.tokenEndpoint,
-                                tokenAuthMethod = storedOAuthSession.tokenAuthMethod,
-                                tokenResponse = OAuthTokenResponse(
-                                    accessToken = "",
-                                    tokenType = "",
-                                    expiresIn = null,
-                                    refreshToken = storedOAuthSession.refreshToken,
-                                    scope = null
-                                )
-                            )
-                        }
-                    }
-                    CoderSetupWizardState.goToStep(WizardStep.CONNECT)
-                    return CoderCliSetupWizardPage(
-                        context, settingsPage, visibilityState,
-                        initialAutoSetup = true,
-                        jumpToMainPageOnError = false,
-                        onConnect = onConnect,
-                        onTokenRefreshed = ::onTokenRefreshed
-                    )
-                } catch (ex: Exception) {
-                    errorBuffer.add(ex)
-                } finally {
-                    firstRun = false
-                }
-            }
+        if (client != null) return null
+        return router.requireWizard { buildSetupWizard() }
+    }
 
-            // Login flow.
-            CoderSetupWizardState.goToFirstStep()
-            val setupWizardPage =
-                CoderCliSetupWizardPage(
-                    context,
-                    settingsPage,
-                    visibilityState,
+    /**
+     * Build the wizard for the current state. Called once per provider lifetime
+     * (until [close] clears the router); subsequent visibility cycles reuse the
+     * same instance, preserving any in-flight connect job.
+     */
+    private fun buildSetupWizard(): CoderCliSetupWizardPage {
+        // When coming back to the application, initializeSession immediately.
+        if (shouldDoAutoSetup()) {
+            try {
+                val storedOAuthSession = context.secrets.oauthSessionFor(context.deploymentUrl.toString())
+                val wizard = CoderCliSetupWizardPage(
+                    context, settingsPage, visibilityState,
+                    initialAutoSetup = true,
+                    jumpToMainPageOnError = false,
                     onConnect = onConnect,
                     onTokenRefreshed = ::onTokenRefreshed
                 )
-            // We might have navigated here due to a polling error.
-            errorBuffer.forEach {
-                setupWizardPage.notify("Error encountered", it)
+                wizard.model.url = context.deploymentUrl
+                wizard.model.token = context.secrets.apiTokenFor(context.deploymentUrl)
+                if (storedOAuthSession != null) {
+                    wizard.model.oauthSession = CoderOAuthSessionContext(
+                        clientId = storedOAuthSession.clientId,
+                        clientSecret = storedOAuthSession.clientSecret,
+                        tokenCodeVerifier = "",
+                        state = "",
+                        tokenEndpoint = storedOAuthSession.tokenEndpoint,
+                        tokenAuthMethod = storedOAuthSession.tokenAuthMethod,
+                        tokenResponse = OAuthTokenResponse(
+                            accessToken = "",
+                            tokenType = "",
+                            expiresIn = null,
+                            refreshToken = storedOAuthSession.refreshToken,
+                            scope = null
+                        )
+                    )
+                }
+                wizard.model.goTo(WizardStep.CONNECT)
+                return wizard
+            } catch (ex: Exception) {
+                errorBuffer.add(ex)
+            } finally {
+                firstRun = false
             }
-            errorBuffer.clear()
-            // and now reset the errors, otherwise we show it every time on the screen
-            return setupWizardPage
         }
-        return null
+
+        // Login flow.
+        val setupWizardPage = CoderCliSetupWizardPage(
+            context,
+            settingsPage,
+            visibilityState,
+            onConnect = onConnect,
+            onTokenRefreshed = ::onTokenRefreshed
+        )
+        setupWizardPage.model.goToFirst()
+        // We might have navigated here due to a polling error.
+        errorBuffer.forEach {
+            setupWizardPage.notify("Error encountered", it)
+        }
+        errorBuffer.clear()
+        return setupWizardPage
     }
 
     /**

src/main/kotlin/com/coder/toolbox/cli/downloader/CoderDownloadService.kt

@@ -8,6 +8,7 @@ import com.coder.toolbox.util.getHeaders
 import com.coder.toolbox.util.getOS
 import com.coder.toolbox.util.sha1
 import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.delay
 import kotlinx.coroutines.withContext
 import okhttp3.ResponseBody
 import retrofit2.Response
@@ -24,6 +25,7 @@ import java.nio.file.StandardOpenOption
 import java.util.zip.GZIPInputStream
 import kotlin.io.path.name
 import kotlin.io.path.notExists
+import kotlin.time.Duration.Companion.seconds
 
 private val SUPPORTED_BIN_MIME_TYPES = listOf(
     "application/octet-stream",
@@ -73,6 +75,7 @@ class CoderDownloadService(
                 }
                 context.logger.info("Downloading binary to temporary $cliTempDst")
                 response.saveToDisk(cliTempDst, showTextProgress, buildVersion)?.makeExecutable()
+                delay(10.seconds)
                 DownloadResult.Downloaded(remoteBinaryURL, cliTempDst)
             }
 

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

@@ -4,7 +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.CoderSetupWizardState
+import com.coder.toolbox.views.state.WizardModel
 import com.coder.toolbox.views.state.WizardStep
 import com.jetbrains.toolbox.api.remoteDev.ProviderVisibilityState
 import com.jetbrains.toolbox.api.ui.actions.RunnableActionDescription
@@ -25,15 +25,18 @@ class CoderCliSetupWizardPage(
     onConnect: SuspendBiConsumer<CoderRestClient, CoderCLIManager>,
     onTokenRefreshed: (suspend (url: URL, oauthSessionCtx: CoderOAuthSessionContext) -> Unit)? = null
 ) : CoderPage(MutableStateFlow(context.i18n.ptrl("Setting up Coder")), false) {
+    val model: WizardModel = WizardModel()
+
     private val shouldAutoSetup = MutableStateFlow(initialAutoSetup)
     private val settingsAction = Action(context, "Settings") {
         context.ui.showUiPage(settingsPage)
     }
 
-    private val deploymentUrlStep = DeploymentUrlStep(context, visibilityState)
-    private val tokenStep = TokenStep(context)
+    private val deploymentUrlStep = DeploymentUrlStep(context, model, visibilityState)
+    private val tokenStep = TokenStep(context, model)
     private val connectStep = ConnectStep(
         context,
+        model,
         shouldAutoLogin = shouldAutoSetup,
         jumpToMainPageOnError = jumpToMainPageOnError,
         visibilityState,
@@ -50,11 +53,10 @@ class CoderCliSetupWizardPage(
     override val fields: MutableStateFlow<List<UiField>> = MutableStateFlow(emptyList())
     override val actionButtons: MutableStateFlow<List<RunnableActionDescription>> = MutableStateFlow(emptyList())
 
-
     override fun beforeShow() {
         stateCollectJob?.cancel()
         stateCollectJob = context.cs.launch {
-            CoderSetupWizardState.step.collect { step ->
+            model.step.collect { step ->
                 context.logger.info("Wizard step changed to $step")
                 displaySteps()
             }
@@ -63,7 +65,7 @@ class CoderCliSetupWizardPage(
     }
 
     private fun displaySteps() {
-        when (CoderSetupWizardState.currentStep()) {
+        when (model.currentStep()) {
             WizardStep.URL_REQUEST -> {
                 fields.update {
                     listOf(deploymentUrlStep.panel)
@@ -120,16 +122,19 @@ class CoderCliSetupWizardPage(
                 }
                 connectStep.onVisible()
             }
-
-            WizardStep.DONE -> {
-                context.logger.info("Closing the Setup Wizard")
-                stateCollectJob?.cancel()
-                context.ui.hideUiPage(this)
-                CoderSetupWizardState.goToFirstStep()
-            }
         }
     }
 
+    /**
+     * Cancels any in-flight work owned by this wizard. Called by the page router
+     * 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.
+     */
+    fun dispose() {
+        stateCollectJob?.cancel()
+        connectStep.dispose()
+    }
+
     override fun afterHide() {
         stateCollectJob?.cancel()
     }

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

@@ -142,8 +142,8 @@ class ConnectStep(
                 oauthSession?.let { session ->
                     onTokenRefreshed?.invoke(client.url, session)
                 }
-                model.clearFormData()
-                model.goToDone()
+                // 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()
             } catch (ex: CancellationException) {
                 if (ex.message != USER_HIT_THE_BACK_BUTTON) {
@@ -208,4 +208,14 @@ class ConnectStep(
             handleNavigation()
         }
     }
+
+    /**
+     * Cancels any in-flight connection without navigating. Used when the wizard
+     * itself is being torn down by an external trigger (e.g. a deep link to a
+     * different deployment).
+     */
+    fun dispose() {
+        signInJob?.cancel(CancellationException(USER_HIT_THE_BACK_BUTTON))
+        signInJob = null
+    }
 }