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

Author: opieter-aws

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

@@ -94,20 +94,14 @@ 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 }>(
+    const selectedCustomization = globals.globalState.tryGet<Customization>(
         'CODEWHISPERER_SELECTED_CUSTOMIZATION',
-        Object,
-        {}
+        Object
     )
-    const selectedCustomization = selectedCustomizationArr[AuthUtil.instance.conn.label]
 
     if (selectedCustomization && selectedCustomization.name !== '') {
         return selectedCustomization
@@ -124,7 +118,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 +127,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,27 +139,27 @@ 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 persistedCustomizationsObj[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
+    persistedCustomizationsObj[AuthUtil.instance.profileName] = customizations
     await globals.globalState.update('CODEWHISPERER_PERSISTED_CUSTOMIZATIONS', persistedCustomizationsObj)
 }