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

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.
Signed-off-by: Nikolas Komonen <[email protected]>

packages/core/src/auth/auth.ts

@@ -1003,9 +1003,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'
@@ -180,10 +178,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)
     }
@@ -217,7 +219,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)
     }
@@ -286,21 +288,18 @@ 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 = cast(globalState.get(this.key), Optional(String))
+        const id = cast(this.state.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 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'

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