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

.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
         )