fix(core): Add globalState migration for CODEWHISPERER_PERSISTED_CUSTOMIZATIONS

Author: opieter-aws

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

@@ -139,25 +139,15 @@ export const getPersistedCustomizations = (): Customization[] => {
     if (!AuthUtil.instance.isIdcConnection()) {
         return []
     }
-    const persistedCustomizationsObj = globals.globalState.tryGet<{ [label: string]: Customization[] }>(
-        'CODEWHISPERER_PERSISTED_CUSTOMIZATIONS',
-        Object,
-        {}
-    )
-    return persistedCustomizationsObj[AuthUtil.instance.profileName] || []
+    return globals.globalState.getCodewhispererPersistedCustomization(AuthUtil.instance.profileName)
 }
 
 export const setPersistedCustomizations = async (customizations: Customization[]) => {
     if (!AuthUtil.instance.isIdcConnection()) {
         return
     }
-    const persistedCustomizationsObj = globals.globalState.tryGet<{ [label: string]: Customization[] }>(
-        'CODEWHISPERER_PERSISTED_CUSTOMIZATIONS',
-        Object,
-        {}
-    )
-    persistedCustomizationsObj[AuthUtil.instance.profileName] = 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

@@ -229,7 +229,7 @@ export class GlobalState implements vscode.Memento {
     }
 
     /**
-     * Get the codewhisperer customerization. If legacy (map of customizations) store the
+     * Get the codewhisperer 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
@@ -249,6 +249,32 @@ export class GlobalState implements vscode.Memento {
         }
     }
 
+    /**
+     * Get the codewhisperer persisted 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 codewhisperer persisted customizations, or empty array if not found.
+     * If legacy, return the codewhisperer persisted customizations for the auth profile name
+     */
+    getCodewhispererPersistedCCustomization(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 persistedCustomizationsArray = customizations[profileName] || []
+            this.tryUpdate('CODEWHISPERER_PERSISTED_CUSTOMIZATIONS', persistedCustomizationsArray)
+            return persistedCustomizationsArray
+        } else {
+            return result
+        }
+    }
+
     /**
      * Sets SSO session creation timestamp for the given session `id`.
      *