Merge aws#5348 use globalState abstraction

Author: justinmk3

packages/amazonq/test/unit/codewhisperer/util/authUtil.test.ts

@@ -21,7 +21,7 @@ import {
     createTestAuth,
 } from 'aws-core-vscode/test'
 import { Auth, Connection, isAnySsoConnection, isBuilderIdConnection } from 'aws-core-vscode/auth'
-import { vscodeComponent } from 'aws-core-vscode/shared'
+import { globals, vscodeComponent } from 'aws-core-vscode/shared'
 
 const enterpriseSsoStartUrl = 'https://enterprise.awsapps.com/start'
 
@@ -30,7 +30,7 @@ describe('AuthUtil', async function () {
     let authUtil: AuthUtil
 
     beforeEach(async function () {
-        auth = createTestAuth()
+        auth = createTestAuth(globals.globalState)
         authUtil = new AuthUtil(auth)
     })
 
@@ -286,7 +286,7 @@ describe('getChatAuthState()', function () {
     let laterDate: Date
 
     beforeEach(async function () {
-        auth = createTestAuth()
+        auth = createTestAuth(globals.globalState)
         authUtil = new AuthUtil(auth)
 
         laterDate = new Date(Date.now() + 10_000_000)

packages/core/src/auth/secondaryAuth.ts

@@ -98,8 +98,7 @@ export class SecondaryAuth<T extends Connection = Connection> {
         public readonly toolId: ToolId,
         public readonly toolLabel: string,
         public readonly isUsable: (conn: Connection) => conn is T,
-        private readonly auth: Auth,
-        private readonly memento = globals.context.globalState
+        private readonly auth: Auth
     ) {
         const handleConnectionChanged = async (newActiveConn?: Connection) => {
             if (newActiveConn === undefined && this.#activeConnection?.id) {
@@ -176,7 +175,9 @@ export class SecondaryAuth<T extends Connection = Connection> {
     }
 
     public async saveConnection(conn: T) {
-        await this.memento.update(this.key, conn.id)
+        // TODO: fix this
+        // eslint-disable-next-line aws-toolkits/no-banned-usages
+        await globals.context.globalState.update(this.key, conn.id)
         this.#savedConnection = conn
         this.#onDidChangeActiveConnection.fire(this.activeConnection)
     }
@@ -207,7 +208,9 @@ export class SecondaryAuth<T extends Connection = Connection> {
 
     /** Stop using the saved connection and fallback to using the active connection, if it is usable. */
     public async clearSavedConnection() {
-        await this.memento.update(this.key, undefined)
+        // TODO: fix this
+        // eslint-disable-next-line aws-toolkits/no-banned-usages
+        await globals.context.globalState.update(this.key, undefined)
         this.#savedConnection = undefined
         this.#onDidChangeActiveConnection.fire(this.activeConnection)
     }
@@ -252,18 +255,21 @@ export class SecondaryAuth<T extends Connection = Connection> {
     })
 
     private async loadSavedConnection() {
-        const id = cast(this.memento.get(this.key), Optional(String))
+        // TODO: fix this
+        // eslint-disable-next-line aws-toolkits/no-banned-usages
+        const globalState = globals.context.globalState
+        const id = cast(globalState.get(this.key), Optional(String))
         if (id === undefined) {
             return
         }
 
         const conn = await this.auth.getConnection({ id })
         if (conn === undefined) {
             getLogger().warn(`auth (${this.toolId}): removing saved connection "${this.key}" as it no longer exists`)
-            await this.memento.update(this.key, undefined)
+            await globalState.update(this.key, undefined)
         } else if (!this.isUsable(conn)) {
             getLogger().warn(`auth (${this.toolId}): saved connection "${this.key}" is not valid`)
-            await this.memento.update(this.key, undefined)
+            await globalState.update(this.key, undefined)
         } else {
             await this.auth.refreshConnectionState(conn)
             return conn

packages/core/src/dev/activation.ts

@@ -48,6 +48,7 @@ export type DevOptions = {
 }
 
 let targetContext: vscode.ExtensionContext
+let globalState: vscode.Memento
 let targetAuth: Auth
 
 /**
@@ -127,7 +128,7 @@ export class DevDocumentProvider implements vscode.TextDocumentContentProvider {
         } else if (uri.path.startsWith('/globalstate')) {
             // lol hax
             // as of November 2023, all of a memento's properties are stored as property `f` when minified
-            return JSON.stringify((targetContext.globalState as any).f, undefined, 4)
+            return JSON.stringify((globalState as any).f, undefined, 4)
         } else {
             return `unknown URI path: ${uri}`
         }
@@ -154,6 +155,8 @@ export async function activate(ctx: vscode.ExtensionContext): Promise<void> {
         // Internal command to open dev menu for a specific context and options
         vscode.commands.registerCommand('_aws.dev.invokeMenu', (opts: DevOptions) => {
             targetContext = opts.context
+            // eslint-disable-next-line aws-toolkits/no-banned-usages
+            globalState = targetContext.globalState
             targetAuth = opts.auth
             void openMenu(
                 entries(menuOptions)
@@ -298,7 +301,7 @@ class ObjectEditor {
             case 'globalsView':
                 return showState('globalstate')
             case 'globals':
-                return this.openState(targetContext.globalState, key)
+                return this.openState(globalState, key)
             case 'secrets':
                 return this.openState(targetContext.secrets, key)
             case 'auth':
@@ -381,7 +384,7 @@ async function openStorageFromInput() {
                     return new SkipPrompter('')
                 } else if (target === 'globals') {
                     // List all globalState keys in the quickpick menu.
-                    const items = targetContext.globalState
+                    const items = globalState
                         .keys()
                         .map((key) => {
                             return {

packages/core/src/shared/credentials/credentialsProfileMru.ts

@@ -3,23 +3,21 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import * as vscode from 'vscode'
+import globals from '../extensionGlobals'
 
 /**
  * Tracks the credentials selected by the user, ordered by most recent.
  */
 export class CredentialsProfileMru {
     public static readonly maxCredentialMruSize = 5
 
-    private static readonly configurationStateName: string = 'recentCredentials'
-
-    public constructor(private readonly _context: vscode.ExtensionContext) {}
+    public constructor() {}
 
     /**
      * @description Returns the most recently used credentials names
      */
     public getMruList(): string[] {
-        return this._context.globalState.get<string[]>(CredentialsProfileMru.configurationStateName, [])
+        return globals.globalState.tryGet<string[]>('recentCredentials', Object, [])
     }
 
     /**
@@ -38,6 +36,6 @@ export class CredentialsProfileMru {
 
         mru.splice(CredentialsProfileMru.maxCredentialMruSize)
 
-        await this._context.globalState.update(CredentialsProfileMru.configurationStateName, mru)
+        await globals.globalState.update('recentCredentials', mru)
     }
 }

packages/core/src/shared/credentials/defaultCredentialSelectionDataProvider.ts

@@ -46,7 +46,7 @@ export class DefaultCredentialSelectionDataProvider implements CredentialSelecti
         public readonly existingProfileNames: string[],
         protected context: vscode.ExtensionContext
     ) {
-        this._credentialsMru = new CredentialsProfileMru(context)
+        this._credentialsMru = new CredentialsProfileMru()
     }
 
     public async pickCredentialProfile(

packages/core/src/shared/extensionGlobals.ts

@@ -134,6 +134,8 @@ function proxyGlobals(globals_: ToolkitGlobals): ToolkitGlobals {
 let globals = proxyGlobals(resolveGlobalsObject())
 
 export function checkDidReload(context: ExtensionContext): boolean {
+    // TODO: fix this
+    // eslint-disable-next-line aws-toolkits/no-banned-usages
     return !!context.globalState.get<string>('ACTIVATION_LAUNCH_PATH_KEY')
 }
 
@@ -147,6 +149,7 @@ export function initialize(context: ExtensionContext, isWeb: boolean = false): T
         context,
         clock: copyClock(),
         didReload: checkDidReload(context),
+        // eslint-disable-next-line aws-toolkits/no-banned-usages
         globalState: new GlobalState(context.globalState),
         manifestPaths: {} as ToolkitGlobals['manifestPaths'],
         visualizationResourcePaths: {} as ToolkitGlobals['visualizationResourcePaths'],

packages/core/src/shared/globalState.ts

@@ -27,6 +27,7 @@ type globalKey =
     | stepFunctionsKey
     | ToolIdStateKey
     | JsonSchemasKey
+    | 'amazonq.telemetry.migrated'
     | 'aws.amazonq.codewhisperer.newCustomizations'
     | 'aws.amazonq.hasShownWalkthrough'
     | 'aws.amazonq.showTryChatCodeLens'
@@ -61,6 +62,8 @@ type globalKey =
     | 'region'
     // TODO: implement this via `PromptSettings` instead of globalState.
     | 'sam.sync.updateMessage'
+    | 'telemetryClientId'
+    | 'telemetryId'
 
 /**
  * Extension-local (not visible to other vscode extensions) shared state which persists after IDE
@@ -75,7 +78,7 @@ type globalKey =
  * - garbage collection
  */
 export class GlobalState implements vscode.Memento {
-    constructor(private readonly memento: vscode.Memento) {}
+    constructor(protected readonly memento: vscode.Memento) {}
 
     keys(): readonly string[] {
         return this.memento.keys()
@@ -93,9 +96,9 @@ export class GlobalState implements vscode.Memento {
      * {@link String}, {@link Boolean}, etc.
      * @param defaultVal Value returned if `key` has no value.
      */
-    getStrict<T>(key: globalKey, type: TypeConstructor<T>, defaulVal?: T) {
+    getStrict<T>(key: globalKey, type: TypeConstructor<T>, defaultVal?: T) {
         try {
-            const val = this.memento.get<T>(key) ?? defaulVal
+            const val = this.memento.get<T>(key) ?? defaultVal
             return !type || val === undefined ? val : cast(val, type)
         } catch (e) {
             const msg = `GlobalState: invalid state (or read failed) for key: "${key}"`
@@ -117,27 +120,27 @@ export class GlobalState implements vscode.Memento {
      * @param key Key name
      * @param defaultVal Value returned if `key` has no value.
      */
-    get<T>(key: globalKey, defaulVal?: T): T | undefined {
+    get<T>(key: globalKey, defaultVal?: T): T | undefined {
         const skip = (o: any) => o as T // Don't type check.
-        return this.getStrict(key, skip, defaulVal)
+        return this.getStrict(key, skip, defaultVal)
     }
 
     /**
-     * Gets the value for `key` if it satisfies the `type` specification, else logs an error and returns `defaulVal`.
+     * Gets the value for `key` if it satisfies the `type` specification, else logs an error and returns `defaultVal`.
      *
      * @param key Key name
      * @param type Type validator function, or primitive type constructor such as {@link Object},
      * {@link String}, {@link Boolean}, etc.
      * @param defaultVal Value returned if `key` has no value.
      */
     tryGet<T>(key: globalKey, type: TypeConstructor<T>): T | undefined
-    tryGet<T>(key: globalKey, type: TypeConstructor<T>, defaulVal: T): T
-    tryGet<T>(key: globalKey, type: TypeConstructor<T>, defaulVal?: T): T | undefined {
+    tryGet<T>(key: globalKey, type: TypeConstructor<T>, defaultVal: T): T
+    tryGet<T>(key: globalKey, type: TypeConstructor<T>, defaultVal?: T): T | undefined {
         try {
-            return this.getStrict(key, type, defaulVal)
+            return this.getStrict(key, type, defaultVal)
         } catch (e) {
             getLogger().error('%s', (e as Error).message)
-            return defaulVal
+            return defaultVal
         }
     }
 

packages/core/src/shared/telemetry/activation.ts

@@ -44,7 +44,7 @@ export async function activate(
     await config.initAmazonQSetting() // TODO: Remove after a few releases.
 
     DefaultTelemetryClient.productName = productName
-    globals.telemetry = await DefaultTelemetryService.create(extensionContext, awsContext, getComputeRegion())
+    globals.telemetry = await DefaultTelemetryService.create(awsContext, getComputeRegion())
 
     const isAmazonQExt = isAmazonQ()
     try {

packages/core/src/shared/telemetry/telemetryService.ts

@@ -49,12 +49,11 @@ export class DefaultTelemetryService {
      * Use {@link create}() to create an instance of this class.
      */
     protected constructor(
-        private readonly context: ExtensionContext,
         private readonly awsContext: AwsContext,
         private readonly computeRegion?: string,
         publisher?: TelemetryPublisher
     ) {
-        this.persistFilePath = path.join(context.globalStorageUri.fsPath, 'telemetryCache')
+        this.persistFilePath = path.join(globals.context.globalStorageUri.fsPath, 'telemetryCache')
 
         this.startTime = new globals.clock.Date()
 
@@ -70,14 +69,9 @@ export class DefaultTelemetryService {
      * This exists since we need to first ensure the global storage
      * path exists before creating the instance.
      */
-    static async create(
-        context: ExtensionContext,
-        awsContext: AwsContext,
-        computeRegion?: string,
-        publisher?: TelemetryPublisher
-    ) {
-        await DefaultTelemetryService.ensureGlobalStorageExists(context)
-        return new DefaultTelemetryService(context, awsContext, computeRegion, publisher)
+    static async create(awsContext: AwsContext, computeRegion?: string, publisher?: TelemetryPublisher) {
+        await DefaultTelemetryService.ensureGlobalStorageExists(globals.context)
+        return new DefaultTelemetryService(awsContext, computeRegion, publisher)
     }
 
     public get logger(): TelemetryLogger {
@@ -219,8 +213,9 @@ export class DefaultTelemetryService {
             const clientId = getClientId(globals.globalState)
             // grab our Cognito identityId
             const poolId = DefaultTelemetryClient.config.identityPool
-            const identityMapJson = this.context.globalState.get<string>(
+            const identityMapJson = globals.globalState.tryGet(
                 DefaultTelemetryService.telemetryCognitoIdKey,
+                String,
                 '[]'
             )
             // Maps don't cleanly de/serialize with JSON.parse/stringify so we need to do it ourselves
@@ -234,7 +229,7 @@ export class DefaultTelemetryService {
 
                 // save it
                 identityMap.set(poolId, identityPublisherTuple.cognitoIdentityId)
-                await this.context.globalState.update(
+                await globals.globalState.update(
                     DefaultTelemetryService.telemetryCognitoIdKey,
                     JSON.stringify(Array.from(identityMap.entries()))
                 )