feat(OAuth2): implemet support for OAuth2 (1)

CHANGELOG.md

@@ -2,6 +2,14 @@
 
 ## Unreleased
 
+### Added
+
+- support for OAuth2
+
+### Changed
+
+- redesigned the Settings page, all the options are now grouped in a couple of sections for easy navigation
+
 ## 0.8.6 - 2026-03-05
 
 ### Changed

build.gradle.kts

@@ -64,11 +64,11 @@ dependencies {
     implementation(libs.retrofit)
     implementation(libs.retrofit.moshi)
     implementation(libs.bundles.bouncycastle)
+
     testImplementation(kotlin("test"))
     testImplementation(libs.coroutines.test)
     testImplementation(libs.mokk)
     testImplementation(libs.bundles.toolbox.plugin.api)
-    testImplementation(libs.coroutines.test)
 }
 
 val extension = ExtensionJson(
@@ -204,21 +204,13 @@ tasks.register("cleanAll", Delete::class.java) {
 
 private fun getPluginInstallDir(): Path {
     val userHome = System.getProperty("user.home").let { Path.of(it) }
-    val toolboxCachesDir = when {
+    val pluginsDir = when {
         SystemInfoRt.isWindows -> System.getenv("LOCALAPPDATA")?.let { Path.of(it) } ?: (userHome / "AppData" / "Local")
         // currently this is the location that TBA uses on Linux
         SystemInfoRt.isLinux -> System.getenv("XDG_DATA_HOME")?.let { Path.of(it) } ?: (userHome / ".local" / "share")
         SystemInfoRt.isMac -> userHome / "Library" / "Caches"
         else -> error("Unknown os")
-    } / "JetBrains" / "Toolbox"
-
-    val pluginsDir = when {
-        SystemInfoRt.isWindows ||
-                SystemInfoRt.isLinux ||
-                SystemInfoRt.isMac -> toolboxCachesDir
-
-        else -> error("Unknown os")
-    } / "plugins"
+    } / "JetBrains" / "Toolbox" / "plugins"
 
     return pluginsDir / extension.id
 }

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

@@ -3,6 +3,8 @@ package com.coder.toolbox
 import com.coder.toolbox.browser.browse
 import com.coder.toolbox.cli.CoderCLIManager
 import com.coder.toolbox.feed.IdeFeedManager
+import com.coder.toolbox.oauth.OAuth2Service
+import com.coder.toolbox.oauth.OAuthTokenResponse
 import com.coder.toolbox.plugin.PluginManager
 import com.coder.toolbox.sdk.CoderRestClient
 import com.coder.toolbox.sdk.ex.APIResponseException
@@ -17,15 +19,16 @@ import com.coder.toolbox.util.toURL
 import com.coder.toolbox.util.token
 import com.coder.toolbox.util.url
 import com.coder.toolbox.util.validateStrictWebUrl
-import com.coder.toolbox.util.waitForTrue
 import com.coder.toolbox.util.withPath
 import com.coder.toolbox.views.Action
 import com.coder.toolbox.views.CoderCliSetupWizardPage
 import com.coder.toolbox.views.CoderDelimiter
 import com.coder.toolbox.views.CoderSettingsPage
 import com.coder.toolbox.views.NewEnvironmentPage
-import com.coder.toolbox.views.state.CoderCliSetupContext
-import com.coder.toolbox.views.state.CoderCliSetupWizardState
+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.WizardStep
 import com.jetbrains.toolbox.api.core.ui.icons.SvgIcon
 import com.jetbrains.toolbox.api.core.ui.icons.SvgIcon.IconType
@@ -48,7 +51,6 @@ import kotlinx.coroutines.selects.onTimeout
 import kotlinx.coroutines.selects.select
 import java.net.URI
 import java.net.URL
-import java.util.concurrent.atomic.AtomicBoolean
 import kotlin.coroutines.cancellation.CancellationException
 import kotlin.time.Duration.Companion.seconds
 import kotlin.time.TimeSource
@@ -79,7 +81,7 @@ class CoderRemoteProvider(
     private var firstRun = true
 
     private val isInitialized: MutableStateFlow<Boolean> = MutableStateFlow(false)
-    private val isHandlingUri: AtomicBoolean = AtomicBoolean(false)
+
     private val coderHeaderPage = NewEnvironmentPage(context.i18n.pnotr(context.deploymentUrl.toString()))
     private val settingsPage: CoderSettingsPage = CoderSettingsPage(context, triggerSshConfig) {
         client?.let { restClient ->
@@ -265,7 +267,6 @@ class CoderRemoteProvider(
         lastEnvironments.clear()
         environments.value = LoadableState.Value(emptyList())
         isInitialized.update { false }
-        CoderCliSetupWizardState.goToFirstStep()
         context.logger.info("Coder plugin is now closed")
     }
 
@@ -274,6 +275,7 @@ class CoderRemoteProvider(
             it.cancel()
             context.logger.info("Cancelled workspace poll job ${pollJob.toString()}")
         }
+        pollJob = null
         client?.let {
             it.close()
             context.logger.info("REST API client closed and resources released")
@@ -341,49 +343,53 @@ 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
+            }
+
             val params = uri.toQueryParameters()
             if (params.isEmpty()) {
                 // probably a plugin installation scenario
                 context.logAndShowInfo("URI will not be handled", "No query parameters were provided")
                 return
             }
-            isHandlingUri.set(true)
-            // this switches to the main plugin screen, even
-            // if last opened provider was not Coder
-            context.envPageManager.showPluginEnvironmentsPage()
-            coderHeaderPage.isBusy.update { true }
             context.logger.info("Handling $uri...")
             val newUrl = resolveDeploymentUrl(params)?.toURL() ?: return
             val newToken = if (context.settingsStore.requiresMTlsAuth) null else resolveToken(params) ?: return
+            coderHeaderPage.isBusy.update { true }
             if (sameUrl(newUrl, client?.url)) {
                 if (context.settingsStore.requiresTokenAuth) {
                     newToken?.let {
                         refreshSession(newUrl, it)
                     }
                 }
+                linkHandler.handle(params, newUrl, this.client!!, this.cli!!)
+                coderHeaderPage.isBusy.update { false }
             } else {
-                CoderCliSetupContext.apply {
+                // 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
                 }
-                CoderCliSetupWizardState.goToStep(WizardStep.CONNECT)
-                CoderCliSetupWizardPage(
-                    context, settingsPage, visibilityState,
-                    initialAutoSetup = true,
-                    jumpToMainPageOnError = true,
-                    connectSynchronously = true,
-                    onConnect = ::onConnect
-                ).apply {
-                    beforeShow()
-                }
+                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
+                    )
+                )
             }
-            // force the poll loop to run
-            triggerProviderVisible.send(true)
-            // wait for environments to be populated
-            isInitialized.waitForTrue()
-
-            linkHandler.handle(params, newUrl, this.client!!, this.cli!!)
-            coderHeaderPage.isBusy.update { false }
         } catch (ex: Exception) {
             val textError = if (ex is APIResponseException) {
                 if (!ex.reason.isNullOrBlank()) {
@@ -394,14 +400,43 @@ class CoderRemoteProvider(
                 "Error encountered while handling Coder URI",
                 textError ?: ""
             )
+            coderHeaderPage.isBusy.update { false }
             context.envPageManager.showPluginEnvironmentsPage()
         } finally {
-            coderHeaderPage.isBusy.update { false }
-            isHandlingUri.set(false)
             firstRun = false
         }
     }
 
+    private suspend fun handleOAuthUri(uri: URI) {
+        val params = uri.toQueryParameters()
+        val code = params["code"] ?: return context.logAndShowError(
+            "Failed to handle OAuth code",
+            "Server responded did not respond back with an access token"
+        )
+
+        params["state"]?.takeIf { it == CoderSetupWizardContext.oauthSession?.state }
+            ?: return context.logAndShowError(
+                "Failed to handle OAuth code",
+                "Server responded back with an invalid state that does not match the initial authorization state sent to the server"
+            )
+
+        exchangeOAuthCodeForToken(code, CoderSetupWizardContext.oauthSession!!)
+    }
+
+    private suspend fun exchangeOAuthCodeForToken(code: String, oauthSessionContext: CoderOAuthSessionContext) {
+        try {
+            context.logger.info("Handling OAuth callback...")
+
+            val tokenResponse = OAuth2Service(context).exchangeCode(oauthSessionContext, code)
+            oauthSessionContext.tokenResponse = tokenResponse
+
+            CoderSetupWizardState.goToStep(WizardStep.CONNECT)
+
+        } catch (e: Exception) {
+            context.logAndShowError("OAuth Error", "Exception during token exchange: ${e.message}", e)
+        }
+    }
+
     private suspend fun resolveDeploymentUrl(params: Map<String, String>): String? {
         val deploymentURL = params.url() ?: askUrl()
         if (deploymentURL.isNullOrBlank()) {
@@ -434,6 +469,7 @@ class CoderRemoteProvider(
             context,
             url,
             token,
+            null,
             PluginManager.pluginInfo.version,
         ).apply { initializeSession() }
         val newCli = CoderCLIManager(context, url).apply {
@@ -462,24 +498,40 @@ class CoderRemoteProvider(
      * list.
      */
     override fun getOverrideUiPage(): UiPage? {
-        if (isHandlingUri.get()) {
-            return null
-        }
         // 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 {
-                    CoderCliSetupContext.apply {
+                    val storedOAuthSession = context.secrets.oauthSessionFor(context.deploymentUrl.toString())
+                    CoderSetupWizardContext.apply {
                         url = context.deploymentUrl
-                        token = context.secrets.tokenFor(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
+                                )
+                            )
+                        }
                     }
-                    CoderCliSetupWizardState.goToStep(WizardStep.CONNECT)
+                    CoderSetupWizardState.goToStep(WizardStep.CONNECT)
                     return CoderCliSetupWizardPage(
                         context, settingsPage, visibilityState,
                         initialAutoSetup = true,
                         jumpToMainPageOnError = false,
-                        onConnect = ::onConnect
+                        onConnect = onConnect,
+                        onTokenRefreshed = ::onTokenRefreshed
                     )
                 } catch (ex: Exception) {
                     errorBuffer.add(ex)
@@ -489,8 +541,15 @@ class CoderRemoteProvider(
             }
 
             // Login flow.
+            CoderSetupWizardState.goToFirstStep()
             val setupWizardPage =
-                CoderCliSetupWizardPage(context, settingsPage, visibilityState, onConnect = ::onConnect)
+                CoderCliSetupWizardPage(
+                    context,
+                    settingsPage,
+                    visibilityState,
+                    onConnect = onConnect,
+                    onTokenRefreshed = ::onTokenRefreshed
+                )
             // We might have navigated here due to a polling error.
             errorBuffer.forEach {
                 setupWizardPage.notify("Error encountered", it)
@@ -508,14 +567,22 @@ class CoderRemoteProvider(
      */
     private fun shouldDoAutoSetup(): Boolean = firstRun && (canAutoLogin() || !settings.requiresTokenAuth)
 
-    fun canAutoLogin(): Boolean = !context.secrets.tokenFor(context.deploymentUrl).isNullOrBlank()
+    fun canAutoLogin(): Boolean = !context.secrets.apiTokenFor(context.deploymentUrl)
+        .isNullOrBlank() || context.secrets.oauthSessionFor(context.deploymentUrl.toString()) != null
+
+    private suspend fun onTokenRefreshed(url: URL, oauthSessionCtx: CoderOAuthSessionContext) {
+        oauthSessionCtx.tokenResponse?.accessToken?.let { cli?.login(it) }
+        context.secrets.storeOAuthFor(url.toString(), oauthSessionCtx)
+    }
 
-    private fun onConnect(client: CoderRestClient, cli: CoderCLIManager) {
+    private val onConnect: SuspendBiConsumer<CoderRestClient, CoderCLIManager> = SuspendBiConsumer { client, cli ->
         // Store the URL and token for use next time.
         close()
         context.settingsStore.updateLastUsedUrl(client.url)
         if (context.settingsStore.requiresTokenAuth) {
-            context.secrets.storeTokenFor(client.url, client.token ?: "")
+            if (client.token != null) {
+                context.secrets.storeApiTokenFor(client.url, client.token)
+            }
             context.logger.info("Deployment URL and token were stored and will be available for automatic connection")
         } else {
             context.logger.info("Deployment URL was stored and will be available for automatic connection")
@@ -538,6 +605,27 @@ class CoderRemoteProvider(
         context.logger.info("Workspace poll job with name ${pollJob.toString()} was created")
     }
 
+    /**
+     * Returns a [SuspendBiConsumer] that handles the given link parameters.
+     * Runs in a background coroutine so it doesn't block the connect step's
+     * post-connection flow.
+     */
+    private fun deferredLinkHandler(
+        params: Map<String, String>,
+        deploymentUrl: URL,
+    ): SuspendBiConsumer<CoderRestClient, CoderCLIManager> = SuspendBiConsumer { client, cli ->
+        context.cs.launch(CoroutineName("Deferred Link Handler")) {
+            try {
+                linkHandler.handle(params, deploymentUrl, client, cli)
+            } catch (ex: Exception) {
+                context.logAndShowError(
+                    "Error handling deferred link",
+                    ex.message ?: ""
+                )
+            }
+        }
+    }
+
     private fun MutableStateFlow<LoadableState<List<CoderRemoteEnvironment>>>.showLoadingMessage() {
         this.update {
             LoadableState.Loading

src/main/kotlin/com/coder/toolbox/cli/CoderCLIManager.kt

@@ -12,14 +12,11 @@ import com.coder.toolbox.cli.gpg.GPGVerifier
 import com.coder.toolbox.cli.gpg.VerificationResult
 import com.coder.toolbox.cli.gpg.VerificationResult.Failed
 import com.coder.toolbox.cli.gpg.VerificationResult.Invalid
-import com.coder.toolbox.plugin.PluginManager
 import com.coder.toolbox.sdk.CoderHttpClientBuilder
-import com.coder.toolbox.sdk.interceptors.Interceptors
 import com.coder.toolbox.sdk.v2.models.Workspace
 import com.coder.toolbox.sdk.v2.models.WorkspaceAgent
 import com.coder.toolbox.settings.SignatureFallbackStrategy.ALLOW
 import com.coder.toolbox.util.InvalidVersionException
-import com.coder.toolbox.util.ReloadableTlsContext
 import com.coder.toolbox.util.SemVer
 import com.coder.toolbox.util.escape
 import com.coder.toolbox.util.escapeSubcommand
@@ -148,15 +145,8 @@ class CoderCLIManager(
     val coderConfigPath: Path = context.settingsStore.dataDir(deploymentURL).resolve("config")
 
     private fun createDownloadService(): CoderDownloadService {
-        val interceptors = buildList {
-            add((Interceptors.userAgent(PluginManager.pluginInfo.version)))
-            add(Interceptors.logging(context))
-        }
-        val okHttpClient = CoderHttpClientBuilder.build(
-            context,
-            interceptors,
-            ReloadableTlsContext(context.settingsStore.readOnly().tls)
-        )
+
+        val okHttpClient = CoderHttpClientBuilder.default(context)
 
         val retrofit = Retrofit.Builder()
             .baseUrl(deploymentURL.toString())