refactor(codecatalyst): use globalState abstraction

Author: justinmk3

packages/core/src/auth/secondaryAuth.ts

@@ -99,7 +99,7 @@ export class SecondaryAuth<T extends Connection = Connection> {
         public readonly toolLabel: string,
         public readonly isUsable: (conn: Connection) => conn is T,
         private readonly auth: Auth,
-        private readonly memento = globals.globalState
+        private readonly memento = globals.context.globalState
     ) {
         const handleConnectionChanged = async (newActiveConn?: Connection) => {
             if (newActiveConn === undefined && this.#activeConnection?.id) {

packages/core/src/codecatalyst/auth.ts

@@ -9,6 +9,7 @@ import { Auth } from '../auth/auth'
 import * as localizedText from '../shared/localizedText'
 import { getSecondaryAuth, setScopes } from '../auth/secondaryAuth'
 import { getLogger } from '../shared/logger'
+import globals from '../shared/extensionGlobals'
 import { ToolkitError, isAwsError } from '../shared/errors'
 import { MetricName, MetricShapes, telemetry } from '../shared/telemetry/telemetry'
 import { openUrl } from '../shared/utilities/vsCodeUtils'
@@ -70,12 +71,8 @@ export class CodeCatalystAuthenticationProvider {
     private readonly onDidChangeEmitter = new vscode.EventEmitter<void>()
     public readonly onDidChange = this.onDidChangeEmitter.event
 
-    private readonly mementoKey = 'codecatalyst.connections'
-
     public constructor(
         protected readonly storage: CodeCatalystAuthStorage,
-        protected readonly memento: vscode.Memento,
-
         public readonly auth = Auth.instance,
         public readonly secondaryAuth = getSecondaryAuth(
             auth,
@@ -364,7 +361,7 @@ export class CodeCatalystAuthenticationProvider {
     }
 
     private getStates(): Record<string, ConnectionState> {
-        return this.memento.get(this.mementoKey, {} as Record<string, ConnectionState>)
+        return globals.globalState.tryGet<Record<string, ConnectionState>>('codecatalyst.connections', Object, {})
     }
 
     public tryGetConnectionState(conn: SsoConnection): ConnectionState | undefined {
@@ -380,7 +377,7 @@ export class CodeCatalystAuthenticationProvider {
     }
 
     private async setConnectionState(conn: SsoConnection, state: ConnectionState) {
-        await this.memento.update(this.mementoKey, {
+        await globals.globalState.update('codecatalyst.connections', {
             ...this.getStates(),
             [conn.id]: state,
         })
@@ -424,7 +421,7 @@ export class CodeCatalystAuthenticationProvider {
         return CodeCatalystAuthenticationProvider.#instance
     }
 
-    public static fromContext(ctx: Pick<vscode.ExtensionContext, 'secrets' | 'globalState'>) {
-        return (this.#instance ??= new this(new CodeCatalystAuthStorage(ctx.secrets), ctx.globalState))
+    public static fromContext(ctx: Pick<vscode.ExtensionContext, 'secrets'>) {
+        return (this.#instance ??= new this(new CodeCatalystAuthStorage(ctx.secrets)))
     }
 }

packages/core/src/codecatalyst/commands.ts

@@ -335,7 +335,7 @@ export class CodeCatalystCommands {
         return devenv
     }
 
-    public static fromContext(ctx: Pick<vscode.ExtensionContext, 'secrets' | 'globalState'>) {
+    public static fromContext(ctx: Pick<vscode.ExtensionContext, 'secrets'>) {
         const auth = CodeCatalystAuthenticationProvider.fromContext(ctx)
 
         return new this(auth)

packages/core/src/codecatalyst/reconnect.ts

@@ -55,7 +55,7 @@ function handleRestart(conn: SsoConnection, ctx: ExtContext, envId: string | und
         getLogger().info('codecatalyst: attempting to poll dev environments')
 
         // Reconnect devenvs (if coming from a restart)
-        reconnectDevEnvs(conn, ctx).catch((err) => {
+        reconnectDevEnvs(conn).catch((err) => {
             getLogger().error(`codecatalyst: error while resuming devenvs: ${err}`)
         })
     }
@@ -66,9 +66,12 @@ function handleRestart(conn: SsoConnection, ctx: ExtContext, envId: string | und
  * @param conn a connection that may be used for CodeCatalyst
  * @param ctx the extension context
  */
-async function reconnectDevEnvs(conn: SsoConnection, ctx: ExtContext): Promise<void> {
-    const memento = ctx.extensionContext.globalState
-    const pendingDevEnvs = memento.get<Record<string, DevEnvMemento>>('CODECATALYST_RECONNECT', {})
+async function reconnectDevEnvs(conn: SsoConnection): Promise<void> {
+    const pendingDevEnvs = globals.globalState.tryGet<Record<string, DevEnvMemento>>(
+        'CODECATALYST_RECONNECT',
+        Object,
+        {}
+    )
     const validDevEnvs = filterInvalidDevEnvs(pendingDevEnvs)
     if (Object.keys(validDevEnvs).length === 0) {
         return
@@ -93,7 +96,7 @@ async function reconnectDevEnvs(conn: SsoConnection, ctx: ExtContext): Promise<v
             progress.report({ message: progressTitle })
             const client = await createClient(conn)
 
-            return pollDevEnvs(client, progress, token, memento, validDevEnvs)
+            return pollDevEnvs(client, progress, token, validDevEnvs)
         }
     )
 }
@@ -115,14 +118,13 @@ function filterInvalidDevEnvs(devenvs: Record<string, DevEnvMemento>) {
 /**
  * Ensure that all devenvs that are currently being looked at set to attempting to reconnect so that they are not looked at
  * by any other instance of VSCode.
- * @param memento
  * @param devenvs
  */
-function setWatchedDevEnvStatus(memento: vscode.Memento, devenvs: Record<string, DevEnvMemento>, watchStatus: boolean) {
+function setWatchedDevEnvStatus(devenvs: Record<string, DevEnvMemento>, watchStatus: boolean) {
     for (const [id, detail] of Object.entries(devenvs)) {
         devenvs[id] = { ...detail, attemptingReconnect: watchStatus }
     }
-    return memento.update('CODECATALYST_RECONNECT', devenvs)
+    return globals.globalState.update('CODECATALYST_RECONNECT', devenvs)
 }
 
 /**
@@ -131,26 +133,24 @@ function setWatchedDevEnvStatus(memento: vscode.Memento, devenvs: Record<string,
  *      2. In a terminating state or expired, in which case no longer watch the devenv
  *      3. Failed to start, in which case notify the user
  * @param client A connected client
- * @param memento vscode global storage library
  * @param devenvs All VALID devenvs that are not being watched by any other VSCode instance
  */
 async function pollDevEnvs(
     client: CodeCatalystClient,
     progress: vscode.Progress<{ message: string }>,
     token: vscode.CancellationToken,
-    memento: vscode.Memento,
     devenvs: Record<string, DevEnvMemento>
 ) {
     // Ensure that all devenvs that you want to look at are attempting reconnection
     // and won't be watched by any other VSCode instance
-    await setWatchedDevEnvStatus(memento, devenvs, true)
+    await setWatchedDevEnvStatus(devenvs, true)
 
     const shouldCloseRootInstance = Object.keys(devenvs).length === 1
     getLogger().info(`codecatalyst: reconnect: pollDevEnvs: ${Object.keys(devenvs).length}`)
 
     while (Object.keys(devenvs).length > 0) {
         if (token.isCancellationRequested) {
-            await setWatchedDevEnvStatus(memento, devenvs, false)
+            await setWatchedDevEnvStatus(devenvs, false)
             return
         }
 
@@ -197,7 +197,7 @@ async function pollDevEnvs(
                     progress.report({ message: `Dev Environment ${devenvName} has expired` })
                 }
             } catch {
-                await failDevEnv(memento, id)
+                await failDevEnv(id)
                 delete devenvs[id]
                 void showViewLogsMessage(
                     localize('AWS.codecatalyst.reconnect', 'Unable to reconnect to ${0}', devenvName)
@@ -220,14 +220,13 @@ function isExpired(previousConnectionTime: number): boolean {
 }
 
 /**
- * When a devenv fails, remove it from the memento so we no longer watch it in the future
- * @param memento The memento instance from vscode
+ * When a devenv fails, remove it from globalState so we no longer watch it in the future
  * @param devenvId the id of the deveng to fail
  */
-function failDevEnv(memento: vscode.Memento, devenvId: string) {
-    const curr = memento.get<Record<string, DevEnvMemento>>('CODECATALYST_RECONNECT', {})
+function failDevEnv(devenvId: string) {
+    const curr = globals.globalState.tryGet<Record<string, DevEnvMemento>>('CODECATALYST_RECONNECT', Object, {})
     delete curr[devenvId]
-    return memento.update('CODECATALYST_RECONNECT', curr)
+    return globals.globalState.update('CODECATALYST_RECONNECT', curr)
 }
 
 async function openReconnectedDevEnv(
@@ -268,8 +267,11 @@ export function isLongReconnect(oldSettings: DevEnvironmentSettings, newSettings
 }
 
 export function saveReconnectionInformation(devenv: DevEnvironmentId & Pick<DevEnvironment, 'alias'>): Thenable<void> {
-    const memento = globals.context.globalState
-    const pendingReconnects = memento.get<Record<string, DevEnvMemento>>('CODECATALYST_RECONNECT', {})
+    const pendingReconnects = globals.globalState.tryGet<Record<string, DevEnvMemento>>(
+        'CODECATALYST_RECONNECT',
+        Object,
+        {}
+    )
     const workspaceFolders = vscode.workspace.workspaceFolders
     const currentWorkspace =
         workspaceFolders !== undefined && workspaceFolders.length > 0 ? workspaceFolders[0].uri.fsPath : '/projects'
@@ -281,12 +283,15 @@ export function saveReconnectionInformation(devenv: DevEnvironmentId & Pick<DevE
         previousConnectionTimestamp: Date.now(),
         alias: devenv.alias,
     }
-    return memento.update('CODECATALYST_RECONNECT', pendingReconnects)
+    return globals.globalState.update('CODECATALYST_RECONNECT', pendingReconnects)
 }
 
 export function removeReconnectionInformation(devenv: DevEnvironmentId): Thenable<void> {
-    const memento = globals.context.globalState
-    const pendingReconnects = memento.get<Record<string, DevEnvMemento>>('CODECATALYST_RECONNECT', {})
+    const pendingReconnects = globals.globalState.tryGet<Record<string, DevEnvMemento>>(
+        'CODECATALYST_RECONNECT',
+        Object,
+        {}
+    )
     delete pendingReconnects[devenv.id]
-    return memento.update('CODECATALYST_RECONNECT', pendingReconnects)
+    return globals.globalState.update('CODECATALYST_RECONNECT', pendingReconnects)
 }

packages/core/src/dev/beta.ts

@@ -22,20 +22,19 @@ import { CancellationError } from '../shared/utilities/timeoutUtils'
 const localize = nls.loadMessageBundle()
 
 const downloadIntervalMs = 1000 * 60 * 60 * 24 // A day in milliseconds
-const betaToolkitKey = 'dev.beta'
 
 interface BetaToolkit {
     readonly needUpdate: boolean
     readonly lastCheck: number
 }
 
 function getBetaToolkitData(vsixUrl: string): BetaToolkit | undefined {
-    return globals.context.globalState.get<Record<string, BetaToolkit>>(betaToolkitKey, {})[vsixUrl]
+    return globals.globalState.tryGet<Record<string, BetaToolkit>>('dev.beta', Object, {})[vsixUrl]
 }
 
 async function updateBetaToolkitData(vsixUrl: string, data: BetaToolkit) {
-    await globals.context.globalState.update(betaToolkitKey, {
-        ...globals.context.globalState.get<Record<string, BetaToolkit>>(betaToolkitKey, {}),
+    await globals.globalState.update('dev.beta', {
+        ...globals.globalState.get<Record<string, BetaToolkit>>('dev.beta', {}),
         [vsixUrl]: data,
     })
 }

packages/core/src/shared/globalState.ts

@@ -45,11 +45,19 @@ type globalKey =
     | 'CODEWHISPERER_SELECTED_CUSTOMIZATION'
     | 'CODEWHISPERER_USER_GROUP'
     | 'awsTelemetryNoticeVersionAck'
+    | 'codecatalyst.connections'
+    | 'dev.beta'
     | 'globalsMostRecentVersion'
     | 'gumby.wasQCodeTransformationUsed'
     | 'hasAlreadyOpenedAmazonQ'
     | 'isExtensionFirstUse'
     | 'lastExtensionVersion'
+    | 'lastSelectedRegion'
+    | 'recentCredentials'
+    // List of regions enabled in AWS Explorer.
+    | 'region'
+    // TODO: implement this via `PromptSettings` instead of globalState.
+    | 'sam.sync.updateMessage'
 
 /**
  * Extension-local (not visible to other vscode extensions) shared state which persists after IDE

packages/core/src/shared/ui/common/region.ts

@@ -27,13 +27,12 @@ export function createRegionPrompter(
     regions = globals.regionProvider.getRegions(),
     options: RegionPrompterOptions = {}
 ): QuickPickPrompter<Region> {
-    const lastRegionKey = 'lastSelectedRegion'
     const defaultRegion = options.defaultRegion ?? globals.regionProvider.defaultRegionId
     const filteredRegions = regions.filter(
         (r) => !options.serviceFilter || globals.regionProvider.isServiceInRegion(options.serviceFilter, r.id)
     )
 
-    const lastRegion = globals.context.globalState.get<Region>(lastRegionKey)
+    const lastRegion = globals.globalState.tryGet<Region>('lastSelectedRegion', Object)
     const items = filteredRegions.map((region) => ({
         label: region.name,
         detail: region.id,
@@ -60,9 +59,7 @@ export function createRegionPrompter(
 
     return prompter.transform((item) => {
         getLogger().debug('createRegionPrompter: selected %O', item)
-        globals.context.globalState.update(lastRegionKey, item).then(undefined, (e) => {
-            getLogger().error('globalState.update() failed: %s', (e as Error).message)
-        })
+        globals.globalState.tryUpdate('lastSelectedRegion', item)
         return item
     })
 }

packages/core/src/shared/utilities/messages.ts

@@ -364,17 +364,14 @@ export async function copyToClipboard(data: string, label?: string): Promise<voi
     getLogger().verbose('copied %s to clipboard: %O', label ?? '', data)
 }
 
-export async function showOnce<T>(
-    key: string,
-    fn: () => Promise<T>,
-    memento = globals.context.globalState
-): Promise<T | undefined> {
-    if (memento.get(key)) {
+/** TODO: eliminate this, callers should use `PromptSettings` instead. */
+export async function showOnce<T>(key: 'sam.sync.updateMessage', fn: () => Promise<T>): Promise<T | undefined> {
+    if (globals.globalState.tryGet(key, Boolean, false)) {
         return
     }
 
     const result = fn()
-    await memento.update(key, true)
+    await globals.globalState.update(key, true)
 
     return result
 }

packages/core/src/test/codecatalyst/auth.test.ts

@@ -6,7 +6,7 @@
 import assert from 'assert'
 import { CodeCatalystAuthStorage, CodeCatalystAuthenticationProvider, defaultScopes } from '../../codecatalyst/auth'
 import { getTestWindow } from '../shared/vscode/window'
-import { FakeMemento, FakeSecretStorage } from '../fakeExtensionContext'
+import { FakeSecretStorage } from '../fakeExtensionContext'
 import { createBuilderIdProfile, createSsoProfile, createTestAuth } from '../credentials/testUtil'
 import Sinon from 'sinon'
 import { isAnySsoConnection } from '../../auth/connection'
@@ -21,7 +21,6 @@ describe('CodeCatalystAuthenticationProvider', async function () {
         auth = createTestAuth()
         codecatalystAuth = new CodeCatalystAuthenticationProvider(
             new CodeCatalystAuthStorage(new FakeSecretStorage()),
-            new FakeMemento(),
             auth
         )
     })

packages/core/src/test/login/webview/vue/backend_toolkit.test.ts

@@ -16,7 +16,7 @@ import {
     CodeCatalystAuthStorage,
     defaultScopes,
 } from '../../../../codecatalyst/auth'
-import { FakeSecretStorage, FakeMemento } from '../../../fakeExtensionContext'
+import { FakeSecretStorage } from '../../../fakeExtensionContext'
 import * as authUtils from '../../../../auth/utils'
 
 // TODO: remove auth page and tests
@@ -37,7 +37,6 @@ describe('Toolkit Login', function () {
         auth = createTestAuth()
         codecatalystAuth = new CodeCatalystAuthenticationProvider(
             new CodeCatalystAuthStorage(new FakeSecretStorage()),
-            new FakeMemento(),
             auth
         )
         sandbox.stub(Auth, 'instance').value(auth)