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

## Problem
With the old VSCode auth, customization was stored in global state
(indexed by connection ID). With the new auth from the Flare identity
server, this no longer applies since there is only a single connection

## Solution
* Consume the new AuthUtil in the customizationUtil
* Change `CODEWHISPERER_SELECTED_CUSTOMIZATION` in the global state to
be a single value instead of a map

## Follow-up
Investigate if it's possible to hook up the customization to the Q
profile from the new Region Manager, so customizations are persisted
between different logins

---

- Treat all work as PUBLIC. Private `feature/x` branches will not be
squash-merged at release time.
- Your code changes must meet the guidelines in
[CONTRIBUTING.md](https://github.com/aws/aws-toolkit-vscode/blob/master/CONTRIBUTING.md#guidelines).
- License: I confirm that my contribution is made under the terms of the
Apache 2.0 license.

Author: opieter-aws

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`.
      *