Fix OAuth callback transition to connect wizard

Route OAuth callbacks through a pending OAuth connection instead of reading
wizard model state directly, then create and schedule a fresh connect-step
wizard after token exchange.

Author: fioan89

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

@@ -29,6 +29,7 @@ import com.coder.toolbox.views.SuspendBiConsumer
 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.PendingOAuthConnection
 import com.coder.toolbox.views.state.toSessionContext
 import com.jetbrains.toolbox.api.core.ui.icons.SvgIcon
 import com.jetbrains.toolbox.api.core.ui.icons.SvgIcon.IconType
@@ -417,15 +418,17 @@ class CoderRemoteProvider(
             )
         }
 
-        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(
+        if (!router.hasActiveWizard) {
+            return context.logAndShowError(
+                FAILED_TO_HANDLE_OAUTH2_TITLE,
+                "OAuth2 callback arrived but the setup wizard is no longer active"
+            )
+        }
+        val pendingOAuthConnection = router.pendingOAuthConnection ?: return context.logAndShowError(
             FAILED_TO_HANDLE_OAUTH2_TITLE,
             "OAuth2 callback arrived but no OAuth session was started"
         )
-        params["state"]?.takeIf { it == activeOAuthSession.state }
+        params["state"]?.takeIf { it == pendingOAuthConnection.session.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"
@@ -443,19 +446,29 @@ class CoderRemoteProvider(
             )
             return
         }
-        exchangeOAuthCodeForToken(code, activeOAuthSession, activeWizard)
+        exchangeOAuthCodeForToken(code, pendingOAuthConnection)
     }
 
     private suspend fun exchangeOAuthCodeForToken(
         code: String,
-        oauthSessionContext: CoderOAuthSessionContext,
-        wizard: CoderSetupWizardPage,
+        pendingOAuthConnection: PendingOAuthConnection,
     ) {
         try {
             context.logger.info("Handling OAuth callback...")
 
+            val oauthSessionContext = pendingOAuthConnection.session
             val tokenResponse = OAuth2Client(context).exchangeCode(oauthSessionContext, code)
-            wizard.advanceToConnectWithOAuth(oauthSessionContext.copy(tokenResponse = tokenResponse))
+            val wizard = CoderSetupWizardPage.connectStep(
+                context, settingsPage, visibilityState,
+                url = pendingOAuthConnection.url,
+                credentials = Credentials.OAuth(oauthSessionContext.copy(tokenResponse = tokenResponse)),
+                onConnect = onConnect,
+                onTokenRefreshed = ::onTokenRefreshed,
+            )
+            router.navigate(wizard)
+
+            context.envPageManager.showPluginEnvironmentsPage(true)
+            context.ui.showUiPage(wizard)
         } catch (e: Exception) {
             context.logAndShowError("OAuth Error", "Exception during token exchange: ${e.message}", e)
         }
@@ -523,7 +536,7 @@ class CoderRemoteProvider(
      */
     override fun getOverrideUiPage(): UiPage? {
         // Show the setup wizard if one is already scheduled.
-        router.activeWizard?.let { return it }
+        router.activePage?.let { return it }
 
         // Let the default workspace UI render if the HTTP client is initialized.
         if (client != null) return null

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

@@ -5,6 +5,7 @@ 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.Credentials
+import com.coder.toolbox.views.state.PendingOAuthConnection
 import com.coder.toolbox.views.state.WizardModel
 import com.coder.toolbox.views.state.WizardStep
 import com.jetbrains.toolbox.api.remoteDev.ProviderVisibilityState
@@ -118,9 +119,10 @@ class CoderSetupWizardPage private constructor(
         }
     }
 
-    fun advanceToConnectWithOAuth(oauthSession: CoderOAuthSessionContext) {
-        model.oauthSession = oauthSession
-        model.goTo(WizardStep.CONNECT)
+    fun pendingOAuthConnection(): PendingOAuthConnection? {
+        val url = model.url ?: return null
+        val oauthSession = model.oauthSession ?: return null
+        return PendingOAuthConnection(url, oauthSession)
     }
 
     private fun navigateBackFromConnect() {

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

@@ -20,9 +20,18 @@ sealed interface PageRoute {
 class PageRouter {
     private val route = MutableStateFlow<PageRoute>(PageRoute.None)
 
-    val activeWizard: CoderSetupWizardPage?
+    private val activeWizard: CoderSetupWizardPage?
         get() = (route.value as? PageRoute.Wizard)?.page
 
+    val activePage: CoderPage?
+        get() = activeWizard
+
+    val hasActiveWizard: Boolean
+        get() = activeWizard != null
+
+    val pendingOAuthConnection: PendingOAuthConnection?
+        get() = activeWizard?.pendingOAuthConnection()
+
     /**
      * Returns the page already on this route, or builds a new one and
      * registers it.
@@ -47,4 +56,4 @@ class PageRouter {
         activeWizard?.dispose()
         route.value = PageRoute.None
     }
-}
+}

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

@@ -64,4 +64,9 @@ sealed interface Credentials {
     object MTls : Credentials
     data class Token(val value: String) : Credentials
     data class OAuth(val session: CoderOAuthSessionContext) : Credentials
-}
+}
+
+data class PendingOAuthConnection(
+    val url: URL,
+    val session: CoderOAuthSessionContext
+)

src/test/kotlin/com/coder/toolbox/views/state/PageRouterTest.kt

@@ -0,0 +1,53 @@
+package com.coder.toolbox.views.state
+
+import com.coder.toolbox.oauth.TokenEndpointAuthMethod
+import com.coder.toolbox.views.CoderSetupWizardPage
+import io.mockk.every
+import io.mockk.mockk
+import java.net.URI
+import kotlin.test.Test
+import kotlin.test.assertEquals
+import kotlin.test.assertNull
+
+class PageRouterTest {
+
+    @Test
+    fun `given no active wizard when pending OAuth connection is requested then null is returned`() {
+        // given
+        val router = PageRouter()
+
+        // when
+        val pendingOAuthConnection = router.pendingOAuthConnection
+
+        // then
+        assertNull(pendingOAuthConnection)
+    }
+
+    @Test
+    fun `given active wizard when pending OAuth connection is requested then wizard connection is returned`() {
+        // given
+        val pendingConnection = PendingOAuthConnection(
+            url = URI("https://coder.example.com").toURL(),
+            session = oauthSession()
+        )
+        val wizard = mockk<CoderSetupWizardPage>(relaxed = true)
+        every { wizard.pendingOAuthConnection() } returns pendingConnection
+        val router = PageRouter()
+        router.navigate(wizard)
+
+        // when
+        val pendingOAuthConnection = router.pendingOAuthConnection
+
+        // then
+        assertEquals(pendingConnection, pendingOAuthConnection)
+    }
+
+    private fun oauthSession(): CoderOAuthSessionContext = CoderOAuthSessionContext(
+        clientId = "client-id",
+        clientSecret = "client-secret",
+        tokenCodeVerifier = "verifier",
+        state = "state",
+        tokenEndpoint = "https://coder.example.com/oauth/token",
+        tokenAuthMethod = TokenEndpointAuthMethod.CLIENT_SECRET_BASIC
+    )
+}