fix: OAuth2 authorization token revocation flow

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

@@ -8,6 +8,7 @@ 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
+import com.coder.toolbox.sdk.ex.OAuthTokenResponseException
 import com.coder.toolbox.sdk.v2.models.WorkspaceStatus
 import com.coder.toolbox.util.CoderProtocolHandler
 import com.coder.toolbox.util.DialogUi
@@ -158,7 +159,7 @@ class CoderRemoteProvider(
                         context.logger.info("wake-up from an OS sleep was detected")
                     } else {
                         context.logger.error(ex, "workspace polling error encountered")
-                        if (ex is APIResponseException && ex.isTokenExpired) {
+                        if (ex is APIResponseException && ex.isTokenExpired || ex is OAuthTokenResponseException) {
                             close()
                             context.envPageManager.showPluginEnvironmentsPage()
                             errorBuffer.add(ex)

src/main/kotlin/com/coder/toolbox/oauth/ClientRegistrationResponse.kt

@@ -43,5 +43,3 @@ data class ClientRegistrationErrorResponse(
         }
     }
 }
-
-class ClientRegistrationException(message: String) : Exception(message)

src/main/kotlin/com/coder/toolbox/oauth/OAuth2Client.kt

@@ -3,6 +3,8 @@ package com.coder.toolbox.oauth
 import com.coder.toolbox.CoderToolboxContext
 import com.coder.toolbox.sdk.CoderHttpClientBuilder
 import com.coder.toolbox.sdk.convertors.LoggingConverterFactory
+import com.coder.toolbox.sdk.ex.ClientRegistrationException
+import com.coder.toolbox.sdk.ex.OAuthTokenResponseException
 import com.coder.toolbox.views.state.CoderOAuthSessionContext
 import com.squareup.moshi.Moshi
 import okhttp3.Credentials
@@ -33,13 +35,9 @@ class OAuth2Client(private val context: CoderToolboxContext) {
         }
 
         val errorBody = response.errorBody()?.string()
-        val registrationError = errorBody?.let { ClientRegistrationErrorResponse.fromJson(it) }
-        val errorMessage = if (registrationError != null) {
-            "OAuth2 client registration failed: ${registrationError.toMessage()}"
-        } else {
-            "OAuth2 client registration failed with status ${response.code()}: ${response.message()}"
-        }
-        context.logger.error(errorMessage)
+        val registrationError =
+            errorBody?.let { ClientRegistrationErrorResponse.fromJson(it) }?.toMessage() ?: "${response.message()}"
+        val errorMessage = "OAuth2 client registration failed with status ${response.code()}: $registrationError"
         throw ClientRegistrationException(errorMessage)
     }
 
@@ -114,14 +112,9 @@ class OAuth2Client(private val context: CoderToolboxContext) {
         }
 
         val errorBody = response.errorBody()?.string()
-        val tokenError = errorBody?.let { OAuthTokenErrorResponse.fromJson(it) }
-        val errorMessage = if (tokenError != null) {
-            "Failed to $action: ${tokenError.toMessage()}"
-        } else {
-            "Failed to $action. Response code: ${response.code()} ${response.message()}"
-        }
-        context.logger.error(errorMessage)
-        throw Exception(errorMessage)
+        val tokenError = errorBody?.let { OAuthTokenErrorResponse.fromJson(it) }?.toMessage() ?: "${response.message()}"
+        val errorMessage = "Failed to $action. Response code: ${response.code()} $tokenError"
+        throw OAuthTokenResponseException(errorMessage)
     }
 
     private fun createAuthorizationService(): CoderAuthorizationApi {

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

@@ -9,6 +9,7 @@ import com.coder.toolbox.sdk.convertors.LoggingConverterFactory
 import com.coder.toolbox.sdk.convertors.OSConverter
 import com.coder.toolbox.sdk.convertors.UUIDConverter
 import com.coder.toolbox.sdk.ex.APIResponseException
+import com.coder.toolbox.sdk.interceptors.CODER_SESSION_TOKEN_HEADER_NAME
 import com.coder.toolbox.sdk.interceptors.Interceptors
 import com.coder.toolbox.sdk.v2.CoderV2RestFacade
 import com.coder.toolbox.sdk.v2.models.ApiErrorResponse
@@ -358,7 +359,7 @@ open class CoderRestClient(
             if (response.code() == HttpURLConnection.HTTP_UNAUTHORIZED && oauthContext.hasRefreshToken()) {
                 val tokenRefreshed = refreshMutex.withLock {
                     // Check if the token was already refreshed while we were waiting for the lock.
-                    if (response.raw().request.header("Authorization") != "Bearer ${oauthContext?.tokenResponse?.accessToken}") {
+                    if (response.raw().request.header(CODER_SESSION_TOKEN_HEADER_NAME) != oauthContext?.tokenResponse?.accessToken) {
                         return@withLock true
                     }
                     return@withLock try {
@@ -372,7 +373,7 @@ open class CoderRestClient(
                         true
                     } catch (e: Exception) {
                         context.logger.error(e, "Failed to refresh access token")
-                        false
+                        throw e
                     }
                 }
                 if (tokenRefreshed) {
@@ -425,7 +426,7 @@ open class CoderRestClient(
             }
         } catch (ex: Exception) {
             context.logger.error(ex, "Failed to execute refresh command")
-            false
+            throw ex
         }
     }
 

src/main/kotlin/com/coder/toolbox/sdk/ex/OAuth2Exceptions.kt

@@ -0,0 +1,6 @@
+package com.coder.toolbox.sdk.ex
+
+sealed class OAuth2ErrorException(message: String?) : Exception(message)
+
+class ClientRegistrationException(message: String) : OAuth2ErrorException(message)
+class OAuthTokenResponseException(message: String?) : OAuth2ErrorException(message)

src/main/kotlin/com/coder/toolbox/sdk/interceptors/Interceptors.kt

@@ -7,6 +7,8 @@ import com.coder.toolbox.util.getOS
 import okhttp3.Interceptor
 import java.net.URL
 
+const val CODER_SESSION_TOKEN_HEADER_NAME = "Coder-Session-Token"
+
 /**
  * Factory of okhttp interceptors
  */
@@ -19,7 +21,7 @@ object Interceptors {
         return Interceptor { chain ->
             chain.proceed(
                 chain.request().newBuilder()
-                    .addHeader("Coder-Session-Token", token)
+                    .addHeader(CODER_SESSION_TOKEN_HEADER_NAME, token)
                     .build()
             )
         }

src/main/kotlin/com/coder/toolbox/util/Dialogs.kt

@@ -11,10 +11,6 @@ import com.jetbrains.toolbox.api.ui.components.TextType
  */
 class DialogUi(private val context: CoderToolboxContext) {
 
-    suspend fun confirm(title: LocalizableString, description: LocalizableString): Boolean {
-        return context.ui.showOkCancelPopup(title, description, context.i18n.ptrl("Yes"), context.i18n.ptrl("No"))
-    }
-
     suspend fun ask(
         title: LocalizableString,
         description: LocalizableString,

src/main/kotlin/com/coder/toolbox/util/StateFlowExtensions.kt

@@ -5,11 +5,6 @@ import kotlinx.coroutines.flow.first
 import kotlinx.coroutines.withTimeoutOrNull
 import kotlin.time.Duration
 
-/**
- * Suspends the coroutine until first true value is received.
- */
-suspend fun StateFlow<Boolean>.waitForTrue() = this.first { it }
-
 /**
  * Suspends the coroutine until first false value is received.
  */