fix(amazonq): clear out cached supplier if connection is invalidated …

plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/toolwindow/AmazonQToolWindowFactory.kt

@@ -83,7 +83,7 @@ class AmazonQToolWindowFactory : ToolWindowFactory, DumbAware {
         project.messageBus.connect(toolWindow.disposable).subscribe(
             BearerTokenProviderListener.TOPIC,
             object : BearerTokenProviderListener {
-                override fun onChange(providerId: String, newScopes: List<String>?) {
+                override fun onProviderChange(providerId: String, newScopes: List<String>?) {
                     if (ToolkitConnectionManager.getInstance(project).connectionStateForFeature(QConnection.getInstance()) == BearerTokenAuthState.AUTHORIZED) {
                         preparePanelContent(project, qPanel)
                     }

plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeScan/CodeScanChatApp.kt

@@ -125,7 +125,7 @@ class CodeScanChatApp(private val scope: CoroutineScope) : AmazonQApp {
         ApplicationManager.getApplication().messageBus.connect(this).subscribe(
             BearerTokenProviderListener.TOPIC,
             object : BearerTokenProviderListener {
-                override fun onChange(providerId: String, newScopes: List<String>?) {
+                override fun onProviderChange(providerId: String, newScopes: List<String>?) {
                     val qProvider = getQTokenProvider(context.project)
                     val isQ = qProvider?.id == providerId
                     val isAuthorized = qProvider?.state() == BearerTokenAuthState.AUTHORIZED

plugins/amazonq/codetransform/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codemodernizer/CodeTransformChatApp.kt

@@ -149,7 +149,7 @@ class CodeTransformChatApp : AmazonQApp {
         ApplicationManager.getApplication().messageBus.connect(this).subscribe(
             BearerTokenProviderListener.TOPIC,
             object : BearerTokenProviderListener {
-                override fun onChange(providerId: String, newScopes: List<String>?) {
+                override fun onProviderChange(providerId: String, newScopes: List<String>?) {
                     val qProvider = getQTokenProvider(context.project)
                     val isQ = qProvider?.id == providerId
                     val isAuthorized = qProvider?.state() == BearerTokenAuthState.AUTHORIZED

plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/status/CodeWhispererStatusBarWidget.kt

@@ -56,7 +56,7 @@ class CodeWhispererStatusBarWidget(project: Project) :
         ApplicationManager.getApplication().messageBus.connect(this).subscribe(
             BearerTokenProviderListener.TOPIC,
             object : BearerTokenProviderListener {
-                override fun onChange(providerId: String, newScopes: List<String>?) {
+                override fun onProviderChange(providerId: String, newScopes: List<String>?) {
                     statusBar.updateWidget(ID)
                 }
             }

plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/auth/DefaultAuthCredentialsService.kt

@@ -19,7 +19,6 @@
 import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
 import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManagerListener
 import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection
-import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.BearerTokenAuthState
 import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.BearerTokenProvider
 import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.BearerTokenProviderListener
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.AmazonQLspService
@@ -43,15 +42,14 @@
     private val project: Project,
     private val encryptionManager: JwtEncryptionManager,
     private val cs: CoroutineScope,
-) : AuthCredentialsService,
-    BearerTokenProviderListener,
+) : BearerTokenProviderListener,
     ToolkitConnectionManagerListener,
     QRegionProfileSelectedListener,
     Disposable {
 
     private val scheduler: ScheduledExecutorService = AppExecutorUtil.getAppScheduledExecutorService()
-    private var tokenSyncTask: ScheduledFuture<*>? = null
-    private val tokenSyncIntervalMinutes = 5L
+    private var tokenRefreshTask: ScheduledFuture<*>? = null
+    private val tokenRefreshInterval = 5L
 
     init {
         project.messageBus.connect(this).apply {
@@ -67,49 +65,37 @@
                 }
         }
 
-        // Start periodic token sync
-        startPeriodicTokenSync()
+        // Start periodic token refresh
+        startPeriodicTokenRefresh()
     }
 
-    private fun startPeriodicTokenSync() {
-        tokenSyncTask = scheduler.scheduleWithFixedDelay(
+    // TODO: we really only need a single application-wide instance of this
+    private fun startPeriodicTokenRefresh() {
+        tokenRefreshTask = scheduler.scheduleWithFixedDelay(
             {
                 try {
                     if (isQConnected(project)) {
-                        if (isQExpired(project)) {
-                            val manager = ToolkitConnectionManager.getInstance(project)
-                            val connection = manager.activeConnectionForFeature(QConnection.getInstance()) ?: return@scheduleWithFixedDelay
-
-                            // Try to refresh the token if it's in NEEDS_REFRESH state
-                            val tokenProvider = (connection.getConnectionSettings() as? TokenConnectionSettings)
-                                ?.tokenProvider
-                                ?.delegate
-                                ?.let { it as? BearerTokenProvider } ?: return@scheduleWithFixedDelay
-
-                            if (tokenProvider.state() == BearerTokenAuthState.NEEDS_REFRESH) {
-                                try {
-                                    tokenProvider.resolveToken()
-                                    // Now that the token is refreshed, update it in Flare
-                                    updateTokenFromActiveConnection()
-                                } catch (e: Exception) {
-                                    LOG.warn(e) { "Failed to refresh bearer token" }
-                                }
-                            }
-                        } else {
-                            updateTokenFromActiveConnection()
-                        }
+                        val manager = ToolkitConnectionManager.getInstance(project)
+                        val connection = manager.activeConnectionForFeature(QConnection.getInstance()) ?: return@scheduleWithFixedDelay
+
+                        // periodically poll token to trigger a background refresh if needed
+                        val tokenProvider = (connection.getConnectionSettings() as? TokenConnectionSettings)
+                            ?.tokenProvider
+                            ?.delegate
+                            ?.let { it as? BearerTokenProvider } ?: return@scheduleWithFixedDelay
+                        tokenProvider.resolveToken()
                     }
                 } catch (e: Exception) {
-                    LOG.warn(e) { "Failed to sync bearer token to Flare" }
+                    LOG.warn(e) { "Failed to refresh bearer token" }
                 }
             },
-            tokenSyncIntervalMinutes,
-            tokenSyncIntervalMinutes,
+            tokenRefreshInterval,
+            tokenRefreshInterval,
             TimeUnit.MINUTES
         )
     }
 
-    override fun updateTokenCredentials(connection: ToolkitConnection, encrypted: Boolean): CompletableFuture<ResponseMessage> {
+    fun updateTokenCredentials(connection: ToolkitConnection, encrypted: Boolean): CompletableFuture<ResponseMessage> {
         val payload = try {
             createUpdateCredentialsPayload(connection, encrypted)
         } catch (e: Exception) {
@@ -129,18 +115,26 @@
         }.asCompletableFuture()
     }
 
-    override fun deleteTokenCredentials() {
+    fun deleteTokenCredentials() {
         cs.launch {
             AmazonQLspService.executeAsyncIfRunning(project) { server ->
                 server.deleteTokenCredentials()
             }
         }
     }
 
-    override fun onChange(providerId: String, newScopes: List<String>?) {
+    override fun onProviderChange(providerId: String, newScopes: List<String>?) {
         updateTokenFromActiveConnection()
     }
 
+    override fun onTokenModified(providerId: String) {
+        updateTokenFromActiveConnection()
+    }
+
+    override fun invalidate(providerId: String) {
+        deleteTokenCredentials()
+    }
+
     override fun activeConnectionChanged(newConnection: ToolkitConnection?) {
         val qConnection = ToolkitConnectionManager.getInstance(project)
             .activeConnectionForFeature(QConnection.getInstance())
@@ -161,10 +155,6 @@
     private fun updateTokenFromConnection(connection: ToolkitConnection): CompletableFuture<ResponseMessage> =
         updateTokenCredentials(connection, true)
 
-    override fun invalidate(providerId: String) {
-        deleteTokenCredentials()
-    }
-
     private fun createUpdateCredentialsPayload(connection: ToolkitConnection, encrypted: Boolean): UpdateCredentialsPayload {
         val token = (connection.getConnectionSettings() as? TokenConnectionSettings)
             ?.tokenProvider
@@ -212,8 +202,8 @@
     }
 
     override fun dispose() {
-        tokenSyncTask?.cancel(false)
-        tokenSyncTask = null
+        tokenRefreshTask?.cancel(false)
+        tokenRefreshTask = null
     }
 
     companion object {

plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/profile/QRegionProfileManager.kt

@@ -114,7 +114,12 @@
                 connectionIdToProfileCount[connection.id] = it.size
             } ?: error("You don't have access to the resource")
         } catch (e: Exception) {
-            LOG.warn(e) { "Failed to list region profiles: ${e.message}" }
+            if (e is AccessDeniedException) {
+                LOG.warn { "Failed to list region profiles: ${e.message}" }
+            } else {
+                LOG.warn(e) { "Failed to list region profiles" }
+            }
+
             throw e
         }
     }

plugins/amazonq/shared/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonq/lsp/auth/DefaultAuthCredentialsServiceTest.kt

@@ -180,7 +180,7 @@ class DefaultAuthCredentialsServiceTest {
         sut = DefaultAuthCredentialsService(project, mockEncryptionManager, this)
         setupMockConnectionManager("updated-token")
 
-        sut.onChange("providerId", listOf("new-scope"))
+        sut.onProviderChange("providerId", listOf("new-scope"))
 
         advanceUntilIdle()
         verify(exactly = 1) { mockLanguageServer.updateTokenCredentials(any()) }

plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/AwsClientManager.kt

@@ -51,7 +51,7 @@ open class AwsClientManager : ToolkitClientManager(), Disposable {
         busConnection.subscribe(
             BearerTokenProviderListener.TOPIC,
             object : BearerTokenProviderListener {
-                override fun onChange(providerId: String, newScopes: List<String>?) {
+                override fun onProviderChange(providerId: String, newScopes: List<String>?) {
                     invalidateSdks(providerId)
                 }
 

plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/AwsResourceCache.kt

@@ -165,7 +165,7 @@ class DefaultAwsResourceCache(
             subscribe(
                 BearerTokenProviderListener.TOPIC,
                 object : BearerTokenProviderListener {
-                    override fun onChange(providerId: String, newScopes: List<String>?) {
+                    override fun onProviderChange(providerId: String, newScopes: List<String>?) {
                         clearByCredential(providerId)
                     }
                 }

plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/credentials/CredentialManager.kt

@@ -64,7 +64,7 @@ class DefaultCredentialManager : CredentialManager(), Disposable {
         ApplicationManager.getApplication().messageBus.connect(this).subscribe(
             BearerTokenProviderListener.TOPIC,
             object : BearerTokenProviderListener {
-                override fun onChange(providerId: String, newScopes: List<String>?) {
+                override fun onProviderChange(providerId: String, newScopes: List<String>?) {
                     modifyDependentProviders(providerId)
                 }
 

plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/credentials/DefaultToolkitAuthManager.kt

@@ -61,7 +61,7 @@
                                 if (it && connection is Disposable) {
                                     // don't invalidate because we kill the token we just retrieved
                                     ApplicationManager.getApplication().messageBus.syncPublisher(BearerTokenProviderListener.TOPIC)
-                                        .onChange(connection.id)
+                                        .onProviderChange(connection.id)
                                     Disposer.dispose(connection)
                                 }
                             }
@@ -147,7 +147,7 @@
             (existOldConn.id == newConnection.id).also { isDuplicate ->
                 if (isDuplicate && existOldConn is Disposable) {
                     ApplicationManager.getApplication().messageBus.syncPublisher(BearerTokenProviderListener.TOPIC)
-                        .onChange(existOldConn.id, newConnection.scopes)
+                        .onProviderChange(existOldConn.id, newConnection.scopes)
                     Disposer.dispose(existOldConn)
                 }
             }
@@ -157,7 +157,7 @@
             (existOldConn.id == newConnection.id).also { isDuplicate ->
                 if (isDuplicate && existOldConn is Disposable) {
                     ApplicationManager.getApplication().messageBus.syncPublisher(BearerTokenProviderListener.TOPIC)
-                        .onChange(existOldConn.id, newConnection.scopes)
+                        .onProviderChange(existOldConn.id, newConnection.scopes)
                     Disposer.dispose(existOldConn)
                 }
             }

plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/credentials/sso/SsoAccessTokenProvider.kt

@@ -194,7 +194,7 @@ class SsoAccessTokenProvider(
 
     @Deprecated("Device authorization grant flow is deprecated")
     private fun registerDAGClient(): ClientRegistration {
-        loadDagClientRegistration(SourceOfLoadRegistration.REGISTER_CLIENT.toString())?.let {
+        loadDagClientRegistration(SourceOfLoadRegistration.REGISTER_CLIENT)?.let {
             return it
         }
 
@@ -235,7 +235,7 @@ class SsoAccessTokenProvider(
     }
 
     private fun registerPkceClient(): PKCEClientRegistration {
-        loadPkceClientRegistration(SourceOfLoadRegistration.REGISTER_CLIENT.toString())?.let {
+        loadPkceClientRegistration(SourceOfLoadRegistration.REGISTER_CLIENT)?.let {
             return it
         }
 
@@ -431,8 +431,8 @@ class SsoAccessTokenProvider(
         stageName = RefreshCredentialStage.LOAD_REGISTRATION
         val registration = try {
             when (currentToken) {
-                is DeviceAuthorizationGrantToken -> loadDagClientRegistration(SourceOfLoadRegistration.REFRESH_TOKEN.toString())
-                is PKCEAuthorizationGrantToken -> loadPkceClientRegistration(SourceOfLoadRegistration.REFRESH_TOKEN.toString())
+                is DeviceAuthorizationGrantToken -> loadDagClientRegistration(SourceOfLoadRegistration.REFRESH_TOKEN)
+                is PKCEAuthorizationGrantToken -> loadPkceClientRegistration(SourceOfLoadRegistration.REFRESH_TOKEN)
             }
         } catch (e: Exception) {
             val message = e.message ?: "$stageName: ${e::class.java.name}"
@@ -519,13 +519,13 @@ class SsoAccessTokenProvider(
         SAVE_TOKEN,
     }
 
-    private fun loadDagClientRegistration(source: String): ClientRegistration? =
-        cache.loadClientRegistration(dagClientRegistrationCacheKey, source)?.let {
+    private fun loadDagClientRegistration(source: SourceOfLoadRegistration): ClientRegistration? =
+        cache.loadClientRegistration(dagClientRegistrationCacheKey, source.toString())?.let {
             return it
         }
 
-    private fun loadPkceClientRegistration(source: String): PKCEClientRegistration? =
-        cache.loadClientRegistration(pkceClientRegistrationCacheKey, source)?.let {
+    private fun loadPkceClientRegistration(source: SourceOfLoadRegistration): PKCEClientRegistration? =
+        cache.loadClientRegistration(pkceClientRegistrationCacheKey, source.toString())?.let {
             return it as PKCEClientRegistration
         }