fix(amazonq): Invalidate Obsolete Customization on Profile Switch (#5575)

When a user switches profiles, the active customization is now verified against the updated list of available customizations.

If the active customization is not present in the newly returned list, it is automatically invalidated.
---------

Co-authored-by: aws-toolkit-automation <43144436+aws-toolkit-automation@users.noreply.github.com>
Co-authored-by: Richard Li <742829+rli@users.noreply.github.com>

Author: 3 people

.changes/next-release/bugfix-a70cb565-5de8-4302-9f1e-71925f78061b.json

@@ -0,0 +1,4 @@
+{
+  "type" : "bugfix",
+  "description" : "Amazon Q: Customization now resets with a warning if unavailable in the selected profile."
+}

plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/customization/CodeWhispererModelConfigurator.kt

@@ -22,6 +22,8 @@ import software.aws.toolkits.core.utils.debug
 import software.aws.toolkits.core.utils.getLogger
 import software.aws.toolkits.jetbrains.services.amazonq.CodeWhispererFeatureConfigService
 import software.aws.toolkits.jetbrains.services.amazonq.calculateIfIamIdentityCenterConnection
+import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfile
+import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileSelectedListener
 import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererClientAdaptor
 import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants
 import software.aws.toolkits.jetbrains.utils.notifyInfo
@@ -48,6 +50,7 @@ private fun notifyInvalidSelectedCustomization(project: Project) {
 }
 
 private fun notifyNewCustomization(project: Project) {
+    if (ApplicationManager.getApplication().isUnitTestMode) return
     notifyInfo(
         title = message("codewhisperer.custom.dialog.title"),
         content = message("codewhisperer.notification.custom.new_customization"),
@@ -81,6 +84,18 @@ class DefaultCodeWhispererModelConfigurator : CodeWhispererModelConfigurator, Pe
 
     private var customizationArnOverrideV2: String? = null
 
+    init {
+        ApplicationManager.getApplication().messageBus.connect(this).subscribe(
+            QRegionProfileSelectedListener.TOPIC,
+            object : QRegionProfileSelectedListener {
+                override fun onProfileSelected(project: Project, profile: QRegionProfile?) {
+                    pluginAwareExecuteOnPooledThread {
+                        CodeWhispererModelConfigurator.getInstance().listCustomizations(project, passive = true)
+                    }
+                }
+            }
+        )
+    }
     override fun showConfigDialog(project: Project) {
         runInEdt {
             calculateIfIamIdentityCenterConnection(project) {

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

@@ -7,8 +7,10 @@ import com.intellij.openapi.application.ApplicationManager
 import com.intellij.testFramework.ApplicationRule
 import com.intellij.testFramework.DisposableRule
 import com.intellij.testFramework.ProjectRule
+import com.intellij.testFramework.registerServiceInstance
 import com.intellij.testFramework.replaceService
 import com.intellij.util.xmlb.XmlSerializer
+import migration.software.aws.toolkits.jetbrains.services.codewhisperer.customization.CodeWhispererModelConfigurator
 import org.assertj.core.api.Assertions.assertThat
 import org.jdom.output.XMLOutputter
 import org.junit.Before
@@ -40,10 +42,14 @@ import software.aws.toolkits.jetbrains.core.credentials.sono.isSono
 import software.aws.toolkits.jetbrains.core.region.MockRegionProviderRule
 import software.aws.toolkits.jetbrains.services.amazonq.CodeWhispererFeatureConfigService
 import software.aws.toolkits.jetbrains.services.amazonq.FeatureContext
+import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileSelectedListener
+import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererClientAdaptor
 import software.aws.toolkits.jetbrains.services.codewhisperer.customization.CodeWhispererCustomization
 import software.aws.toolkits.jetbrains.services.codewhisperer.customization.CodeWhispererCustomizationState
 import software.aws.toolkits.jetbrains.services.codewhisperer.customization.DefaultCodeWhispererModelConfigurator
 import software.aws.toolkits.jetbrains.utils.xmlElement
+import java.util.concurrent.CountDownLatch
+import java.util.concurrent.TimeUnit
 import kotlin.reflect.full.memberProperties
 import kotlin.reflect.jvm.isAccessible
 
@@ -75,6 +81,7 @@ class CodeWhispererModelConfiguratorTest {
     private lateinit var sut: DefaultCodeWhispererModelConfigurator
     private lateinit var mockClient: CodeWhispererRuntimeClient
     private lateinit var abManager: CodeWhispererFeatureConfigService
+    private lateinit var mockClintAdaptor: CodeWhispererClientAdaptor
 
     @Before
     fun setup() {
@@ -83,7 +90,11 @@ class CodeWhispererModelConfiguratorTest {
         regionProvider.addRegion(Region.US_EAST_1)
         regionProvider.addRegion(Region.US_EAST_2)
 
-        sut = DefaultCodeWhispererModelConfigurator()
+        sut = spy(CodeWhispererModelConfigurator.getInstance() as DefaultCodeWhispererModelConfigurator).also { spyInstance ->
+            ApplicationManager.getApplication().replaceService(
+                DefaultCodeWhispererModelConfigurator::class.java, spyInstance, disposableRule.disposable
+            )
+        }
 
         (ToolkitConnectionManager.getInstance(projectRule.project) as DefaultToolkitConnectionManager).loadState(ToolkitConnectionManagerState())
         mockClient.stub {
@@ -110,6 +121,9 @@ class CodeWhispererModelConfiguratorTest {
             abManager,
             disposableRule.disposable
         )
+
+        mockClintAdaptor = mock()
+        projectRule.project.registerServiceInstance(CodeWhispererClientAdaptor::class.java, mockClintAdaptor)
     }
 
     @Test
@@ -550,4 +564,48 @@ class CodeWhispererModelConfiguratorTest {
         assertThat(actual.previousAvailableCustomizations).hasSize(1)
         assertThat(actual.previousAvailableCustomizations["fake-sso-url"]).isEqualTo(listOf("arn_1", "arn_2", "arn_3"))
     }
+
+    @Test
+    fun `profile switch should keep using existing customization if new list still contains that arn`() {
+        val ssoConn = spy(LegacyManagedBearerSsoConnection(region = "us-east-1", startUrl = "url 1", scopes = Q_SCOPES))
+        ToolkitConnectionManager.getInstance(projectRule.project).switchConnection(ssoConn)
+        val oldCustomization = CodeWhispererCustomization("oldArn", "oldName", "oldDescription")
+        sut.switchCustomization(projectRule.project, oldCustomization)
+
+        assertThat(sut.activeCustomization(projectRule.project)).isEqualTo(oldCustomization)
+
+        val fakeCustomizations = listOf(
+            CodeWhispererCustomization("oldArn", "oldName", "oldDescription")
+        )
+        mockClintAdaptor.stub { on { listAvailableCustomizations() } doReturn fakeCustomizations }
+
+        ApplicationManager.getApplication().messageBus
+            .syncPublisher(QRegionProfileSelectedListener.TOPIC)
+            .onProfileSelected(projectRule.project, null)
+
+        assertThat(sut.activeCustomization(projectRule.project)).isEqualTo(oldCustomization)
+    }
+
+    @Test
+    fun `profile switch should invalidate obsolete customization if it's not in the new list`() {
+        val ssoConn = spy(LegacyManagedBearerSsoConnection(region = "us-east-1", startUrl = "url 1", scopes = Q_SCOPES))
+        ToolkitConnectionManager.getInstance(projectRule.project).switchConnection(ssoConn)
+        val oldCustomization = CodeWhispererCustomization("oldArn", "oldName", "oldDescription")
+        sut.switchCustomization(projectRule.project, oldCustomization)
+        assertThat(sut.activeCustomization(projectRule.project)).isEqualTo(oldCustomization)
+        val fakeCustomizations = listOf(
+            CodeWhispererCustomization("newArn", "newName", "newDescription")
+        )
+        mockClintAdaptor.stub { on { listAvailableCustomizations() } doReturn fakeCustomizations }
+
+        val latch = CountDownLatch(1)
+
+        ApplicationManager.getApplication().messageBus
+            .syncPublisher(QRegionProfileSelectedListener.TOPIC)
+            .onProfileSelected(projectRule.project, null)
+
+        latch.await(2, TimeUnit.SECONDS)
+
+        assertThat(sut.activeCustomization(projectRule.project)).isNull()
+    }
 }

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

@@ -144,7 +144,7 @@ class QRegionProfileManager : PersistentStateComponent<QProfileState>, Disposabl
             }
         }
 
-        project.messageBus
+        ApplicationManager.getApplication().messageBus
             .syncPublisher(QRegionProfileSelectedListener.TOPIC)
             .onProfileSelected(project, newProfile)
     }

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

@@ -8,7 +8,7 @@ import com.intellij.util.messages.Topic
 
 interface QRegionProfileSelectedListener {
     companion object {
-        @Topic.ProjectLevel
+        @Topic.AppLevel
         val TOPIC = Topic.create("QRegionProfileSelected", QRegionProfileSelectedListener::class.java)
     }