fix(amazonq): Use the correct state for Q in remote environment

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

@@ -233,7 +233,11 @@ describe('AuthUtil', async function () {
         // Switch to unsupported connection
         const cwAuthUpdatedConnection = captureEventOnce(authUtil.secondaryAuth.onDidChangeActiveConnection)
         await auth.useConnection(unsupportedConn)
+        // This is triggered when the main Auth connection is switched
         await cwAuthUpdatedConnection
+        // This is triggered by registerAuthListener() when it saves the previous active connection as a fallback.
+        // TODO in a refactor see if we can simplify multiple multiple triggers on the same event.
+        await captureEventOnce(authUtil.secondaryAuth.onDidChangeActiveConnection)
 
         // Is using the fallback connection
         assert.ok(authUtil.isConnected())

packages/core/src/auth/auth.ts

@@ -1022,9 +1022,20 @@ export class Auth implements AuthService, ConnectionManager {
         }
     }
 
+    /**
+     * Auth uses its own state, separate from the default {@link globalState}
+     * that is normally used throughout the codebase.
+     *
+     * IMPORTANT: Anything involving auth should ONLY use this state since the state
+     * can vary on certain conditions. So if you see something explicitly using
+     * globalState verify if it should actually be using that.
+     */
+    public getStateMemento: () => vscode.Memento = () => Auth._getStateMemento()
+    private static _getStateMemento = once(() => getEnvironmentSpecificMemento())
+
     static #instance: Auth | undefined
     public static get instance() {
-        return (this.#instance ??= new Auth(new ProfileStore(getEnvironmentSpecificMemento())))
+        return (this.#instance ??= new Auth(new ProfileStore(Auth._getStateMemento())))
     }
 
     private getSsoProfileLabel(profile: SsoProfile) {

packages/core/src/auth/secondaryAuth.ts

@@ -3,8 +3,6 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import globals from '../shared/extensionGlobals'
-
 import * as vscode from 'vscode'
 import { getLogger } from '../shared/logger'
 import { cast, Optional } from '../shared/utilities/typeConstructors'
@@ -18,6 +16,7 @@ import { AuthStatus, telemetry } from '../shared/telemetry/telemetry'
 import { asStringifiedStack } from '../shared/telemetry/spans'
 import { withTelemetryContext } from '../shared/telemetry/util'
 import { isNetworkError } from '../shared/errors'
+import globals from '../shared/extensionGlobals'
 
 export type ToolId = 'codecatalyst' | 'codewhisperer' | 'testId'
 
@@ -214,10 +213,14 @@ export class SecondaryAuth<T extends Connection = Connection> {
         return !!this.activeConnection && this.auth.getConnectionState(this.activeConnection) === 'invalid'
     }
 
+    public get state() {
+        return this.auth.getStateMemento()
+    }
+
     public async saveConnection(conn: T) {
         // TODO: fix this
         // eslint-disable-next-line aws-toolkits/no-banned-usages
-        await globals.context.globalState.update(this.key, conn.id)
+        await this.state.update(this.key, conn.id)
         this.#savedConnection = conn
         this.#onDidChangeActiveConnection.fire(this.activeConnection)
     }
@@ -251,7 +254,7 @@ export class SecondaryAuth<T extends Connection = Connection> {
     public async clearSavedConnection() {
         // TODO: fix this
         // eslint-disable-next-line aws-toolkits/no-banned-usages
-        await globals.context.globalState.update(this.key, undefined)
+        await this.state.update(this.key, undefined)
         this.#savedConnection = undefined
         this.#onDidChangeActiveConnection.fire(this.activeConnection)
     }
@@ -326,9 +329,6 @@ export class SecondaryAuth<T extends Connection = Connection> {
      * Provides telemetry if called by restoreConnection() (or another auth_modifyConnection context)
      */
     private async _loadSavedConnection() {
-        // TODO: fix this
-        // eslint-disable-next-line aws-toolkits/no-banned-usages
-        const globalState = globals.context.globalState
         const id = this.getStateConnectionId()
         if (id === undefined) {
             return
@@ -337,10 +337,10 @@ export class SecondaryAuth<T extends Connection = Connection> {
         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 globalState.update(this.key, undefined)
+            await this.state.update(this.key, undefined)
         } else if (!this.isUsable(conn)) {
             getLogger().warn(`auth (${this.toolId}): saved connection "${this.key}" is not valid`)
-            await globalState.update(this.key, undefined)
+            await this.state.update(this.key, undefined)
         } else {
             const getAuthStatus = (state: ReturnType<typeof this.auth.getConnectionState>): AuthStatus => {
                 return state === 'invalid' ? 'expired' : 'connected'
@@ -377,7 +377,7 @@ export class SecondaryAuth<T extends Connection = Connection> {
     }
 
     private getStateConnectionId() {
-        return cast(globals.globalState.get(this.key), Optional(String))
+        return cast(this.state.get(this.key), Optional(String))
     }
 }
 

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

@@ -29,7 +29,14 @@ export function partition(memento: vscode.Memento, key: string): vscode.Memento
     }
 }
 
-/** Avoids sharing globalState with remote vscode instances. */
+/**
+ * Resolves the appropriate memento/state for the current runtime environment.
+ *
+ * Why?
+ * - In remote instances where we ssh in to them, we do not always want to share
+ *   with the local globalState. We want certain functionality to be isolated to
+ *   the remote instance.
+ */
 export function getEnvironmentSpecificMemento(): vscode.Memento {
     if (!vscode.env.remoteName) {
         // local compute: no further partitioning

packages/core/src/test/credentials/secondaryAuth.test.ts

@@ -9,6 +9,7 @@ import { createBuilderIdProfile, createTestAuth } from './testUtil'
 import { Connection, createSsoProfile, hasScopes, isSsoConnection } from '../../auth/connection'
 import assert from 'assert'
 import globals from '../../shared/extensionGlobals'
+import { waitUntil } from '../../shared/utilities/timeoutUtils'
 
 describe('SecondaryAuth', function () {
     let auth: ReturnType<typeof createTestAuth>
@@ -100,8 +101,13 @@ describe('SecondaryAuth', function () {
         // delete SecondaryAuth
         await auth.deleteConnection(conn)
 
+        // currently both Auth onDidChangeConnectionState and onDidDeleteConnection trigger
+        // and we need to wait for both of the callbacks defined through them in SecondaryAuth to complete.
+        // We know they all completed when SecondaryAuth.onDidChangeActiveConnection has been called twice
+        await waitUntil(async () => onDidChangeActiveConnection.callCount === 2, { interval: 10, timeout: 10000 })
+
         // we fallback to the PrimaryAuth connection
-        assert.strictEqual(onDidChangeActiveConnection.called, true)
+        assert.strictEqual(onDidChangeActiveConnection.callCount, 2)
         assert.deepStrictEqual(
             {
                 id: secondaryAuth.activeConnection?.id,