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

## Problem:

As part of our design we intentionally separate the Auth state between a
local environment, versus a remote (ssh'd) one. We do this so that if
the user has 2 IDE instances, one local and the other remote, if one
changes its auth state the other one will not be affected.

The reason for separating auth that I can remember is due to how we
store the actual tokens on disk. Because the remote does not have access
to the disk of the local, we need to ensure the remote operates its auth
independently of the local auth.

By design globalState is shared by both local and remote instances. So
for local we need to ensure we do not use the base globalState,
otherwise the same Auth state as the remote will be used.

The problem, SecondaryAuth is not doing this for all cases. It is using
the globalState when it should be asking Auth which state it should use.
The happy path was working though, but there can be potential issues
with this bug at some other point.

## Solution:

Expose the state that Auth is using so that something like SecondaryAuth
can ask Auth for the correct state object depending on the local vs
remote.


---

<!--- REMINDER: Ensure that your PR meets the guidelines in
CONTRIBUTING.md -->

License: I confirm that my contribution is made under the terms of the
Apache 2.0 license.

---------

Signed-off-by: Nikolas Komonen <[email protected]>
Signed-off-by: nkomonen-amazon <[email protected]>

Author: nkomonen-amazon

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,