customization override should only happen once for a specific override

Author: Will-ShaoHua

packages/core/src/codewhisperer/util/customizationUtil.ts

@@ -91,10 +91,7 @@ export const baseCustomization = {
 }
 
 /**
- * Gets the customization that should be used for user requests. If a user has manually selected
- * a customization, always respect that choice. If not, check if the user is part of an AB
- * group assigned a specific customization. If so, use that customization. If not, use the
- * base customization.
+ * @returns customization selected by users, `baseCustomization` if none is selected
  */
 export const getSelectedCustomization = (): Customization => {
     if (
@@ -119,10 +116,23 @@ export const getSelectedCustomization = (): Customization => {
     }
 }
 
-export const setSelectedCustomization = async (customization: Customization) => {
+/**
+ * @param customization customization to select
+ * @param isOverride if the API call is made from us (Q) but not users' intent, set isOverride to TRUE
+ * Override happens when ALL following conditions are met
+ *  1. service returns non-empty override customization arn, refer to [featureConfig.ts]
+ *  2. the override customization arn is different from the previous override customization if any. The purpose is to only do override once on users' behalf.
+ */
+export const setSelectedCustomization = async (customization: Customization, isOverride: boolean = false) => {
     if (!AuthUtil.instance.isValidEnterpriseSsoInUse() || !AuthUtil.instance.conn) {
         return
     }
+    if (isOverride) {
+        const previousOverride = globals.globalState.tryGet<string>('aws.amazonq.customization.override', String)
+        if (customization.arn === previousOverride) {
+            return
+        }
+    }
     const selectedCustomizationObj = globals.globalState.tryGet<{ [label: string]: Customization }>(
         'CODEWHISPERER_SELECTED_CUSTOMIZATION',
         Object,
@@ -132,6 +142,9 @@ export const setSelectedCustomization = async (customization: Customization) =>
     getLogger().debug(`Selected customization ${customization.name} for ${AuthUtil.instance.conn.label}`)
 
     await globals.globalState.update('CODEWHISPERER_SELECTED_CUSTOMIZATION', selectedCustomizationObj)
+    if (isOverride) {
+        await globals.globalState.update('aws.amazonq.customization.override', customization.arn)
+    }
     vsCodeState.isFreeTierLimitReached = false
     await Commands.tryExecute('aws.amazonq.refreshStatusBar')
 }

packages/core/src/shared/globalState.ts

@@ -45,6 +45,7 @@ export type globalKey =
     | 'aws.toolkit.amazonq.dismissed'
     | 'aws.toolkit.amazonqInstall.dismissed'
     | 'aws.amazonq.workspaceIndexToggleOn'
+    | 'aws.amazonq.customization.override'
     // Deprecated/legacy names. New keys should start with "aws.".
     | '#sessionCreationDates' // Legacy name from `ssoAccessTokenProvider.ts`.
     | 'CODECATALYST_RECONNECT'