feedback

Author: rli

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

@@ -4,7 +4,6 @@
 package software.aws.toolkits.jetbrains.services.amazonq.toolwindow
 
 import com.intellij.openapi.actionSystem.ActionManager
-import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.application.runInEdt
 import com.intellij.openapi.project.DumbAware
 import com.intellij.openapi.project.Project
@@ -66,9 +65,6 @@ class AmazonQToolWindowFactory : ToolWindowFactory, DumbAware {
                 override fun activeConnectionChanged(newConnection: ToolkitConnection?) {
                     ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance())?.let { qConn ->
                         openMeetQPage(project)
-                        ApplicationManager.getApplication().executeOnPooledThread {
-                            QRegionProfileManager.getInstance().validateProfile(project)
-                        }
                     }
                     prepareChatContent(project, qPanel)
                 }

plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/QRegionProfileManagerTest.kt

@@ -40,6 +40,7 @@ import software.aws.toolkits.jetbrains.services.amazonq.profile.QProfileSwitchIn
 import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfile
 import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileManager
 import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileSelectedListener
+import software.aws.toolkits.jetbrains.utils.satisfiesKt
 import software.aws.toolkits.jetbrains.utils.xmlElement
 import java.net.URI
 import java.util.function.Consumer
@@ -105,6 +106,23 @@ class QRegionProfileManagerTest {
         assertThat(sut.activeProfile(project)).isNull()
     }
 
+    @Test
+    fun `data is cleared when user logs out`() {
+        sut.switchProfile(project, QRegionProfile(arn = "arn", profileName = "foo_profile"), QProfileSwitchIntent.User)
+        assertThat(sut.activeProfile(project)).isEqualTo(QRegionProfile(arn = "arn", profileName = "foo_profile"))
+
+        ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance())?.let {
+            if (it is AwsBearerTokenConnection) {
+                logoutFromSsoConnection(project, it)
+            }
+        }
+
+        assertThat(sut.state).satisfiesKt {
+            assertThat(it.connectionIdToActiveProfile).isEmpty()
+            assertThat(it.connectionIdToProfileList).isEmpty()
+        }
+    }
+
     @Test
     fun `switch should send message onProfileChanged for active switch`() {
         var cnt = 0

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

@@ -4,6 +4,7 @@
 package software.aws.toolkits.jetbrains.services.amazonq.profile
 
 import com.intellij.openapi.Disposable
+import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.components.BaseState
 import com.intellij.openapi.components.PersistentStateComponent
 import com.intellij.openapi.components.Service
@@ -26,6 +27,7 @@ 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
 import software.aws.toolkits.jetbrains.core.credentials.sono.isSono
+import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.BearerTokenProviderListener
 import software.aws.toolkits.jetbrains.core.region.AwsRegionProvider
 import software.aws.toolkits.jetbrains.utils.notifyInfo
 import software.aws.toolkits.resources.AmazonQBundle.message
@@ -41,7 +43,20 @@ class QRegionProfileManager : PersistentStateComponent<QProfileState>, Disposabl
 
     // Map to store connectionId to its active profile
     private val connectionIdToActiveProfile = Collections.synchronizedMap<String, QRegionProfile>(mutableMapOf())
-    private val connectionIdToProfileList = mutableMapOf<String, Int>()
+    private val connectionIdToProfileCount = mutableMapOf<String, Int>()
+
+    init {
+        ApplicationManager.getApplication().messageBus.connect(this)
+            .subscribe(
+                BearerTokenProviderListener.TOPIC,
+                object : BearerTokenProviderListener {
+                    override fun invalidate(providerId: String) {
+                        connectionIdToActiveProfile.remove(providerId)
+                        connectionIdToProfileCount.remove(providerId)
+                    }
+                }
+            )
+    }
 
     // should be call on project startup to validate if profile is still active
     @RequiresBackgroundThread
@@ -80,7 +95,7 @@ class QRegionProfileManager : PersistentStateComponent<QProfileState>, Disposabl
                 switchProfile(project, mappedProfiles.first(), intent = QProfileSwitchIntent.Update)
             }
             mappedProfiles.takeIf { it.isNotEmpty() }?.also {
-                connectionIdToProfileList[connection.id] = it.size
+                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}" }
@@ -112,7 +127,7 @@ class QRegionProfileManager : PersistentStateComponent<QProfileState>, Disposabl
                 Telemetry.amazonq.didSelectProfile.use { span ->
                     span.source(intent.value)
                         .amazonQProfileRegion(newProfile.region)
-                        .profileCount(connectionIdToProfileList[conn.id])
+                        .profileCount(connectionIdToProfileCount[conn.id])
                         .ssoRegion(conn.region)
                         .credentialStartUrl(conn.startUrl)
                         .result(MetricResult.Succeeded)
@@ -141,13 +156,13 @@ class QRegionProfileManager : PersistentStateComponent<QProfileState>, Disposabl
 
     // for each idc connection, user should have a profile, otherwise should show the profile selection error page
     fun isPendingProfileSelection(project: Project): Boolean = getIdcConnectionOrNull(project)?.let { conn ->
-        val profileCounts = connectionIdToProfileList[conn.id] ?: 0
+        val profileCounts = connectionIdToProfileCount[conn.id] ?: 0
         val activeProfile = connectionIdToActiveProfile[conn.id]
         profileCounts == 0 || (profileCounts > 1 && activeProfile?.arn.isNullOrEmpty())
     } ?: false
 
     fun shouldDisplayProfileInfo(project: Project): Boolean = getIdcConnectionOrNull(project)?.let { conn ->
-        (connectionIdToProfileList[conn.id] ?: 0) > 1
+        (connectionIdToProfileCount[conn.id] ?: 0) > 1
     } ?: false
 
     fun getQClientSettings(project: Project): TokenConnectionSettings {
@@ -193,16 +208,16 @@ class QRegionProfileManager : PersistentStateComponent<QProfileState>, Disposabl
     override fun getState(): QProfileState {
         val state = QProfileState()
         state.connectionIdToActiveProfile.putAll(this.connectionIdToActiveProfile)
-        state.connectionIdToProfileList.putAll(this.connectionIdToProfileList)
+        state.connectionIdToProfileList.putAll(this.connectionIdToProfileCount)
         return state
     }
 
     override fun loadState(state: QProfileState) {
         connectionIdToActiveProfile.clear()
         connectionIdToActiveProfile.putAll(state.connectionIdToActiveProfile)
 
-        connectionIdToProfileList.clear()
-        connectionIdToProfileList.putAll(state.connectionIdToProfileList)
+        connectionIdToProfileCount.clear()
+        connectionIdToProfileCount.putAll(state.connectionIdToProfileList)
     }
 }