diff --git a/packages/core/src/codewhisperer/util/customizationUtil.ts b/packages/core/src/codewhisperer/util/customizationUtil.ts index 7898b7a17ba..e9ec68988fd 100644 --- a/packages/core/src/codewhisperer/util/customizationUtil.ts +++ b/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 = () => { diff --git a/packages/core/src/shared/globalState.ts b/packages/core/src/shared/globalState.ts index 44d848ec69d..23a67667d88 100644 --- a/packages/core/src/shared/globalState.ts +++ b/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('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`. *