refactor(codewhisperer): use globalState abstraction

Author: justinmk3

packages/core/src/amazonq/util/viewBadgeHandler.ts

@@ -71,8 +71,7 @@ async function showInitialViewBadge() {
  * @returns True if the badge should be shown, false otherwise
  */
 export async function shouldShowBadge(): Promise<boolean> {
-    const memento = globals.context.globalState
-    const hasAlreadyShown = memento.get('hasAlreadyOpenedAmazonQ')
+    const hasAlreadyShown = globals.globalState.get('hasAlreadyOpenedAmazonQ')
     if (!hasAlreadyShown) {
         const state = await AuthUtil.instance.getChatAuthState()
         if (state.codewhispererCore === 'connected' && state.codewhispererChat !== 'connected') {

packages/core/src/auth/auth.ts

@@ -302,7 +302,7 @@ export class Auth implements AuthService, ConnectionManager {
 
         // Remove the split session logout prompt, if it exists.
         if (!isAmazonQ()) {
-            await globals.context.globalState.update(SessionSeparationPrompt.instance.dismissKey, true)
+            await globals.globalState.update('aws.toolkit.separationPromptDismissed', true)
         }
 
         try {
@@ -1029,9 +1029,6 @@ type LoginCommand = 'aws.toolkit.auth.manageConnections' | 'aws.codecatalyst.man
  * TODO: Remove after some time.
  */
 export class SessionSeparationPrompt {
-    public readonly dismissKey = 'aws.toolkit.separationPromptDismissed'
-    private readonly loginCmdKey = 'aws.toolkit.separationPromptCommand'
-
     // Local variable handles per session displays, e.g. we forgot a CodeCatalyst connection AND
     // an Explorer only connection. We only want to display once in this case.
     // However, we don't want to set this at the global state level until a user interacts with the
@@ -1043,7 +1040,7 @@ export class SessionSeparationPrompt {
      * which is useful to redisplay the prompt after reloads in case a user misses it.
      */
     public async showAnyPreviousPrompt() {
-        const cmd = globals.context.globalState.get<string>(this.loginCmdKey)
+        const cmd = globals.globalState.tryGet('aws.toolkit.separationPromptCommand', String)
         return cmd ? await this.showForCommand(cmd as LoginCommand) : undefined
     }
 
@@ -1053,11 +1050,14 @@ export class SessionSeparationPrompt {
      * (e.g. codecatalyst sign in vs explorer)
      */
     public async showForCommand(cmd: LoginCommand) {
-        if (this.#separationPromptDisplayed || globals.context.globalState.get<boolean>(this.dismissKey)) {
+        if (
+            this.#separationPromptDisplayed ||
+            globals.globalState.get<boolean>('aws.toolkit.separationPromptDismissed')
+        ) {
             return
         }
 
-        await globals.context.globalState.update(this.loginCmdKey, cmd)
+        await globals.globalState.update('aws.toolkit.separationPromptCommand', cmd)
 
         await telemetry.toolkit_showNotification.run(async () => {
             telemetry.record({ id: 'sessionSeparation' })
@@ -1077,7 +1077,7 @@ export class SessionSeparationPrompt {
                             telemetry.record({ action: 'dismiss' })
                         }
 
-                        await globals.context.globalState.update(this.dismissKey, true)
+                        await globals.globalState.update('aws.toolkit.separationPromptDismissed', true)
                     })
                 })
         })

packages/core/src/awsService/s3/commands/uploadFile.ts

@@ -437,7 +437,7 @@ export async function promptUserForBucket(
         }
     })
 
-    const lastTouchedFolder = globals.context.globalState.get<SavedFolder | undefined>('aws.lastTouchedS3Folder')
+    const lastTouchedFolder = globals.globalState.tryGet<SavedFolder>('aws.lastTouchedS3Folder', Object)
     let lastFolderItem: BucketQuickPickItem | undefined = undefined
     if (lastTouchedFolder) {
         lastFolderItem = {
@@ -448,7 +448,7 @@ export async function promptUserForBucket(
         }
     }
 
-    const lastUploadedToFolder = globals.context.globalState.get<SavedFolder | undefined>('aws.lastUploadedToS3Folder')
+    const lastUploadedToFolder = globals.globalState.tryGet<SavedFolder>('aws.lastUploadedToS3Folder', Object)
     let lastUploadedFolderItem: BucketQuickPickItem | undefined = undefined
     if (lastUploadedToFolder) {
         lastUploadedFolderItem = {

packages/core/src/codewhisperer/index.ts

@@ -37,7 +37,7 @@ export {
 } from './service/transformByQ/transformationResultsViewProvider'
 export { parseVersionsListFromPomFile } from './service/transformByQ/transformFileHandler'
 export { SecurityPanelViewProvider } from './views/securityPanelViewProvider'
-export { get, set, isInlineCompletionEnabled } from './util/commonUtil'
+export { isInlineCompletionEnabled } from './util/commonUtil'
 export { validateOpenProjects, getOpenProjects } from './service/transformByQ/transformProjectValidationHandler'
 export {
     DefaultCodeWhispererClient,

packages/core/src/codewhisperer/models/constants.ts

@@ -134,16 +134,8 @@ export const unsupportedLanguagesCacheTTL = 10 * 60 * 60 * 1000
 
 export const unsupportedLanguagesKey = 'CODEWHISPERER_UNSUPPORTED_LANGUAGES_KEY'
 
-export const autoTriggerEnabledKey = 'CODEWHISPERER_AUTO_TRIGGER_ENABLED'
-
-export const autoScansEnabledKey = 'CODEWHISPERER_AUTO_SCANS_ENABLED'
-
 export const serviceActiveKey = 'CODEWHISPERER_SERVICE_ACTIVE'
 
-export const persistedCustomizationsKey = 'CODEWHISPERER_PERSISTED_CUSTOMIZATIONS'
-
-export const selectedCustomizationKey = 'CODEWHISPERER_SELECTED_CUSTOMIZATION'
-
 export const inlinehintKey = 'CODEWHISPERER_HINT_DISPLAYED'
 
 export type AnnotationChangeSource = 'codewhisperer' | 'selection' | 'editor' | 'content'
@@ -329,8 +321,6 @@ export const updateInlineLockKey = 'CODEWHISPERER_INLINE_UPDATE_LOCK_KEY'
 
 export const newCustomizationMessage = 'You have access to new Amazon Q customizations.'
 
-export const newCustomizationsAvailableKey = 'aws.amazonq.codewhisperer.newCustomizations'
-
 // Start of QCT Strings
 export const uploadZipSizeLimitInBytes = 2000000000 // 2GB
 

packages/core/src/codewhisperer/models/model.ts

@@ -15,8 +15,6 @@ import {
 } from '../../shared/telemetry/telemetry'
 import { References } from '../client/codewhisperer'
 import globals from '../../shared/extensionGlobals'
-import { autoScansEnabledKey, autoTriggerEnabledKey } from './constants'
-import { get, set } from '../util/commonUtil'
 import { ChatControllerEventEmitters } from '../../amazonqGumby/chat/controller/controller'
 import { TransformationSteps } from '../client/codewhispereruserclient'
 
@@ -76,7 +74,6 @@ export interface GetRecommendationsResponse {
 
 /** Manages the state of CodeWhisperer code suggestions */
 export class CodeSuggestionsState {
-    #context: vscode.Memento
     /** The initial state if suggestion state was not defined */
     #fallback: boolean
     #onDidChangeState = new vscode.EventEmitter<boolean>()
@@ -88,15 +85,14 @@ export class CodeSuggestionsState {
         return (this.#instance ??= new this())
     }
 
-    protected constructor(context: vscode.Memento = globals.context.globalState, fallback: boolean = true) {
-        this.#context = context
+    protected constructor(fallback: boolean = true) {
         this.#fallback = fallback
     }
 
     async toggleSuggestions() {
         const autoTriggerEnabled = this.isSuggestionsEnabled()
         const toSet: boolean = !autoTriggerEnabled
-        await set(autoTriggerEnabledKey, toSet, this.#context)
+        await globals.globalState.tryUpdate('CODEWHISPERER_AUTO_TRIGGER_ENABLED', toSet)
         this.#onDidChangeState.fire(toSet)
         return toSet
     }
@@ -108,13 +104,12 @@ export class CodeSuggestionsState {
     }
 
     isSuggestionsEnabled(): boolean {
-        const isEnabled = get(autoTriggerEnabledKey, this.#context)
+        const isEnabled = globals.globalState.tryGet('CODEWHISPERER_AUTO_TRIGGER_ENABLED', Boolean)
         return isEnabled !== undefined ? isEnabled : this.#fallback
     }
 }
 
 export class CodeScansState {
-    #context: vscode.Memento
     /** The initial state if scan state was not defined */
     #fallback: boolean
     #onDidChangeState = new vscode.EventEmitter<boolean>()
@@ -129,15 +124,14 @@ export class CodeScansState {
         return (this.#instance ??= new this())
     }
 
-    protected constructor(context: vscode.Memento = globals.context.globalState, fallback: boolean = true) {
-        this.#context = context
+    protected constructor(fallback: boolean = true) {
         this.#fallback = fallback
     }
 
     async toggleScans() {
         const autoScansEnabled = this.isScansEnabled()
         const toSet: boolean = !autoScansEnabled
-        await set(autoScansEnabledKey, toSet, this.#context)
+        await globals.globalState.tryUpdate('CODEWHISPERER_AUTO_SCANS_ENABLED', toSet)
         this.#onDidChangeState.fire(toSet)
         return toSet
     }
@@ -149,7 +143,7 @@ export class CodeScansState {
     }
 
     isScansEnabled(): boolean {
-        const isEnabled = get(autoScansEnabledKey, this.#context)
+        const isEnabled = globals.globalState.tryGet('CODEWHISPERER_AUTO_SCANS_ENABLED', Boolean)
         return isEnabled !== undefined ? isEnabled : this.#fallback
     }
 

packages/core/src/codewhisperer/ui/codeWhispererNodes.ts

@@ -29,21 +29,21 @@ import { submitFeedback } from '../../feedback/vue/submitFeedback'
 import { focusAmazonQPanel } from '../../codewhispererChat/commands/registerCommands'
 import { isWeb } from '../../shared/extensionGlobals'
 
-export function createAutoSuggestions(pause: boolean): DataQuickPickItem<'autoSuggestions'> {
+export function createAutoSuggestions(running: boolean): DataQuickPickItem<'autoSuggestions'> {
     const labelResume = localize('AWS.codewhisperer.resumeCodeWhispererNode.label', 'Resume Auto-Suggestions')
     const iconResume = getIcon('vscode-debug-start')
     const labelPause = localize('AWS.codewhisperer.pauseCodeWhispererNode.label', 'Pause Auto-Suggestions')
     const iconPause = getIcon('vscode-debug-pause')
 
     return {
         data: 'autoSuggestions',
-        label: pause ? codicon`${iconPause} ${labelPause}` : codicon`${iconResume} ${labelResume}`,
-        description: pause ? 'Currently RUNNING' : 'Currently PAUSED',
+        label: running ? codicon`${iconPause} ${labelPause}` : codicon`${iconResume} ${labelResume}`,
+        description: running ? 'Currently RUNNING' : 'Currently PAUSED',
         onClick: () => toggleCodeSuggestions.execute(placeholder, cwQuickPickSource),
     } as DataQuickPickItem<'autoSuggestions'>
 }
 
-export function createAutoScans(pause: boolean): DataQuickPickItem<'autoScans'> {
+export function createAutoScans(running: boolean): DataQuickPickItem<'autoScans'> {
     const labelResume = localize('AWS.codewhisperer.resumeCodeWhispererNode.label', 'Resume Auto-Scans')
     const iconResume = getIcon('vscode-debug-alt')
     const labelPause = localize('AWS.codewhisperer.pauseCodeWhispererNode.label', 'Pause Auto-Scans')
@@ -52,8 +52,8 @@ export function createAutoScans(pause: boolean): DataQuickPickItem<'autoScans'>
 
     return {
         data: 'autoScans',
-        label: pause ? codicon`${iconPause} ${labelPause}` : codicon`${iconResume} ${labelResume}`,
-        description: monthlyQuotaExceeded ? 'Monthly quota exceeded' : pause ? 'RUNNING' : 'PAUSED',
+        label: running ? codicon`${iconPause} ${labelPause}` : codicon`${iconResume} ${labelResume}`,
+        description: monthlyQuotaExceeded ? 'Monthly quota exceeded' : running ? 'RUNNING' : 'PAUSED',
         onClick: () => toggleCodeScans.execute(placeholder, cwQuickPickSource),
     } as DataQuickPickItem<'autoScans'>
 }

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

@@ -7,7 +7,6 @@ import * as vscode from 'vscode'
 import * as semver from 'semver'
 import { isCloud9 } from '../../shared/extensionUtilities'
 import { getInlineSuggestEnabled } from '../../shared/utilities/editorUtilities'
-import { getLogger } from '../../shared/logger'
 import { AWSTemplateCaseInsensitiveKeyWords, AWSTemplateKeyWords } from '../models/constants'
 
 export function getLocalDatetime() {
@@ -62,19 +61,6 @@ export function getPrefixSuffixOverlap(firstString: string, secondString: string
     return secondString.slice(0, i)
 }
 
-export function get(key: string, context: vscode.Memento): any {
-    return context.get(key)
-}
-
-export async function set(key: string, value: any, context: vscode.Memento): Promise<void> {
-    await context.update(key, value).then(
-        () => {},
-        (error) => {
-            getLogger().verbose(`Failed to update global state: ${error}`)
-        }
-    )
-}
-
 export function checkLeftContextKeywordsForJsonAndYaml(leftFileContent: string, language: string): boolean {
     if (
         (language === 'json' || language === 'yaml') &&

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

@@ -4,16 +4,9 @@
  */
 
 import globals from '../../shared/extensionGlobals'
-import {
-    customLearnMoreUri,
-    newCustomizationsAvailableKey,
-    newCustomizationMessage,
-    persistedCustomizationsKey,
-    selectedCustomizationKey,
-} from '../models/constants'
+import { customLearnMoreUri, newCustomizationMessage } from '../models/constants'
 import { localize, openUrl } from '../../shared/utilities/vsCodeUtils'
 import { AuthUtil } from './authUtil'
-import { set } from './commonUtil'
 import * as vscode from 'vscode'
 import { createCommonButtons } from '../../shared/ui/buttons'
 import { DataQuickPickItem, showQuickPick } from '../../shared/ui/pickerPrompter'
@@ -107,8 +100,11 @@ export const getSelectedCustomization = (): Customization => {
         return baseCustomization
     }
 
-    const selectedCustomizationArr =
-        globals.context.globalState.get<{ [label: string]: Customization }>(selectedCustomizationKey) || {}
+    const selectedCustomizationArr = globals.globalState.tryGet<{ [label: string]: Customization }>(
+        'CODEWHISPERER_SELECTED_CUSTOMIZATION',
+        Object,
+        {}
+    )
     const result = selectedCustomizationArr[AuthUtil.instance.conn.label] || baseCustomization
 
     // A/B case
@@ -130,12 +126,15 @@ export const setSelectedCustomization = async (customization: Customization) =>
     if (!AuthUtil.instance.isValidEnterpriseSsoInUse() || !AuthUtil.instance.conn) {
         return
     }
-    const selectedCustomizationObj =
-        globals.context.globalState.get<{ [label: string]: Customization }>(selectedCustomizationKey) || {}
+    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 set(selectedCustomizationKey, selectedCustomizationObj, globals.context.globalState)
+    await globals.globalState.tryUpdate('CODEWHISPERER_SELECTED_CUSTOMIZATION', selectedCustomizationObj)
     vsCodeState.isFreeTierLimitReached = false
     await Commands.tryExecute('aws.amazonq.refreshStatusBar')
 }
@@ -144,27 +143,33 @@ export const getPersistedCustomizations = (): Customization[] => {
     if (!AuthUtil.instance.isValidEnterpriseSsoInUse() || !AuthUtil.instance.conn) {
         return []
     }
-    const persistedCustomizationsObj =
-        globals.context.globalState.get<{ [label: string]: Customization[] }>(persistedCustomizationsKey) || {}
+    const persistedCustomizationsObj = globals.globalState.tryGet<{ [label: string]: Customization[] }>(
+        'CODEWHISPERER_PERSISTED_CUSTOMIZATIONS',
+        Object,
+        {}
+    )
     return persistedCustomizationsObj[AuthUtil.instance.conn.label] || []
 }
 
 export const setPersistedCustomizations = async (customizations: Customization[]) => {
     if (!AuthUtil.instance.isValidEnterpriseSsoInUse() || !AuthUtil.instance.conn) {
         return
     }
-    const persistedCustomizationsObj =
-        globals.context.globalState.get<{ [label: string]: Customization[] }>(persistedCustomizationsKey) || {}
+    const persistedCustomizationsObj = globals.globalState.tryGet<{ [label: string]: Customization[] }>(
+        'CODEWHISPERER_PERSISTED_CUSTOMIZATIONS',
+        Object,
+        {}
+    )
     persistedCustomizationsObj[AuthUtil.instance.conn.label] = customizations
-    await set(persistedCustomizationsKey, persistedCustomizationsObj, globals.context.globalState)
+    await globals.globalState.tryUpdate('CODEWHISPERER_PERSISTED_CUSTOMIZATIONS', persistedCustomizationsObj)
 }
 
 export const getNewCustomizationsAvailable = () => {
-    return globals.context.globalState.get<number>(newCustomizationsAvailableKey) ?? 0
+    return globals.globalState.tryGet('aws.amazonq.codewhisperer.newCustomizations', Number, 0)
 }
 
 export const setNewCustomizationsAvailable = async (num: number) => {
-    await set(newCustomizationsAvailableKey, num, globals.context.globalState)
+    await globals.globalState.tryUpdate('aws.amazonq.codewhisperer.newCustomizations', num)
     vsCodeState.isFreeTierLimitReached = false
 }