Add Customization A/B

Use the customization Arn value if and only if:

user is an IdC user
arn is in the list of available ones

packages/core/src/codewhisperer/service/featureConfigProvider.ts

@@ -7,17 +7,24 @@ import { FeatureValue } from '../client/codewhispereruserclient'
 import { codeWhispererClient as client } from '../client/codewhisperer'
 import { AuthUtil } from '../util/authUtil'
 import { getLogger } from '../../shared/logger'
+import { isBuilderIdConnection, isIdcSsoConnection } from '../../auth/connection'
+import { getAvailableCustomizationsList } from '../util/customizationUtil'
 
 export class FeatureContext {
     constructor(public name: string, public variation: string, public value: FeatureValue) {}
 }
 
 const testFeatureName = 'testFeature'
+const customizationArnOverrideName = 'customizationArnOverride'
 const featureConfigPollIntervalInMs = 30 * 60 * 1000 // 30 mins
 
 // TODO: add real feature later
 export const featureDefinitions = new Map([
     [testFeatureName, new FeatureContext(testFeatureName, 'CONTROL', { stringValue: 'testValue' })],
+    [
+        customizationArnOverrideName,
+        new FeatureContext(customizationArnOverrideName, 'customizationARN', { stringValue: '' }),
+    ],
 ])
 
 export class FeatureConfigProvider {
@@ -53,6 +60,29 @@ export class FeatureConfigProvider {
                     new FeatureContext(evaluation.feature, evaluation.variation, evaluation.value)
                 )
             })
+
+            const customizationArnOverride = this.featureConfigs.get(customizationArnOverrideName)?.value?.stringValue
+            if (customizationArnOverride !== undefined) {
+                // Double check if server-side wrongly returns a customizationArn to BID users
+                if (isBuilderIdConnection(AuthUtil.instance.conn)) {
+                    this.featureConfigs.delete(customizationArnOverrideName)
+                } else if (isIdcSsoConnection(AuthUtil.instance.conn)) {
+                    let availableCustomizations = null
+                    try {
+                        availableCustomizations = (await getAvailableCustomizationsList()).map(c => c.arn)
+                    } catch (e) {
+                        getLogger().debug('amazonq: Failed to list available customizations')
+                    }
+
+                    // If customizationArn from A/B is not available in listAvailableCustomizations response, don't use this value
+                    if (!availableCustomizations?.includes(customizationArnOverride)) {
+                        getLogger().debug(
+                            `Customization arn ${customizationArnOverride} not available in listAvailableCustomizations, not using`
+                        )
+                        this.featureConfigs.delete(customizationArnOverrideName)
+                    }
+                }
+            }
         } catch (e) {
             getLogger().error(`CodeWhisperer: Error when fetching feature configs ${e}`, e)
         }
@@ -82,6 +112,10 @@ export class FeatureConfigProvider {
         return this.getFeatureValueForKey(testFeatureName).stringValue
     }
 
+    getCustomizationArnOverride(): string | undefined {
+        return this.getFeatureValueForKey(customizationArnOverrideName).stringValue
+    }
+
     // Get the feature value for the given key.
     // In case of a misconfiguration, it will return a default feature value of Boolean true.
     private getFeatureValueForKey(name: string): FeatureValue {

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

@@ -25,6 +25,7 @@ import { showMessageWithUrl } from '../../shared/utilities/messages'
 import { parse } from '@aws-sdk/util-arn-parser'
 import { Commands } from '../../shared/vscode/commands2'
 import { vsCodeState } from '../models/model'
+import { FeatureConfigProvider } from '../service/featureConfigProvider'
 
 /**
  *
@@ -108,7 +109,21 @@ export const getSelectedCustomization = (): Customization => {
 
     const selectedCustomizationArr =
         globals.context.globalState.get<{ [label: string]: Customization }>(selectedCustomizationKey) || {}
-    return selectedCustomizationArr[AuthUtil.instance.conn.label] || baseCustomization
+    const result = selectedCustomizationArr[AuthUtil.instance.conn.label] || baseCustomization
+
+    // A/B case
+    const arnOverride = FeatureConfigProvider.instance.getCustomizationArnOverride()
+    if (arnOverride === undefined || arnOverride === '') {
+        return result
+    } else {
+        // A trick to prioritize arn from A/B over user's currently selected(for request and telemetry)
+        // but still shows customization info of user's currently selected.
+        return {
+            arn: arnOverride,
+            name: result.name,
+            description: result.description,
+        }
+    }
 }
 
 export const setSelectedCustomization = async (customization: Customization) => {