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

when a user reauthenticates while reusing a connection, the cached supplier can incorrectly retain the last token instead of immediately returning the new one

additionally, the cached supplier can infinitely attempt token refresh since we are retaining invalid tokens instead of destroying them

Author: rli

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/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonq/AmazonQTestBase.kt

@@ -17,17 +17,14 @@ import org.mockito.kotlin.spy
 import org.mockito.kotlin.whenever
 import software.aws.toolkits.core.TokenConnectionSettings
 import software.aws.toolkits.core.credentials.ToolkitBearerTokenProvider
-import software.aws.toolkits.core.utils.test.aString
 import software.aws.toolkits.jetbrains.core.credentials.AwsBearerTokenConnection
 import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
-import software.aws.toolkits.jetbrains.core.credentials.sso.DeviceAuthorizationGrantToken
 import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.BearerTokenProvider
 import software.aws.toolkits.jetbrains.services.amazonq.clients.AmazonQStreamingClient
 import software.aws.toolkits.jetbrains.utils.rules.CodeInsightTestFixtureRule
 import software.aws.toolkits.jetbrains.utils.rules.HeavyJavaCodeInsightTestFixtureRule
 import software.aws.toolkits.jetbrains.utils.rules.JavaCodeInsightTestFixtureRule
 import software.aws.toolkits.jetbrains.utils.rules.addModule
