diff --git a/.changes/next-release/bugfix-1162ea20-1fcf-43a6-b502-16d8b7d588d7.json b/.changes/next-release/bugfix-1162ea20-1fcf-43a6-b502-16d8b7d588d7.json new file mode 100644 index 00000000000..771f6ca2c00 --- /dev/null +++ b/.changes/next-release/bugfix-1162ea20-1fcf-43a6-b502-16d8b7d588d7.json @@ -0,0 +1,4 @@ +{ + "type" : "bugfix", + "description" : "Fix issue where IDE freezes when logging into Amazon Q" +} \ No newline at end of file diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/toolwindow/AmazonQToolWindowFactory.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/toolwindow/AmazonQToolWindowFactory.kt index 0985d88b8ab..123534f38dc 100644 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/toolwindow/AmazonQToolWindowFactory.kt +++ b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/toolwindow/AmazonQToolWindowFactory.kt @@ -65,7 +65,6 @@ class AmazonQToolWindowFactory : ToolWindowFactory, DumbAware { override fun activeConnectionChanged(newConnection: ToolkitConnection?) { ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance())?.let { qConn -> openMeetQPage(project) - QRegionProfileManager.getInstance().validateProfile(project) } prepareChatContent(project, qPanel) } diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/QRegionProfileManagerTest.kt b/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/QRegionProfileManagerTest.kt index 0453cf3a766..e8bb0398ad7 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/QRegionProfileManagerTest.kt +++ b/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 diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/profile/QRegionProfileManager.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/profile/QRegionProfileManager.kt index 2d345cc6df0..c9014d16d20 100644 --- a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/profile/QRegionProfileManager.kt +++ b/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 @@ -11,6 +12,7 @@ import com.intellij.openapi.components.State import com.intellij.openapi.components.Storage import com.intellij.openapi.components.service import com.intellij.openapi.project.Project +import com.intellij.util.concurrency.annotations.RequiresBackgroundThread import com.intellij.util.xmlb.annotations.MapAnnotation import com.intellij.util.xmlb.annotations.Property import software.amazon.awssdk.core.SdkClient @@ -25,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 @@ -40,9 +43,23 @@ class QRegionProfileManager : PersistentStateComponent, Disposabl // Map to store connectionId to its active profile private val connectionIdToActiveProfile = Collections.synchronizedMap(mutableMapOf()) - private val connectionIdToProfileList = mutableMapOf() + private val connectionIdToProfileCount = mutableMapOf() + + 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 fun validateProfile(project: Project) { val conn = getIdcConnectionOrNull(project) val selected = activeProfile(project) ?: return @@ -78,7 +95,7 @@ class QRegionProfileManager : PersistentStateComponent, 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}" } @@ -110,7 +127,7 @@ class QRegionProfileManager : PersistentStateComponent, 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) @@ -139,13 +156,13 @@ class QRegionProfileManager : PersistentStateComponent, 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 { @@ -191,7 +208,7 @@ class QRegionProfileManager : PersistentStateComponent, 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 } @@ -199,8 +216,8 @@ class QRegionProfileManager : PersistentStateComponent, Disposabl connectionIdToActiveProfile.clear() connectionIdToActiveProfile.putAll(state.connectionIdToActiveProfile) - connectionIdToProfileList.clear() - connectionIdToProfileList.putAll(state.connectionIdToProfileList) + connectionIdToProfileCount.clear() + connectionIdToProfileCount.putAll(state.connectionIdToProfileList) } }