ab

Author: Will-ShaoHua

packages/core/src/shared/featureConfig.ts

@@ -4,7 +4,6 @@
  */
 
 import {
-    Customization,
     FeatureValue,
     ListFeatureEvaluationsRequest,
     ListFeatureEvaluationsResponse,
@@ -21,7 +20,7 @@ import { getClientId, getOperatingSystem } from './telemetry/util'
 import { extensionVersion } from './vscode/env'
 import { telemetry } from './telemetry/telemetry'
 import { Commands } from './vscode/commands2'
-import { setSelectedCustomization } from '../codewhisperer/util/customizationUtil'
+import { getAvailableCustomizationsList, setSelectedCustomization } from '../codewhisperer/util/customizationUtil'
 
 const localize = nls.loadMessageBundle()
 
@@ -145,29 +144,15 @@ export class FeatureConfigProvider {
             }
             getLogger().info('AB Testing Cohort Assignments %O', response.featureEvaluations)
 
-            const profile = AuthUtil.instance.regionProfileManager.activeRegionProfile
             const customizationArnOverride = this.featureConfigs.get(Features.customizationArnOverride)?.value
                 ?.stringValue
             const previousOverride = globals.globalState.tryGet<string>('aws.amazonq.customization.overrideV2', String)
-            if (profile && customizationArnOverride !== undefined && customizationArnOverride !== previousOverride) {
+            if (customizationArnOverride !== undefined && customizationArnOverride !== previousOverride) {
                 // Double check if server-side wrongly returns a customizationArn to BID users
                 if (isBuilderIdConnection(AuthUtil.instance.conn)) {
                     this.featureConfigs.delete(Features.customizationArnOverride)
                 } else if (isIdcSsoConnection(AuthUtil.instance.conn)) {
-                    let availableCustomizations: Customization[] = []
-                    try {
-                        const items: Customization[] = []
-                        // TODO: list all customization across different profiles
-                        const response = await client.listAvailableCustomizations(profile)
-                        for (const customizations of response.map(
-                            (listAvailableCustomizationsResponse) => listAvailableCustomizationsResponse.customizations
-                        )) {
-                            items.push(...customizations)
-                        }
-                        availableCustomizations = items
-                    } catch (e) {
-                        getLogger().debug('amazonq: Failed to list available customizations')
-                    }
+                    const availableCustomizations = await getAvailableCustomizationsList()
 
                     // If customizationArn from A/B is not available in listAvailableCustomizations response, don't use this value
                     const targetCustomization = availableCustomizations?.find((c) => c.arn === customizationArnOverride)
@@ -178,6 +163,16 @@ export class FeatureConfigProvider {
                         this.featureConfigs.delete(Features.customizationArnOverride)
                     } else {
                         await setSelectedCustomization(targetCustomization, true)
+                        // note that we should also switch profile if either
+                        // 1. user has not selected a profile yet
+                        // 2. user's selected profile is not the same as the one of customizationOverride
+                        const profile = AuthUtil.instance.regionProfileManager.activeRegionProfile
+                        if (!profile || (profile && profile.arn !== targetCustomization.profile.arn)) {
+                            await AuthUtil.instance.regionProfileManager.switchRegionProfile(
+                                targetCustomization.profile,
+                                'customization'
+                            )
+                        }
                     }
 
                     await vscode.commands.executeCommand('aws.amazonq.refreshStatusBar')