refactor(codewhisperer): use globalState abstraction

Author: justinmk3

packages/core/src/codewhisperer/commands/onInlineAcceptance.ts

@@ -50,31 +50,25 @@ export const acceptSuggestion = Commands.declare(
             const editor = vscode.window.activeTextEditor
             await Container.instance.lineAnnotationController.refresh(editor, 'codewhisperer')
             const onAcceptanceFunc = isInlineCompletionEnabled() ? onInlineAcceptance : onAcceptance
-            await onAcceptanceFunc(
-                {
-                    editor,
-                    range,
-                    effectiveRange,
-                    acceptIndex,
-                    recommendation,
-                    requestId,
-                    sessionId,
-                    triggerType,
-                    completionType,
-                    language,
-                    references,
-                },
-                context.extensionContext.globalState
-            )
+            await onAcceptanceFunc({
+                editor,
+                range,
+                effectiveRange,
+                acceptIndex,
+                recommendation,
+                requestId,
+                sessionId,
+                triggerType,
+                completionType,
+                language,
+                references,
+            })
         }
 )
 /**
  * This function is called when user accepts a intelliSense suggestion or an inline suggestion
  */
-export async function onInlineAcceptance(
-    acceptanceEntry: OnRecommendationAcceptanceEntry,
-    globalStorage: vscode.Memento
-) {
+export async function onInlineAcceptance(acceptanceEntry: OnRecommendationAcceptanceEntry) {
     RecommendationHandler.instance.cancelPaginatedRequest()
     RecommendationHandler.instance.disposeInlineCompletion()
 
@@ -121,7 +115,7 @@ export async function onInlineAcceptance(
             language: languageContext.language,
         })
         const insertedCoderange = new vscode.Range(start, end)
-        CodeWhispererCodeCoverageTracker.getTracker(languageContext.language, globalStorage)?.countAcceptedTokens(
+        CodeWhispererCodeCoverageTracker.getTracker(languageContext.language)?.countAcceptedTokens(
             insertedCoderange,
             acceptanceEntry.editor.document.getText(insertedCoderange),
             acceptanceEntry.editor.document.fileName

packages/core/src/codewhisperer/tracker/codewhispererCodeCoverageTracker.ts

@@ -383,10 +383,7 @@ export class CodeWhispererCodeCoverageTracker {
 
     public static readonly instances = new Map<CodewhispererLanguage, CodeWhispererCodeCoverageTracker>()
 
-    public static getTracker(
-        language: string,
-        memeto: vscode.Memento = globals.context.globalState
-    ): CodeWhispererCodeCoverageTracker | undefined {
+    public static getTracker(language: string): CodeWhispererCodeCoverageTracker | undefined {
         if (!runtimeLanguageContext.isLanguageSupported(language)) {
             return undefined
         }

packages/core/src/codewhispererChat/editor/codelens.ts

@@ -33,7 +33,6 @@ export class TryChatCodeLensProvider implements vscode.CodeLensProvider {
     /** How many times we want to show the CodeLens */
     static readonly maxCount = 10
     static readonly debounceMillis: number = 700
-    static readonly showCodeLensId = `aws.amazonq.showTryChatCodeLens`
 
     private static providerDisposable: vscode.Disposable | undefined = undefined
     private disposables: vscode.Disposable[] = []
@@ -58,7 +57,7 @@ export class TryChatCodeLensProvider implements vscode.CodeLensProvider {
     }
 
     static async register(isAmazonQVisible: vscode.Event<boolean>): Promise<boolean> {
-        const shouldShow = globals.context.globalState.get(this.showCodeLensId, true)
+        const shouldShow = globals.globalState.tryGet('aws.amazonq.showTryChatCodeLens', Boolean, true)
         if (!shouldShow) {
             return false
         }
@@ -85,7 +84,7 @@ export class TryChatCodeLensProvider implements vscode.CodeLensProvider {
              * - lineAnnotationController is visible (i.e not in end state)
              * - Amazon Q chat is visible
              */
-            const isLineAnnotationVisible = globals.context.globalState.get<string>(inlinehintKey)
+            const isLineAnnotationVisible = globals.globalState.tryGet(inlinehintKey, String)
             if ((isLineAnnotationVisible && isLineAnnotationVisible !== EndState.id) || this.isAmazonQVisible) {
                 return resolve([])
             }
@@ -150,7 +149,7 @@ export class TryChatCodeLensProvider implements vscode.CodeLensProvider {
     }
 
     dispose() {
-        void globals.context.globalState.update(TryChatCodeLensProvider.showCodeLensId, false)
+        globals.globalState.tryUpdate('aws.amazonq.showTryChatCodeLens', false)
         TryChatCodeLensProvider.providerDisposable?.dispose()
         this.disposables.forEach((d) => d.dispose())
     }

packages/core/src/shared/globalState.ts

@@ -26,6 +26,7 @@ type globalKey =
     | ToolIdStateKey
     | 'aws.amazonq.codewhisperer.newCustomizations'
     | 'aws.amazonq.hasShownWalkthrough'
+    | 'aws.amazonq.showTryChatCodeLens'
     | 'aws.downloadPath'
     | 'aws.lastTouchedS3Folder'
     | 'aws.lastUploadedToS3Folder'

packages/core/src/test/codewhispererChat/editor/codelens.test.ts

@@ -93,12 +93,12 @@ describe('TryChatCodeLensProvider', () => {
     it('does not register the provider if we do not want to show the code lens', async function () {
         await TryChatCodeLensProvider.register(isAmazonQVisibleEvent)
         // indicate we do not want to show it
-        await globals.context.globalState.update(TryChatCodeLensProvider.showCodeLensId, false)
+        await globals.globalState.update('aws.amazonq.showTryChatCodeLens', false)
         // ensure we do not show it
         assert.deepStrictEqual(await TryChatCodeLensProvider.register(isAmazonQVisibleEvent), false)
 
         // indicate we want to show it
-        await globals.context.globalState.update(TryChatCodeLensProvider.showCodeLensId, true)
+        await globals.globalState.update('aws.amazonq.showTryChatCodeLens', true)
         // The general toolkit activation will have already registered this provider, so it throws when we try again
         // But if it throws it implies it tried to register it.
         await assert.rejects(TryChatCodeLensProvider.register(isAmazonQVisibleEvent), {
@@ -125,7 +125,7 @@ describe('TryChatCodeLensProvider', () => {
         stubConnection('connected')
         isAmazonQVisibleEventEmitter.fire(false)
         // indicate lineAnnotationController is not visible and in end state
-        await globals.context.globalState.update(inlinehintKey, EndState.id)
+        await globals.globalState.update(inlinehintKey, EndState.id)
 
         let codeLensCount = 0
         const modifierKey = resolveModifierKey()
@@ -165,7 +165,7 @@ describe('TryChatCodeLensProvider', () => {
 
         lineAnnotationControllerStates.forEach((id: string) => {
             it(`id - ${id}`, async () => {
-                await globals.context.globalState.update(inlinehintKey, id)
+                await globals.globalState.update(inlinehintKey, id)
 
                 const emptyResult = await instance.provideCodeLenses(
                     {} as any,