remove composite key & cr

Author: evanliu048

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

@@ -109,6 +109,8 @@ 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,

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

@@ -27,7 +27,6 @@ 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
@@ -65,13 +64,6 @@ private fun notifyNewCustomization(project: Project) {
     )
 }
 
-/**
- * generate a stable composite key for profile & customization combination for diff/deduplicate/search
- * rule: `<profileArn>::<customizationArn>`
- */
-private fun CodeWhispererCustomization.compositeKey(): String =
-    "${profile?.arn.orEmpty()}::$arn"
-
 @Service(Service.Level.APP)
 @State(name = "codewhispererCustomizationStates", storages = [Storage("aws.xml")])
 class DefaultCodeWhispererModelConfigurator : CodeWhispererModelConfigurator, PersistentStateComponent<CodeWhispererCustomizationState>, Disposable {
@@ -80,7 +72,6 @@ class DefaultCodeWhispererModelConfigurator : CodeWhispererModelConfigurator, Pe
     private val connectionIdToActiveCustomizationArn = Collections.synchronizedMap<String, CodeWhispererCustomization>(mutableMapOf())
 
     // Map to store connectionId to its listAvailableCustomizations result last time
-    // Format of the customization key: profileArn::customizationArn
     private val connectionToCustomizationsShownLastTime = mutableMapOf<String, MutableList<String>>()
 
     private val connectionIdToIsAllowlisted = Collections.synchronizedMap<String, Boolean>(mutableMapOf())
@@ -126,23 +117,15 @@ class DefaultCodeWhispererModelConfigurator : CodeWhispererModelConfigurator, Pe
                     CodeWhispererClientAdaptor.getInstance(project).listAvailableCustomizations(profile)
                 }.onFailure { e ->
                     val requestId = (e as? CodeWhispererRuntimeException)?.requestId()
-                    val isNotAllowlisted = (e as? Exception)
-                        ?.let { CodeWhispererConstants.Customization.noAccessToCustomizationExceptionPredicate(e) }
-                        ?: false
-                    val logMessage = if (isNotAllowlisted) {
-                        // 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()}"
-                    }
+                    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() }
-            // calculate new profile+customization list using compositeKey
-            val diff = aggregatedCustomizations.filterNot { customization -> previousCustomizationsShapshot.contains(customization.compositeKey()) }.toSet()
+            val diff = aggregatedCustomizations.filterNot { customization -> previousCustomizationsShapshot.contains(customization.arn) }.toSet()
 
             // 3 if passive,
             //   (1) update allowlisting
@@ -156,18 +139,18 @@ class DefaultCodeWhispererModelConfigurator : CodeWhispererModelConfigurator, Pe
                     notifyNewCustomization(project)
                 }
             } else {
-                connectionToCustomizationsShownLastTime[it.id] = aggregatedCustomizations.map { customization -> customization.compositeKey() }.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 q region profile associated with the selected customization does not match the currently active profile
+            //    (3) the existing q region profile associated with the selected customization does not match the currently active profile
             activeCustomization(project)?.let { activeCustom ->
                 if (aggregatedCustomizations.isEmpty()) {
                     invalidateSelectedAndNotify(project)
-                } else if (!aggregatedCustomizations.any { latestCustom -> latestCustom.compositeKey() == activeCustom.compositeKey() } ||
-                    activeCustom.profile != QRegionProfileManager.getInstance().activeProfile(project)
+                } else if (!aggregatedCustomizations.any { latestCustom -> latestCustom.arn == activeCustom.arn } ||
+                    (activeCustom.profile != null && activeCustom.profile != QRegionProfileManager.getInstance().activeProfile(project))
                 ) {
                     invalidateSelectedAndNotify(project)
                 }

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

@@ -8,7 +8,7 @@ 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
+ * 'customization' -> users selected a customization tied to a different profile, triggering a profile switch
  */
 enum class QProfileSwitchIntent(val value: String) {
     User("user"),

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

@@ -182,7 +182,7 @@ class QRegionProfileManager : PersistentStateComponent<QProfileState>, Disposabl
     }
 
     inline fun <reified T : SdkClient> getQClient(project: Project): T = getQClient(project, null, T::class)
-    inline fun <reified T : SdkClient> getQClient(project: Project, profile: QRegionProfile?): T = getQClient(project, profile, T::class)
+    inline fun <reified T : SdkClient> getQClient(project: Project, profile: QRegionProfile): T = getQClient(project, profile, T::class)
 
     fun <T : SdkClient> getQClient(project: Project, profile: QRegionProfile?, sdkClass: KClass<T>): T {
         val settings = getQClientSettings(project, profile)