Merge pull request #5676 from aws/autoMerge/feature/q-lsp-chat

Merge main into feature/q-lsp-chat

Author: rli

.changes/next-release/feature-ae7d15f4-ae07-4faf-963a-6d83a85852e4.json

@@ -0,0 +1,4 @@
+{
+  "type" : "feature",
+  "description" : "Amazon Q: Support selecting customizations across all Q profiles with automatic profile switching for enterprise users"
+}

plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/startup/CodeWhispererProjectStartupActivity.kt

@@ -13,7 +13,6 @@ import software.aws.toolkits.jetbrains.core.coroutines.projectCoroutineScope
 import software.aws.toolkits.jetbrains.services.amazonq.CodeWhispererFeatureConfigService
 import software.aws.toolkits.jetbrains.services.amazonq.calculateIfIamIdentityCenterConnection
 import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.CodeWhispererCodeScanManager
-import software.aws.toolkits.jetbrains.services.codewhisperer.customization.CodeWhispererCustomization
 import software.aws.toolkits.jetbrains.services.codewhisperer.customization.CodeWhispererModelConfigurator
 import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.CodeWhispererExplorerActionManager
 import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.isUserBuilderId
@@ -79,18 +78,16 @@ class CodeWhispererProjectStartupActivity : StartupActivity.DumbAware {
         projectCoroutineScope(project).launch {
             while (isActive) {
                 CodeWhispererFeatureConfigService.getInstance().fetchFeatureConfigs(project)
-                CodeWhispererFeatureConfigService.getInstance().getCustomizationFeature()?.let { customization ->
+                CodeWhispererFeatureConfigService.getInstance().getCustomizationFeature()?.let { overrideContext ->
                     val persistedCustomizationOverride = CodeWhispererModelConfigurator.getInstance().getPersistedCustomizationOverride()
-                    val latestCustomizationOverride = customization.value.stringValue()
+                    val latestCustomizationOverride = overrideContext.value.stringValue()
                     if (persistedCustomizationOverride == latestCustomizationOverride) return@let
 
                     // latest is different from the currently persisted, need update
-                    CodeWhispererFeatureConfigService.getInstance().validateCustomizationOverride(project, customization)
-                    CodeWhispererModelConfigurator.getInstance().switchCustomization(
-                        project,
-                        CodeWhispererCustomization(arn = customization.value.stringValue(), name = customization.variation),
-                        isOverride = true
-                    )
+                    val customization = CodeWhispererFeatureConfigService.getInstance().validateCustomizationOverride(project, overrideContext)
+                    if (customization != null) {
+                        CodeWhispererModelConfigurator.getInstance().switchCustomization(project, customization, isOverride = true)
+                    }
                 }
 
                 delay(FEATURE_CONFIG_POLL_INTERVAL_IN_MS)

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

@@ -18,7 +18,6 @@ import org.mockito.kotlin.argumentCaptor
 import org.mockito.kotlin.doReturn
 import org.mockito.kotlin.mock
 import org.mockito.kotlin.stub
-import org.mockito.kotlin.times
 import org.mockito.kotlin.verify
 import software.amazon.awssdk.services.codewhispererruntime.CodeWhispererRuntimeClient
 import software.amazon.awssdk.services.codewhispererruntime.model.ArtifactType

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

@@ -35,6 +35,7 @@ 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.SONO_URL
 import software.aws.toolkits.jetbrains.services.amazonq.CodeWhispererFeatureConfigService
+import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfile
 import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileManager
 import kotlin.reflect.full.memberFunctions
 import kotlin.test.Test
@@ -78,7 +79,7 @@ class CodeWhispererFeatureConfigServiceTest {
 
         projectRule.project.replaceService(
             QRegionProfileManager::class.java,
-            mock<QRegionProfileManager> { on { getQClient(any<Project>(), eq(CodeWhispererRuntimeClient::class)) } doReturn mockClient },
+            mock<QRegionProfileManager> { on { getQClient(any<Project>(), eq(QRegionProfile()), eq(CodeWhispererRuntimeClient::class)) } doReturn mockClient },
             disposableRule.disposable
         )
 

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

@@ -14,6 +14,7 @@ import migration.software.aws.toolkits.jetbrains.services.codewhisperer.customiz
 import org.assertj.core.api.Assertions.assertThat
 import org.jdom.output.XMLOutputter
 import org.junit.Before
+import org.junit.Ignore
 import org.junit.Rule
 import org.junit.Test
 import org.mockito.kotlin.any
@@ -42,6 +43,8 @@ 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.QRegionProfile
+import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileManager
 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
@@ -82,6 +85,7 @@ class CodeWhispererModelConfiguratorTest {
     private lateinit var mockClient: CodeWhispererRuntimeClient
     private lateinit var abManager: CodeWhispererFeatureConfigService
     private lateinit var mockClintAdaptor: CodeWhispererClientAdaptor
+    private lateinit var mockQRegionProfileManager: QRegionProfileManager
 
     @Before
     fun setup() {
@@ -124,6 +128,11 @@ class CodeWhispererModelConfiguratorTest {
 
         mockClintAdaptor = mock()
         projectRule.project.registerServiceInstance(CodeWhispererClientAdaptor::class.java, mockClintAdaptor)
+
+        mockQRegionProfileManager = mock {
+            on { listRegionProfiles(any()) }.thenReturn(listOf(QRegionProfile("fake_name", "fake_arn")))
+        }
+        ApplicationManager.getApplication().replaceService(QRegionProfileManager::class.java, mockQRegionProfileManager, disposableRule.disposable)
     }
 
     @Test
@@ -431,7 +440,10 @@ class CodeWhispererModelConfiguratorTest {
 
             this.connectionIdToActiveCustomizationArn.putAll(
                 mapOf(
-                    "fake-sso-url" to CodeWhispererCustomization(arn = "arn_2", name = "name_2", description = "description_2")
+                    "fake-sso-url" to CodeWhispererCustomization(
+                        arn = "arn_2", name = "name_2", description = "description_2",
+                        profile = QRegionProfile(profileName = "myActiveProfile", arn = "arn:aws:codewhisperer:us-west-2:123456789012:profile/myActiveProfile")
+                    )
                 )
             )
 
@@ -450,6 +462,12 @@ class CodeWhispererModelConfiguratorTest {
             "<option name=\"arn\" value=\"arn_2\" />" +
             "<option name=\"name\" value=\"name_2\" />" +
             "<option name=\"description\" value=\"description_2\" />" +
+            "<option name=\"profile\">" +
+            "<QRegionProfile>" +
+            "<option name=\"arn\" value=\"arn:aws:codewhisperer:us-west-2:123456789012:profile/myActiveProfile\" />" +
+            "<option name=\"profileName\" value=\"myActiveProfile\" />" +
+            "</QRegionProfile>" +
+            "</option>" +
             "</CodeWhispererCustomization>" +
             "</value>" +
             "</entry>" +
@@ -500,6 +518,10 @@ class CodeWhispererModelConfiguratorTest {
                               <option name="arn" value="arn_2" />
                               <option name="name" value="name_2" />
                               <option name="description" value="description_2" />
+                              <option name="profile"><QRegionProfile>
+                              <option name="profileName" value="myActiveProfile"/>
+                              <option name="arn" value="arn:aws:codewhisperer:us-west-2:123456789012:profile/myActiveProfile"/>
+                            </QRegionProfile></option>
                             </CodeWhispererCustomization>
                           </value>
                         </entry>
@@ -529,14 +551,42 @@ class CodeWhispererModelConfiguratorTest {
             CodeWhispererCustomization(
                 arn = "arn_2",
                 name = "name_2",
-                description = "description_2"
+                description = "description_2",
+                profile = QRegionProfile("myActiveProfile", "arn:aws:codewhisperer:us-west-2:123456789012:profile/myActiveProfile")
             )
         )
 
         assertThat(actual.previousAvailableCustomizations).hasSize(1)
         assertThat(actual.previousAvailableCustomizations["fake-sso-url"]).isEqualTo(listOf("arn_1", "arn_2", "arn_3"))
     }
 
+    @Test
+    fun `backward compatibility - should still be deseriealizable where profile field is not present`() {
+        val xml = xmlElement(
+            """
+                <component name="codewhispererCustomizationStates">
+                        <option name="connectionIdToActiveCustomizationArn">
+                            <map>
+                                <entry key="sso-session:foo">
+                                    <value>
+                                        <CodeWhispererCustomization>
+                                            <option name="arn" value="arn:foo" />
+                                            <option name="name" value="Customization-foo" />
+                                            <option name="description" value="Foo foo foo foo" />
+                                        </CodeWhispererCustomization>
+                                    </value>
+                                </entry>
+                            </map>
+                        </option>
+                    </component>
+            """.trimIndent()
+        )
+
+        val actual = XmlSerializer.deserialize(xml, CodeWhispererCustomizationState::class.java)
+        val cnt = actual.connectionIdToActiveCustomizationArn.size
+        assertThat(cnt).isEqualTo(1)
+    }
+
     @Test
     fun `deserialize users choosing default customization`() {
         val element = xmlElement(
@@ -565,6 +615,7 @@ class CodeWhispererModelConfiguratorTest {
         assertThat(actual.previousAvailableCustomizations["fake-sso-url"]).isEqualTo(listOf("arn_1", "arn_2", "arn_3"))
     }
 
+    @Ignore
     @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))
@@ -573,6 +624,11 @@ class CodeWhispererModelConfiguratorTest {
         sut.switchCustomization(projectRule.project, oldCustomization)
 
         assertThat(sut.activeCustomization(projectRule.project)).isEqualTo(oldCustomization)
+        // TODO: mock sdk client to fix the test
+//        val fakeCustomizations = listOf(
+//            CodeWhispererCustomization("oldArn", "oldName", "oldDescription")
+//        )
+//        mockClintAdaptor.stub { on { listAvailableCustomizations(QRegionProfile("fake_name", "fake_arn")) } doReturn fakeCustomizations }
 
         ApplicationManager.getApplication().messageBus
             .syncPublisher(QRegionProfileSelectedListener.TOPIC)
@@ -581,6 +637,7 @@ class CodeWhispererModelConfiguratorTest {
         assertThat(sut.activeCustomization(projectRule.project)).isEqualTo(oldCustomization)
     }
 
+    @Ignore
     @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))
@@ -589,6 +646,12 @@ class CodeWhispererModelConfiguratorTest {
         sut.switchCustomization(projectRule.project, oldCustomization)
         assertThat(sut.activeCustomization(projectRule.project)).isEqualTo(oldCustomization)
 
+        // TODO: mock sdk client to fix the test
+//        val fakeCustomizations = listOf(
+//            CodeWhispererCustomization("newArn", "newName", "newDescription")
+//        )
+//        mockClintAdaptor.stub { on { listAvailableCustomizations(QRegionProfile("fake_name", "fake_arn")) } doReturn fakeCustomizations }
+
         val latch = CountDownLatch(1)
 
         ApplicationManager.getApplication().messageBus

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

@@ -178,7 +178,7 @@ class QRegionProfileManagerTest {
         client.stub {
             onGeneric { listAvailableProfilesPaginator(any<Consumer<ListAvailableProfilesRequest.Builder>>()) } doReturn iterable
         }
-        val connectionSettings = sut.getQClientSettings(project)
+        val connectionSettings = sut.getQClientSettings(project, null)
         resourceCache.addEntry(connectionSettings, QProfileResources.LIST_REGION_PROFILES, QProfileResources.LIST_REGION_PROFILES.fetch(connectionSettings))
 
         assertThat(sut.listRegionProfiles(project))
@@ -247,7 +247,7 @@ class QRegionProfileManagerTest {
             sut.activeProfile(project)
         ).isEqualTo(QRegionProfile(arn = "arn:aws:codewhisperer:eu-central-1:123456789012:profile/FOO_PROFILE", profileName = "FOO_PROFILE"))
 
-        val settings = sut.getQClientSettings(project)
+        val settings = sut.getQClientSettings(project, null)
         assertThat(settings.region.id).isEqualTo(Region.EU_CENTRAL_1.id())
 
         sut.switchProfile(
@@ -259,7 +259,7 @@ class QRegionProfileManagerTest {
             sut.activeProfile(project)
         ).isEqualTo(QRegionProfile(arn = "arn:aws:codewhisperer:us-east-1:123456789012:profile/BAR_PROFILE", profileName = "BAR_PROFILE"))
 
-        val settings2 = sut.getQClientSettings(project)
+        val settings2 = sut.getQClientSettings(project, null)
         assertThat(settings2.region.id).isEqualTo(Region.US_EAST_1.id())
     }
 
@@ -275,7 +275,7 @@ class QRegionProfileManagerTest {
         assertThat(
             sut.activeProfile(project)
         ).isEqualTo(QRegionProfile(arn = "arn:aws:codewhisperer:eu-central-1:123456789012:profile/FOO_PROFILE", profileName = "FOO_PROFILE"))
-        assertThat(sut.getQClientSettings(project).region.id).isEqualTo(Region.EU_CENTRAL_1.id())
+        assertThat(sut.getQClientSettings(project, null).region.id).isEqualTo(Region.EU_CENTRAL_1.id())
 
         val client = sut.getQClient<CodeWhispererRuntimeClient>(project)
         assertThat(client).isInstanceOf(CodeWhispererRuntimeClient::class.java)
@@ -292,7 +292,7 @@ class QRegionProfileManagerTest {
         assertThat(
             sut.activeProfile(project)
         ).isEqualTo(QRegionProfile(arn = "arn:aws:codewhisperer:us-east-1:123456789012:profile/BAR_PROFILE", profileName = "BAR_PROFILE"))
-        assertThat(sut.getQClientSettings(project).region.id).isEqualTo(Region.US_EAST_1.id())
+        assertThat(sut.getQClientSettings(project, null).region.id).isEqualTo(Region.US_EAST_1.id())
 
         val client2 = sut.getQClient<CodeWhispererRuntimeClient>(project)
         assertThat(client2).isInstanceOf(CodeWhispererRuntimeClient::class.java)

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

@@ -15,6 +15,7 @@ import software.aws.toolkits.core.utils.getLogger
 import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
 import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection
 import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileManager
+import software.aws.toolkits.jetbrains.services.codewhisperer.customization.CodeWhispererCustomization
 import software.aws.toolkits.jetbrains.utils.isQExpired
 
 @Service
@@ -103,37 +104,49 @@ class CodeWhispererFeatureConfigService {
         }
     }
 
-    fun validateCustomizationOverride(project: Project, customization: FeatureContext) {
-        val customizationArnOverride = customization.value.stringValue()
-        val connection = connection(project) ?: return
-        if (customizationArnOverride != null) {
-            // Double check if server-side wrongly returns a customizationArn to BID users
-            calculateIfBIDConnection(project) {
-                featureConfigs.remove(CUSTOMIZATION_ARN_OVERRIDE_NAME)
-            }
-            val availableCustomizations =
-                calculateIfIamIdentityCenterConnection(project) {
-                    try {
+    fun validateCustomizationOverride(project: Project, featOverrideContext: FeatureContext): CodeWhispererCustomization? {
+        val customizationArnOverride = featOverrideContext.value.stringValue()
+        connection(project) ?: return null
+        customizationArnOverride ?: return null
+
+        // Double check if server-side wrongly returns a customizationArn to BID users
+        calculateIfBIDConnection(project) {
+            featureConfigs.remove(CUSTOMIZATION_ARN_OVERRIDE_NAME)
+        }
+        val availableCustomizations =
+            calculateIfIamIdentityCenterConnection(project) {
+                try {
+                    val profiles = QRegionProfileManager.getInstance().listRegionProfiles(project)
+                        ?: error("Attempted to fetch profiles while there does not exist")
+
+                    val customs = profiles.flatMap { profile ->
                         QRegionProfileManager.getInstance().getQClient<CodeWhispererRuntimeClient>(project)
-                            .listAvailableCustomizationsPaginator {}
-                            .flatMap { resp ->
-                                resp.customizations().map {
-                                    it.arn()
-                                }
+                            .listAvailableCustomizations { it.profileArn(profile.arn) }.customizations().map { originalCustom ->
+                                CodeWhispererCustomization(
+                                    arn = originalCustom.arn(),
+                                    name = originalCustom.name(),
+                                    description = originalCustom.description(),
+                                    profile = profile
+                                )
                             }
-                    } catch (e: Exception) {
-                        LOG.debug(e) { "Failed to list available customizations" }
-                        null
                     }
-                }
 
-            // If customizationArn from A/B is not available in listAvailableCustomizations response, don't use this value
-            if (availableCustomizations?.contains(customizationArnOverride) == false) {
-                LOG.debug {
-                    "Customization arn $customizationArnOverride not available in listAvailableCustomizations, not using"
+                    customs
+                } catch (e: Exception) {
+                    LOG.debug(e) { "encountered error while validating customization override" }
+                    null
                 }
-                featureConfigs.remove(CUSTOMIZATION_ARN_OVERRIDE_NAME)
             }
+
+        val isValidOverride = availableCustomizations != null && availableCustomizations.any { it.arn == customizationArnOverride }
+
+        // If customizationArn from A/B is not available in listAvailableCustomizations response, don't use this value
+        return if (!isValidOverride) {
+            LOG.debug { "Customization arn $customizationArnOverride not available in listAvailableCustomizations, not using" }
+            featureConfigs.remove(CUSTOMIZATION_ARN_OVERRIDE_NAME)
+            null
+        } else {
+            availableCustomizations?.find { it.arn == customizationArnOverride }
         }
     }