-import java.time.Instant
 
 open class AmazonQTestBase(
     @Rule @JvmField
@@ -47,11 +44,7 @@ open class AmazonQTestBase(
         project = projectRule.project
         toolkitConnectionManager = spy(ToolkitConnectionManager.getInstance(project))
 
-        val accessToken = DeviceAuthorizationGrantToken(aString(), aString(), aString(), aString(), Instant.MAX, Instant.now())
-
-        val provider = mock<BearerTokenProvider> {
-            doReturn(accessToken).whenever(it).refresh()
-        }
+        val provider = mock<BearerTokenProvider>()
 
         val mockBearerProvider = mock<ToolkitBearerTokenProvider> {
             doReturn(provider).whenever(it).delegate

plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/FeatureDevTestBase.kt

@@ -28,10 +28,8 @@ import software.amazon.awssdk.services.codewhispererruntime.model.SendTelemetryE
 import software.amazon.awssdk.services.codewhispererruntime.model.StartTaskAssistCodeGenerationResponse
 import software.aws.toolkits.core.TokenConnectionSettings
 import software.aws.toolkits.core.credentials.ToolkitBearerTokenProvider
-import software.aws.toolkits.core.utils.test.aString
 import software.aws.toolkits.jetbrains.core.credentials.AwsBearerTokenConnection
 import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
-import software.aws.toolkits.jetbrains.core.credentials.sso.DeviceAuthorizationGrantToken
 import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.BearerTokenProvider
 import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.clients.FeatureDevClient
 import software.aws.toolkits.jetbrains.services.amazonqFeatureDev.session.CodeGenerationStreamResult
@@ -41,7 +39,6 @@ import software.aws.toolkits.jetbrains.utils.rules.HeavyJavaCodeInsightTestFixtu
 import software.aws.toolkits.jetbrains.utils.rules.JavaCodeInsightTestFixtureRule
 import software.aws.toolkits.jetbrains.utils.rules.addModule
 import java.io.File
-import java.time.Instant
 
 open class FeatureDevTestBase(
     @Rule @JvmField
@@ -164,11 +161,7 @@ open class FeatureDevTestBase(
     open fun setup() {
         project = projectRule.project
         toolkitConnectionManager = spy(ToolkitConnectionManager.getInstance(project))
-        val accessToken = DeviceAuthorizationGrantToken(aString(), aString(), aString(), aString(), Instant.MAX, Instant.now())
-        val provider =
-            mock<BearerTokenProvider> {
-                doReturn(accessToken).whenever(it).refresh()
-            }
+        val provider = mock<BearerTokenProvider>()
         val mockBearerProvider =
             mock<ToolkitBearerTokenProvider> {
                 doReturn(provider).whenever(it).delegate

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/codetransform/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codemodernizer/utils/CodeTransformApiUtils.kt

@@ -87,11 +87,6 @@ suspend fun JobId.pollTransformationStatusAndPlan(
     var transformationPlan: TransformationPlan? = null
     var didSleepOnce = false
     var hasSeenTransforming = false
-    val maxRefreshes = 10
-    var numRefreshes = 0
-
-    // refresh token at start of polling since local build just prior can take a long time
-    refreshToken(project)
 
     try {
         waitUntil(
@@ -138,13 +133,10 @@ suspend fun JobId.pollTransformationStatusAndPlan(
                     onStateChange(state, newStatus, transformationPlan)
                 }
                 state = newStatus
-                numRefreshes = 0
-                return@waitUntil state
-            } catch (e: AccessDeniedException) {
-                if (numRefreshes++ > maxRefreshes) throw e
-                refreshToken(project)
                 return@waitUntil state
-            } catch (e: InvalidGrantException) {
+            } catch (e: Exception) {
+                if (e !is AccessDeniedException && e !is InvalidGrantException) throw e
+
                 CodeTransformMessageListener.instance.onReauthStarted()
                 notifyStickyWarn(
                     message("codemodernizer.notification.warn.expired_credentials.title"),

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

@@ -8,7 +8,6 @@ import com.intellij.openapi.vfs.VfsUtilCore
 import com.intellij.openapi.vfs.VirtualFileManager
 import software.amazon.awssdk.services.codewhispererruntime.model.TransformationLanguage
 import software.amazon.awssdk.services.codewhispererruntime.model.TransformationStatus
-import software.aws.toolkits.core.TokenConnectionSettings
 import software.aws.toolkits.jetbrains.core.credentials.AwsBearerTokenConnection
 import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
 import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection
@@ -43,12 +42,6 @@ val STATES_AFTER_STARTED = setOf(
     *STATES_AFTER_INITIAL_BUILD.toTypedArray(),
 )
 
-fun refreshToken(project: Project) {
-    val connection = ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance())
-    val provider = (connection?.getConnectionSettings() as TokenConnectionSettings).tokenProvider.delegate as BearerTokenProvider
-    provider.refresh()
-}
-
 fun getAuthType(project: Project): CredentialSourceId? {
     val connection = checkBearerConnectionValidity(project, BearerTokenFeatureSet.Q)
     var authType: CredentialSourceId? = null

plugins/amazonq/codetransform/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codemodernizer/CodeWhispererCodeModernizerTestBase.kt

@@ -45,7 +45,6 @@ import software.aws.toolkits.core.credentials.ToolkitBearerTokenProvider
 import software.aws.toolkits.core.utils.test.aString
 import software.aws.toolkits.jetbrains.core.credentials.AwsBearerTokenConnection
 import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
-import software.aws.toolkits.jetbrains.core.credentials.sso.DeviceAuthorizationGrantToken
 import software.aws.toolkits.jetbrains.core.credentials.sso.PKCEAuthorizationGrantToken
 import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.BearerTokenAuthState
 import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.BearerTokenProvider
@@ -250,11 +249,7 @@ open class CodeWhispererCodeModernizerTestBase(
         project = projectRule.project
         toolkitConnectionManager = spy(ToolkitConnectionManager.getInstance(project))
 
-        val accessToken = DeviceAuthorizationGrantToken(aString(), aString(), aString(), aString(), Instant.MAX, Instant.now())
-        val provider =
-            mock<BearerTokenProvider> {
-                doReturn(accessToken).whenever(it).refresh()
-            }
+        val provider = mock<BearerTokenProvider> { }
         val mockBearerProvider =
             mock<ToolkitBearerTokenProvider> {
                 doReturn(provider).whenever(it).delegate
@@ -340,7 +335,6 @@ open class CodeWhispererCodeModernizerTestBase(
         val accessToken = PKCEAuthorizationGrantToken(aString(), aString(), aString(), aString(), Instant.MAX, Instant.now())
 
         val provider = mock<BearerTokenProvider> {
-            doReturn(accessToken).whenever(it).refresh()
             doReturn(accessToken).whenever(it).currentToken()
             doReturn(authState).whenever(it).state()
         }

plugins/amazonq/codetransform/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codemodernizer/CodeWhispererCodeModernizerUtilsTest.kt

@@ -38,7 +38,6 @@ import software.aws.toolkits.jetbrains.services.codemodernizer.utils.getTableMap
 import software.aws.toolkits.jetbrains.services.codemodernizer.utils.isPlanComplete
 import software.aws.toolkits.jetbrains.services.codemodernizer.utils.parseBuildFile
 import software.aws.toolkits.jetbrains.services.codemodernizer.utils.pollTransformationStatusAndPlan
-import software.aws.toolkits.jetbrains.services.codemodernizer.utils.refreshToken
 import software.aws.toolkits.jetbrains.services.codemodernizer.utils.validateCustomVersionsFile
 import software.aws.toolkits.jetbrains.services.codemodernizer.utils.validateSctMetadata
 import software.aws.toolkits.jetbrains.utils.notifyStickyWarn
@@ -90,18 +89,18 @@ class CodeWhispererCodeModernizerUtilsTest : CodeWhispererCodeModernizerTestBase
     }
 
     @Test
-    fun `refresh on access denied`() {
+    fun `show re-auth notification on access denied`() {
         val mockAccessDeniedException = Mockito.mock(AccessDeniedException::class.java)
 
-        mockkStatic(::refreshToken)
-        every { refreshToken(any()) } just runs
+        mockkStatic(::notifyStickyWarn)
+        every { notifyStickyWarn(any(), any(), any(), any(), any()) } just runs
 
         Mockito.doThrow(
             mockAccessDeniedException
         ).doReturn(
             exampleGetCodeMigrationResponse,
             exampleGetCodeMigrationResponse.replace(TransformationStatus.STARTED),
-            exampleGetCodeMigrationResponse.replace(TransformationStatus.COMPLETED), // Should stop before this point
+            exampleGetCodeMigrationResponse.replace(TransformationStatus.COMPLETED),
         ).whenever(clientAdaptorSpy).getCodeModernizationJob(any())
 
         Mockito.doReturn(exampleGetCodeMigrationPlanResponse)
@@ -128,7 +127,7 @@ class CodeWhispererCodeModernizerUtilsTest : CodeWhispererCodeModernizerTestBase
                 TransformationStatus.STARTED,
             )
         assertThat(expected).isEqualTo(mutableList)
-        io.mockk.verify { refreshToken(any()) }
+        verify { notifyStickyWarn(message("codemodernizer.notification.warn.expired_credentials.title"), any(), any(), any(), any()) }
     }
 
     @Test

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)
                 }
             }