feat(core): Store codewhisperer customization in global state as single value and consume new authUtil

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

@@ -94,20 +94,11 @@ export const baseCustomization = {
  * @returns customization selected by users, `baseCustomization` if none is selected
  */
 export const getSelectedCustomization = (): Customization => {
-    if (
-        !AuthUtil.instance.isCustomizationFeatureEnabled ||
-        !AuthUtil.instance.isValidEnterpriseSsoInUse() ||
-        !AuthUtil.instance.conn
-    ) {
+    if (!AuthUtil.instance.isCustomizationFeatureEnabled || !AuthUtil.instance.isIdcConnection()) {
         return baseCustomization
     }
 
-    const selectedCustomizationArr = globals.globalState.tryGet<{ [label: string]: Customization }>(
-        'CODEWHISPERER_SELECTED_CUSTOMIZATION',
-        Object,
-        {}
-    )
-    const selectedCustomization = selectedCustomizationArr[AuthUtil.instance.conn.label]
+    const selectedCustomization = globals.globalState.getAmazonQCustomization(AuthUtil.instance.profileName)
 
     if (selectedCustomization && selectedCustomization.name !== '') {
         return selectedCustomization
@@ -124,7 +115,7 @@ export const getSelectedCustomization = (): Customization => {
  *  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) {
+    if (!AuthUtil.instance.isIdcConnection()) {
         return
     }
     if (isOverride) {
@@ -133,15 +124,10 @@ export const setSelectedCustomization = async (customization: Customization, isO
             return
         }
     }
-    const selectedCustomizationObj = globals.globalState.tryGet<{ [label: string]: Customization }>(
-        'CODEWHISPERER_SELECTED_CUSTOMIZATION',
-        Object,
-        {}
-    )
-    selectedCustomizationObj[AuthUtil.instance.conn.label] = customization
-    getLogger().debug(`Selected customization ${customization.name} for ${AuthUtil.instance.conn.label}`)
 
-    await globals.globalState.update('CODEWHISPERER_SELECTED_CUSTOMIZATION', selectedCustomizationObj)
+    await globals.globalState.update('CODEWHISPERER_SELECTED_CUSTOMIZATION', customization)
+    getLogger().debug(`Selected customization ${customization.name} for ${AuthUtil.instance.profileName}`)
+
     if (isOverride) {
         await globals.globalState.update('aws.amazonq.customization.overrideV2', customization.arn)
     }
@@ -150,28 +136,18 @@ export const setSelectedCustomization = async (customization: Customization, isO
 }
 
 export const getPersistedCustomizations = (): Customization[] => {
-    if (!AuthUtil.instance.isValidEnterpriseSsoInUse() || !AuthUtil.instance.conn) {
+    if (!AuthUtil.instance.isIdcConnection()) {
         return []
     }
-    const persistedCustomizationsObj = globals.globalState.tryGet<{ [label: string]: Customization[] }>(
-        'CODEWHISPERER_PERSISTED_CUSTOMIZATIONS',
-        Object,
-        {}
-    )
-    return persistedCustomizationsObj[AuthUtil.instance.conn.label] || []
+    return globals.globalState.getAmazonQCachedCustomization(AuthUtil.instance.profileName)
 }
 
 export const setPersistedCustomizations = async (customizations: Customization[]) => {
-    if (!AuthUtil.instance.isValidEnterpriseSsoInUse() || !AuthUtil.instance.conn) {
+    if (!AuthUtil.instance.isIdcConnection()) {
         return
     }
-    const persistedCustomizationsObj = globals.globalState.tryGet<{ [label: string]: Customization[] }>(
-        'CODEWHISPERER_PERSISTED_CUSTOMIZATIONS',
-        Object,
-        {}
-    )
-    persistedCustomizationsObj[AuthUtil.instance.conn.label] = customizations
-    await globals.globalState.update('CODEWHISPERER_PERSISTED_CUSTOMIZATIONS', persistedCustomizationsObj)
+
+    await globals.globalState.update('CODEWHISPERER_PERSISTED_CUSTOMIZATIONS', customizations)
 }
 
 export const getNewCustomizationsAvailable = () => {

packages/core/src/shared/globalState.ts

@@ -7,6 +7,7 @@ import * as vscode from 'vscode'
 import { getLogger } from './logger/logger'
 import * as redshift from '../awsService/redshift/models/models'
 import { TypeConstructor, cast } from './utilities/typeConstructors'
+import { Customization } from '../codewhisperer/client/codewhispereruserclient'
 
 type ToolId = 'codecatalyst' | 'codewhisperer' | 'testId'
 export type ToolIdStateKey = `${ToolId}.savedConnectionId`
@@ -227,6 +228,53 @@ export class GlobalState implements vscode.Memento {
         return all?.[warehouseArn]
     }
 
+    /**
+     * Get the Amazon Q customization. If legacy (map of customizations) store the
+     * customization with label of profile name
+     *
+     * @param profileName name of profile, only used in case legacy customization is found
+     * @returns Amazon Q customization, or undefined if not found.
+     * If legacy, return the Amazon Q customization for the auth profile name
+     */
+    getAmazonQCustomization(profileName: string): Customization | undefined {
+        const result = this.tryGet('CODEWHISPERER_SELECTED_CUSTOMIZATION', Object, undefined)
+
+        // Legacy migration for old customization map of type { [label: string]: Customization[] }
+        if (typeof result === 'object' && Object.values(result).every(Array.isArray)) {
+            const selectedCustomization = result[profileName]
+            this.tryUpdate('CODEWHISPERER_SELECTED_CUSTOMIZATION', selectedCustomization)
+            return selectedCustomization
+        } else {
+            return result
+        }
+    }
+
+    /**
+     * Get the Amazon Q cached customizations. If legacy (map of customizations) store the
+     * customizations with label of profile name
+     *
+     * @param profileName name of profile, only used in case legacy customization is found
+     * @returns array of Amazon Q cached customizations, or empty array if not found.
+     * If legacy, return the Amazon Q persisted customizations for the auth profile name
+     */
+    getAmazonQCachedCustomization(profileName: string): Customization[] {
+        const result = this.tryGet<Customization[]>('CODEWHISPERER_PERSISTED_CUSTOMIZATIONS', Array, [])
+
+        // Legacy migration for old customization map of type { [label: string]: Customization[] }
+        if (result.length === 0) {
+            const customizations = this.tryGet<{ [label: string]: Customization[] }>(
+                'CODEWHISPERER_PERSISTED_CUSTOMIZATIONS',
+                Object,
+                {}
+            )
+            const cachedCustomizationsArray = customizations[profileName] || []
+            this.tryUpdate('CODEWHISPERER_PERSISTED_CUSTOMIZATIONS', cachedCustomizationsArray)
+            return cachedCustomizationsArray
+        } else {
+            return result
+        }
+    }
+
     /**
      * Sets SSO session creation timestamp for the given session `id`.
      *