customization override

Author: Will-ShaoHua

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/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 }
         }
     }