From e0c91fd47b819ba38a4f46edf4f06f9f313e2556 Mon Sep 17 00:00:00 2001 From: Will Lo <96078566+Will-ShaoHua@users.noreply.github.com> Date: Tue, 6 May 2025 14:42:23 -0700 Subject: [PATCH 1/3] feat(amazonq): show all customizations across different profiles (#5646) Co-authored-by: evanliu048 --- ...-ae7d15f4-ae07-4faf-963a-6d83a85852e4.json | 4 + .../credentials/CodeWhispererClientAdaptor.kt | 12 +-- .../CodeWhispererCustomizationDialog.kt | 7 +- .../CodeWhispererModelConfigurator.kt | 89 +++++++++++-------- .../CodeWhispererProjectStartupActivity.kt | 15 ++-- .../CodeWhispererClientAdaptorTest.kt | 9 +- .../CodeWhispererFeatureConfigServiceTest.kt | 3 +- .../CodeWhispererModelConfiguratorTest.kt | 57 +++++++++++- .../QRegionProfileManagerTest.kt | 10 +-- .../CodeWhispererFeatureConfigService.kt | 63 +++++++------ .../amazonq/profile/QProfileSwitchIntent.kt | 4 +- .../amazonq/profile/QRegionProfileManager.kt | 20 ++--- .../CodeWhispererCustomization.kt | 4 + 13 files changed, 194 insertions(+), 103 deletions(-) create mode 100644 .changes/next-release/feature-ae7d15f4-ae07-4faf-963a-6d83a85852e4.json diff --git a/.changes/next-release/feature-ae7d15f4-ae07-4faf-963a-6d83a85852e4.json b/.changes/next-release/feature-ae7d15f4-ae07-4faf-963a-6d83a85852e4.json new file mode 100644 index 00000000000..d11656f96b8 --- /dev/null +++ b/.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" +} \ No newline at end of file diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/credentials/CodeWhispererClientAdaptor.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/credentials/CodeWhispererClientAdaptor.kt index 9f7a91a90e3..4bde496b98c 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/credentials/CodeWhispererClientAdaptor.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/credentials/CodeWhispererClientAdaptor.kt @@ -39,6 +39,7 @@ import software.aws.toolkits.core.utils.debug import software.aws.toolkits.core.utils.getLogger import software.aws.toolkits.jetbrains.core.credentials.sono.isInternalUser import software.aws.toolkits.jetbrains.services.amazonq.codeWhispererUserContext +import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfile import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileManager import software.aws.toolkits.jetbrains.services.codewhisperer.customization.CodeWhispererCustomization import software.aws.toolkits.jetbrains.services.codewhisperer.language.CodeWhispererProgrammingLanguage @@ -83,7 +84,7 @@ interface CodeWhispererClientAdaptor { fun getCodeFixJob(request: GetCodeFixJobRequest): GetCodeFixJobResponse - fun listAvailableCustomizations(): List + fun listAvailableCustomizations(profile: QRegionProfile): List fun startTestGeneration(uploadId: String, targetCode: List, userInput: String): StartTestGenerationResponse @@ -287,9 +288,9 @@ open class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeW override fun getCodeFixJob(request: GetCodeFixJobRequest): GetCodeFixJobResponse = bearerClient().getCodeFixJob(request) // DO NOT directly use this method to fetch customizations, use wrapper [CodeWhispererModelConfigurator.listCustomization()] instead - override fun listAvailableCustomizations(): List = - bearerClient().listAvailableCustomizationsPaginator( - ListAvailableCustomizationsRequest.builder().profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn).build() + override fun listAvailableCustomizations(profile: QRegionProfile): List = + QRegionProfileManager.getInstance().getQClient(project, profile).listAvailableCustomizationsPaginator( + ListAvailableCustomizationsRequest.builder().profileArn(profile.arn).build() ) .stream() .toList() @@ -303,7 +304,8 @@ open class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeW CodeWhispererCustomization( arn = it.arn(), name = it.name(), - description = it.description() + description = it.description(), + profile = profile ) } } diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/customization/CodeWhispererCustomizationDialog.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/customization/CodeWhispererCustomizationDialog.kt index eedff33e2c3..4bab85bceea 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/customization/CodeWhispererCustomizationDialog.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/customization/CodeWhispererCustomizationDialog.kt @@ -26,6 +26,7 @@ import software.amazon.awssdk.arns.Arn import software.aws.toolkits.core.utils.debug import software.aws.toolkits.core.utils.getLogger import software.aws.toolkits.core.utils.tryOrNull +import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfile import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants.Q_CUSTOM_LEARN_MORE_URI import software.aws.toolkits.jetbrains.ui.AsyncComboBox import software.aws.toolkits.jetbrains.utils.notifyInfo @@ -34,7 +35,7 @@ import javax.swing.JComponent import javax.swing.JList private val NoDataToDisplay = CustomizationUiItem( - CodeWhispererCustomization("", message("codewhisperer.custom.dialog.option.no_data"), ""), + CodeWhispererCustomization("", message("codewhisperer.custom.dialog.option.no_data"), "", QRegionProfile("", "")), false, false ) @@ -259,6 +260,10 @@ private object CustomizationRenderer : ColoredListCellRenderer? = calculateIfIamIdentityCenterConnection(project) { - // 1. invoke API and get result - val listAvailableCustomizationsResult = try { - CodeWhispererClientAdaptor.getInstance(project).listAvailableCustomizations() - } catch (e: Exception) { - val requestId = (e as? CodeWhispererRuntimeException)?.requestId() - val logMessage = if (CodeWhispererConstants.Customization.noAccessToCustomizationExceptionPredicate(e)) { - // TODO: not required for non GP users - "ListAvailableCustomizations: connection ${it.id} is not allowlisted, requestId: ${requestId.orEmpty()}" - } else { - "ListAvailableCustomizations: failed due to unknown error ${e.message}, requestId: ${requestId.orEmpty()}" - } - - LOG.debug { logMessage } - null + // 1. fetch all profiles, invoke fetch customizations API and get result for each profile and aggregate all the results + val profiles = QRegionProfileManager.getInstance().listRegionProfiles(project) + ?: error("Attempted to fetch profiles while there does not exist") + + val customizations = profiles.flatMap { profile -> + runCatching { + CodeWhispererClientAdaptor.getInstance(project).listAvailableCustomizations(profile) + }.onFailure { e -> + val requestId = (e as? CodeWhispererRuntimeException)?.requestId() + val logMessage = "ListAvailableCustomizations: failed due to unknown error ${e.message}, " + + "requestId: ${requestId.orEmpty()}, profileName: ${profile.profileName}" + LOG.debug { logMessage } + }.getOrDefault(emptyList()) } // 2. get diff val previousCustomizationsShapshot = connectionToCustomizationsShownLastTime.getOrElse(it.id) { emptyList() } - val diff = listAvailableCustomizationsResult?.filterNot { customization -> previousCustomizationsShapshot.contains(customization.arn) }?.toSet() + val diff = customizations.filterNot { customization -> previousCustomizationsShapshot.contains(customization.arn) }.toSet() // 3 if passive, // (1) update allowlisting @@ -135,42 +135,45 @@ class DefaultCodeWhispererModelConfigurator : CodeWhispererModelConfigurator, Pe // if not passive, // (1) update the customization list snapshot (seen by users last time) if it will be displayed if (passive) { - connectionIdToIsAllowlisted[it.id] = listAvailableCustomizationsResult != null - if (diff?.isNotEmpty() == true && !hasShownNewCustomizationNotification.getAndSet(true)) { + connectionIdToIsAllowlisted[it.id] = customizations.isNotEmpty() + if (diff.isNotEmpty() && !hasShownNewCustomizationNotification.getAndSet(true)) { notifyNewCustomization(project) } } else { - listAvailableCustomizationsResult?.let { customizations -> - connectionToCustomizationsShownLastTime[it.id] = customizations.map { customization -> customization.arn }.toMutableList() - } + connectionToCustomizationsShownLastTime[it.id] = customizations.map { customization -> customization.arn }.toMutableList() } // 4. invalidate selected customization if // (1) the API call failed // (2) the selected customization is not in the resultset of API call + // (3) the existing q region profile associated with the selected customization does not match the currently active profile activeCustomization(project)?.let { activeCustom -> - if (listAvailableCustomizationsResult == null) { + if (customizations.isEmpty()) { invalidateSelectedAndNotify(project) - } else if (!listAvailableCustomizationsResult.any { latestCustom -> latestCustom.arn == activeCustom.arn }) { + } else if (customizations.none { latestCustom -> latestCustom.arn == activeCustom.arn }) { invalidateSelectedAndNotify(project) + } else { + // for backward compatibility, previous schema didn't have profile arn, so backfill profile here if it's null + if (activeCustom.profile == null) { + customizations.find { c -> c.arn == activeCustom.arn }?.profile?.let { p -> + activeCustom.profile = p + } + } + + if (activeCustom.profile != null && activeCustom.profile != QRegionProfileManager.getInstance().activeProfile(project)) { + invalidateSelectedAndNotify(project) + } } } // 5. transform result to UI items and return - val customizationUiItems = if (diff != null) { - listAvailableCustomizationsResult.let { customizations -> - val nameToCount = customizations.groupingBy { customization -> customization.name }.eachCount() - - customizations.map { customization -> - CustomizationUiItem( - customization, - isNew = diff.contains(customization), - shouldPrefixAccountId = (nameToCount[customization.name] ?: 0) > 1 - ) - } - } - } else { - null + val nameToCount = customizations.groupingBy { customization -> customization.name }.eachCount() + val customizationUiItems = customizations.map { customization -> + CustomizationUiItem( + customization, + isNew = diff.contains(customization), + shouldPrefixAccountId = (nameToCount[customization.name] ?: 0) > 1 + ) } connectionToCustomizationUiItems[it.id] = customizationUiItems @@ -212,6 +215,18 @@ class DefaultCodeWhispererModelConfigurator : CodeWhispererModelConfigurator, Pe LOG.debug { "Switch from customization $oldCus to $newCustomization" } + // Switch profile if it doesn't match the customization's profile. + // Customizations are profile-scoped and must be used under the correct context. + newCustomization?.profile?.let { p -> + if (p.arn != QRegionProfileManager.getInstance().activeProfile(project)?.arn) { + QRegionProfileManager.getInstance().switchProfile( + project, + p, + QProfileSwitchIntent.Customization + ) + } + } + CodeWhispererCustomizationListener.notifyCustomUiUpdate() } if (isOverride) { diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/startup/CodeWhispererProjectStartupActivity.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/startup/CodeWhispererProjectStartupActivity.kt index 1900239222a..eb730d12ecf 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/startup/CodeWhispererProjectStartupActivity.kt +++ b/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) diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererClientAdaptorTest.kt b/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererClientAdaptorTest.kt index 9a3a455d6e1..246d0ed8c00 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererClientAdaptorTest.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererClientAdaptorTest.kt @@ -65,6 +65,7 @@ import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection import software.aws.toolkits.jetbrains.core.credentials.sono.Q_SCOPES import software.aws.toolkits.jetbrains.core.credentials.sono.SONO_REGION import software.aws.toolkits.jetbrains.services.amazonq.FEATURE_EVALUATION_PRODUCT_NAME +import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfile import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.metadata import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.pythonRequest import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.pythonResponseWithToken @@ -189,13 +190,13 @@ class CodeWhispererClientAdaptorTest { on { client.listAvailableCustomizationsPaginator(any()) } doReturn sdkIterable } - val actual = sut.listAvailableCustomizations() + val actual = sut.listAvailableCustomizations(QRegionProfile("fake_profile", "fake arn")) assertThat(actual).hasSize(3) assertThat(actual).isEqualTo( listOf( - CodeWhispererCustomization(name = "custom-1", arn = "arn-1"), - CodeWhispererCustomization(name = "custom-2", arn = "arn-2"), - CodeWhispererCustomization(name = "custom-3", arn = "arn-3") + CodeWhispererCustomization(name = "custom-1", arn = "arn-1", profile = QRegionProfile("fake_profile", "fake arn")), + CodeWhispererCustomization(name = "custom-2", arn = "arn-2", profile = QRegionProfile("fake_profile", "fake arn")), + CodeWhispererCustomization(name = "custom-3", arn = "arn-3", profile = QRegionProfile("fake_profile", "fake arn")) ) ) } diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererFeatureConfigServiceTest.kt b/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererFeatureConfigServiceTest.kt index 378e2703756..89de36d53fb 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererFeatureConfigServiceTest.kt +++ b/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 { on { getQClient(any(), eq(CodeWhispererRuntimeClient::class)) } doReturn mockClient }, + mock { on { getQClient(any(), eq(QRegionProfile()), eq(CodeWhispererRuntimeClient::class)) } doReturn mockClient }, disposableRule.disposable ) diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererModelConfiguratorTest.kt b/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererModelConfiguratorTest.kt index 751db478b5b..d6900a556cb 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererModelConfiguratorTest.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererModelConfiguratorTest.kt @@ -42,6 +42,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 +84,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 +127,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 +439,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 +461,12 @@ class CodeWhispererModelConfiguratorTest { "" + "" + "" + "" + @@ -500,6 +517,10 @@ class CodeWhispererModelConfiguratorTest { @@ -529,7 +550,8 @@ 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") ) ) @@ -537,6 +559,33 @@ class CodeWhispererModelConfiguratorTest { 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( + """ + + + + """.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( @@ -577,7 +626,7 @@ class CodeWhispererModelConfiguratorTest { val fakeCustomizations = listOf( CodeWhispererCustomization("oldArn", "oldName", "oldDescription") ) - mockClintAdaptor.stub { on { listAvailableCustomizations() } doReturn fakeCustomizations } + mockClintAdaptor.stub { on { listAvailableCustomizations(QRegionProfile("fake_name", "fake_arn")) } doReturn fakeCustomizations } ApplicationManager.getApplication().messageBus .syncPublisher(QRegionProfileSelectedListener.TOPIC) @@ -596,7 +645,7 @@ class CodeWhispererModelConfiguratorTest { val fakeCustomizations = listOf( CodeWhispererCustomization("newArn", "newName", "newDescription") ) - mockClintAdaptor.stub { on { listAvailableCustomizations() } doReturn fakeCustomizations } + mockClintAdaptor.stub { on { listAvailableCustomizations(QRegionProfile("fake_name", "fake_arn")) } doReturn fakeCustomizations } val latch = CountDownLatch(1) 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 53f5d3170f2..4b1f470b67e 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 @@ -178,7 +178,7 @@ class QRegionProfileManagerTest { client.stub { onGeneric { listAvailableProfilesPaginator(any>()) } 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(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(project) assertThat(client2).isInstanceOf(CodeWhispererRuntimeClient::class.java) diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/CodeWhispererFeatureConfigService.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/CodeWhispererFeatureConfigService.kt index d77663430fb..0fa47c27024 100644 --- a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/CodeWhispererFeatureConfigService.kt +++ b/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(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 } } } diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/profile/QProfileSwitchIntent.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/profile/QProfileSwitchIntent.kt index c1a7b4627ab..a48f33a2692 100644 --- a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/profile/QProfileSwitchIntent.kt +++ b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/profile/QProfileSwitchIntent.kt @@ -8,12 +8,14 @@ package software.aws.toolkits.jetbrains.services.amazonq.profile * 'auth' -> users change the profile through webview profile selector page * 'update' -> plugin auto select the profile on users' behalf as there is only 1 profile * 'reload' -> on plugin restart, plugin will try to reload previous selected profile + * 'customization' -> users selected a customization tied to a different profile, triggering a profile switch */ enum class QProfileSwitchIntent(val value: String) { User("user"), Auth("auth"), Update("update"), - Reload("reload"), ; + Reload("reload"), + Customization("customization"), ; override fun toString() = value } 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 32f814ab2b4..995bea0efe7 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 @@ -181,27 +181,25 @@ class QRegionProfileManager : PersistentStateComponent, Disposabl (connectionIdToProfileCount[conn.id] ?: 0) > 1 } ?: false - fun getQClientSettings(project: Project): TokenConnectionSettings { + fun getQClientSettings(project: Project, profile: QRegionProfile?): TokenConnectionSettings { val conn = ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance()) if (conn !is AwsBearerTokenConnection) { error("not a bearer connection") } val settings = conn.getConnectionSettings() - val awsRegion = AwsRegionProvider.getInstance()[QDefaultServiceConfig.REGION] ?: error("unknown region from Q default service config") + val defaultRegion = AwsRegionProvider.getInstance()[QDefaultServiceConfig.REGION] ?: error("unknown region from Q default service config") + val regionId = profile?.region ?: activeProfile(project)?.region + val awsRegion = regionId?.let { AwsRegionProvider.getInstance()[it] } ?: defaultRegion - // TODO: different window should be able to select different profile - return activeProfile(project)?.let { profile -> - AwsRegionProvider.getInstance()[profile.region]?.let { region -> - settings.withRegion(region) - } - } ?: settings.withRegion(awsRegion) + return settings.withRegion(awsRegion) } - inline fun getQClient(project: Project): T = getQClient(project, T::class) + inline fun getQClient(project: Project): T = getQClient(project, null, T::class) + inline fun getQClient(project: Project, profile: QRegionProfile): T = getQClient(project, profile, T::class) - fun getQClient(project: Project, sdkClass: KClass): T { - val settings = getQClientSettings(project) + fun getQClient(project: Project, profile: QRegionProfile?, sdkClass: KClass): T { + val settings = getQClientSettings(project, profile) val client = AwsClientManager.getInstance().getClient(sdkClass, settings) return client } diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/customization/CodeWhispererCustomization.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/customization/CodeWhispererCustomization.kt index 92335fec559..1ca14e6e025 100644 --- a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/customization/CodeWhispererCustomization.kt +++ b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/customization/CodeWhispererCustomization.kt @@ -2,6 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 package software.aws.toolkits.jetbrains.services.codewhisperer.customization +import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfile data class CodeWhispererCustomization( @JvmField @@ -12,4 +13,7 @@ data class CodeWhispererCustomization( @JvmField var description: String? = null, + + @JvmField + var profile: QRegionProfile? = null, ) From c9aba71290bfcc28051a39ba44921a09b0557a90 Mon Sep 17 00:00:00 2001 From: Richard Li <742829+rli@users.noreply.github.com> Date: Tue, 6 May 2025 15:02:40 -0700 Subject: [PATCH 2/3] fix(amazonq): align q server location to VSC (#5490) location was off by a single folder --- .../jetbrains/services/amazonq/lsp/artifacts/ArtifactHelper.kt | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/artifacts/ArtifactHelper.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/artifacts/ArtifactHelper.kt index 8787259bf08..42c14f59517 100644 --- a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/artifacts/ArtifactHelper.kt +++ b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/artifacts/ArtifactHelper.kt @@ -19,12 +19,13 @@ import software.aws.toolkits.jetbrains.core.saveFileFromUrl import software.aws.toolkits.jetbrains.services.amazonq.project.manifest.ManifestManager import software.aws.toolkits.resources.AwsCoreBundle import java.nio.file.Path +import java.nio.file.Paths import java.util.concurrent.atomic.AtomicInteger class ArtifactHelper(private val lspArtifactsPath: Path = DEFAULT_ARTIFACT_PATH, private val maxDownloadAttempts: Int = MAX_DOWNLOAD_ATTEMPTS) { companion object { - private val DEFAULT_ARTIFACT_PATH = getToolkitsCommonCacheRoot().resolve("aws").resolve("toolkits").resolve("language-servers") + private val DEFAULT_ARTIFACT_PATH = getToolkitsCommonCacheRoot().resolve(Paths.get("aws", "toolkits", "language-servers", "AmazonQ")) private val logger = getLogger() private const val MAX_DOWNLOAD_ATTEMPTS = 3 } From 95b797c8669223e2ee74c4154db91116ab11462a Mon Sep 17 00:00:00 2001 From: Richard Li Date: Tue, 6 May 2025 16:36:15 -0700 Subject: [PATCH 3/3] lint --- .../services/codewhisperer/CodeWhispererClientAdaptorTest.kt | 1 - .../customization/CodeWhispererModelConfigurator.kt | 1 - 2 files changed, 2 deletions(-) diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererClientAdaptorTest.kt b/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererClientAdaptorTest.kt index 1f4428136d4..d8745123336 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererClientAdaptorTest.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererClientAdaptorTest.kt @@ -54,7 +54,6 @@ import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection import software.aws.toolkits.jetbrains.core.credentials.sono.Q_SCOPES import software.aws.toolkits.jetbrains.core.credentials.sono.SONO_REGION import software.aws.toolkits.jetbrains.services.amazonq.FEATURE_EVALUATION_PRODUCT_NAME -import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfile import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.metadata import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.sdkHttpResponse import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererClientAdaptorImpl diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/customization/CodeWhispererModelConfigurator.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/customization/CodeWhispererModelConfigurator.kt index 734f13eef3b..e968aaab481 100644 --- a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/customization/CodeWhispererModelConfigurator.kt +++ b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/customization/CodeWhispererModelConfigurator.kt @@ -119,7 +119,6 @@ class DefaultCodeWhispererModelConfigurator : CodeWhispererModelConfigurator, Pe ) } - @RequiresBackgroundThread override fun listCustomizations(project: Project, passive: Boolean): List? = calculateIfIamIdentityCenterConnection(project) {