fix(amazonq): support displaying customizations across all profiles of an Idc connection (#5591)

Follow-up to #5575.

As mentioned in that PR:

Before: Customizations were bound to a specific IDC instance.
After: Customizations are now bound to a specific Q profile. An IDC instance can have multiple Q profiles.

In other words, each Q profile has access to its own set of customizations.

The product team has requested that we show all available customizations across profiles, rather than only those tied to the currently connected profile.

This means that when a user selects a customization that belongs to a different Q profile, the plugin should implicitly switch to that profile. This behavior helps reduce confusion and user churn, since users might not clearly understand which profile grants access to which customizations.

Author: evanliu048

.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/credentials/CodeWhispererClientAdaptor.kt

@@ -38,6 +38,7 @@ import software.amazon.awssdk.services.codewhispererruntime.model.UserIntent
 import software.aws.toolkits.core.utils.debug
 import software.aws.toolkits.core.utils.getLogger
 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
@@ -78,7 +79,7 @@ interface CodeWhispererClientAdaptor {
 
     fun getCodeFixJob(request: GetCodeFixJobRequest): GetCodeFixJobResponse
 
-    fun listAvailableCustomizations(): List<CodeWhispererCustomization>
+    fun listAvailableCustomizations(profile: QRegionProfile): List<CodeWhispererCustomization>
 
     fun startTestGeneration(uploadId: String, targetCode: List<TargetCode>, userInput: String): StartTestGenerationResponse
 
@@ -282,9 +283,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<CodeWhispererCustomization> =
-        bearerClient().listAvailableCustomizationsPaginator(
-            ListAvailableCustomizationsRequest.builder().profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn).build()
+    override fun listAvailableCustomizations(profile: QRegionProfile): List<CodeWhispererCustomization> =
+        QRegionProfileManager.getInstance().getQClient<CodeWhispererRuntimeClient>(project, profile).listAvailableCustomizationsPaginator(
+            ListAvailableCustomizationsRequest.builder().profileArn(profile.arn).build()
         )
             .stream()
             .toList()
@@ -298,7 +299,8 @@ open class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeW
                     CodeWhispererCustomization(
                         arn = it.arn(),
                         name = it.name(),
-                        description = it.description()
+                        description = it.description(),
+                        profile = profile
                     )
                 }
             }

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

@@ -26,6 +26,9 @@ 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.QProfileSwitchIntent
+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.util.CodeWhispererConstants.Q_CUSTOM_LEARN_MORE_URI
 import software.aws.toolkits.jetbrains.ui.AsyncComboBox
 import software.aws.toolkits.jetbrains.utils.notifyInfo
@@ -34,7 +37,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
 )
@@ -106,6 +109,15 @@ class CodeWhispererCustomizationDialog(
             RadioButtonOption.Customization -> run {
                 CodeWhispererModelConfigurator.getInstance().switchCustomization(project, modal.selectedCustomization?.customization)
                 notifyCustomizationIsSelected(project, modal.selectedCustomization)
+                // Switch profile if it doesn't match the customization's profile.
+                // Customizations are profile-scoped and must be used under the correct context.
+                if (modal.selectedCustomization?.customization?.profile?.arn != QRegionProfileManager.getInstance().activeProfile(project)?.arn) {
+                    QRegionProfileManager.getInstance().switchProfile(
+                        project,
+                        modal.selectedCustomization?.customization?.profile,
+                        QProfileSwitchIntent.Customization
+                    )
+                }
             }
         }
 
@@ -179,12 +191,14 @@ class CodeWhispererCustomizationDialog(
                         proposeModelUpdate { model ->
                             val activeCustomization = CodeWhispererModelConfigurator.getInstance().activeCustomization(project)
                             val unsorted = myCustomizations ?: CodeWhispererModelConfigurator.getInstance().listCustomizations(project).orEmpty()
-
-                            val sorted = activeCustomization?.let {
-                                unsorted.putPickedUpFront(setOf(it))
-                            } ?: run {
-                                unsorted.sortedBy { it.customization.name }
-                            }
+                            val activeProfile = QRegionProfileManager.getInstance().activeProfile(project)
+                            // Group customizations by profile name (active profile first, then alphabetical), with the active customization on top
+                            val sorted = unsorted.sortedWith(
+                                compareBy<CustomizationUiItem> { it.customization.profile?.profileName != activeProfile?.profileName }
+                                    .thenBy { it.customization.profile?.profileName.orEmpty() }
+                                    .thenBy { it.customization.name }
+                            )
+                                .let { list -> activeCustomization?.let { list.putPickedUpFront(setOf(it)) } ?: list }
 
                             if (
                                 sorted.isNotEmpty() &&
@@ -259,6 +273,10 @@ private object CustomizationRenderer : ColoredListCellRenderer<CustomizationUiIt
                 }
             }
 
+            if (it.customization.profile?.profileName?.isNotEmpty() == true) {
+                append("  [${it.customization.profile?.profileName}]", SimpleTextAttributes.REGULAR_ATTRIBUTES)
+            }
+
             if (it.isNew) {
                 append("  New", SimpleTextAttributes.GRAYED_SMALL_ATTRIBUTES)
             }

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

@@ -23,9 +23,9 @@ 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.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.util.CodeWhispererConstants
 import software.aws.toolkits.jetbrains.utils.notifyInfo
 import software.aws.toolkits.jetbrains.utils.notifyWarn
 import software.aws.toolkits.jetbrains.utils.pluginAwareExecuteOnPooledThread
@@ -108,25 +108,24 @@ class DefaultCodeWhispererModelConfigurator : CodeWhispererModelConfigurator, Pe
     @RequiresBackgroundThread
     override fun listCustomizations(project: Project, passive: Boolean): List<CustomizationUiItem>? =
         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 listAvailableProfilesResult = QRegionProfileManager.getInstance().listRegionProfiles(project)
+                ?: error("Attempted to fetch profiles while there does not exist")
+
+            val aggregatedCustomizations = listAvailableProfilesResult.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 = aggregatedCustomizations.filterNot { customization -> previousCustomizationsShapshot.contains(customization.arn) }.toSet()
 
             // 3 if passive,
             //   (1) update allowlisting
@@ -135,42 +134,36 @@ 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] = aggregatedCustomizations.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] = aggregatedCustomizations.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 (aggregatedCustomizations.isEmpty()) {
                     invalidateSelectedAndNotify(project)
-                } else if (!listAvailableCustomizationsResult.any { latestCustom -> latestCustom.arn == activeCustom.arn }) {
+                } else if (!aggregatedCustomizations.any { latestCustom -> latestCustom.arn == activeCustom.arn } ||
+                    (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 = aggregatedCustomizations.groupingBy { customization -> customization.name }.eachCount()
+            val customizationUiItems = aggregatedCustomizations.map { customization ->
+                CustomizationUiItem(
+                    customization,
+                    isNew = diff.contains(customization),
+                    shouldPrefixAccountId = (nameToCount[customization.name] ?: 0) > 1
+                )
             }
             connectionToCustomizationUiItems[it.id] = customizationUiItems
 

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<ListAvailableCustomizationsRequest>()) } 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"))
             )
         )
     }

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
         )