fix(auth): login state updates across windows (#5547)

Problem: Users can open multiple VSC sessions. If they sign out or the
connection is lost in another window, then interacting with chat or
in-line suggestions in the original window will only produce cryptic
"Profile does not exist" errors, with no clear way to get back to a
working state.

Solution: Log out once this state is detected (uses chat, invokes manual
or auto inline suggestion).

3 commits:

[fix(amazonq): user stays signed in if they log out in other
windows](270a321)
Problem: Users can open multiple VSC sessions. If they sign out or the
connection is lost in another window, then interacting with chat or
in-line suggestions in the original window will only produce cryptic
"Profile does not exist" errors, with no clear way to get back to a
working state.

Solution: Log out once this state is detected (uses chat, invokes manual
or auto inline suggestion).

[fix(amazonq): disable 'connection already exists' error for login
page](5593784)
Error would appear due to certain race conditions. If the auth page is
visible then we know that we have no connections. (for amazon q only).

[fix(auth): login state updates across
windows](ced109f)
- Uses a vscode filwatcher to detect changes to `.aws/sso/cache`.
- Try to 'reconnect' on new changes based on stored state key.
  - Global state is stored across windows.
- Login or log out will happen in all vsc instances.

---

<!--- 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.

Author: hayemaxi

packages/amazonq/.changes/next-release/Bug Fix-bab65b05-77c0-4a6d-ae51-4af37160d957.json

@@ -0,0 +1,4 @@
+{
+	"type": "Bug Fix",
+	"description": "Auth: Login state not updating across multiple VS Code windows."
+}

packages/core/src/auth/auth.ts

@@ -14,7 +14,7 @@ import { Credentials } from '@aws-sdk/types'
 import { SsoAccessTokenProvider } from './sso/ssoAccessTokenProvider'
 import { Timeout } from '../shared/utilities/timeoutUtils'
 import { errorCode, isAwsError, isNetworkError, ToolkitError, UnknownError } from '../shared/errors'
-import { getCache } from './sso/cache'
+import { getCache, getCacheFileWatcher } from './sso/cache'
 import { isNonNullable, Mutable } from '../shared/utilities/tsUtils'
 import { builderIdStartUrl, SsoToken, truncateStartUrl } from './sso/model'
 import { SsoClient } from './sso/clients'
@@ -58,6 +58,7 @@ import {
     scopesSsoAccountAccess,
     AwsConnection,
     scopesCodeWhispererCore,
+    ProfileNotFoundError,
 } from './connection'
 import { isSageMaker, isCloud9, isAmazonQ } from '../shared/extensionUtilities'
 import { telemetry } from '../shared/telemetry/telemetry'
@@ -132,6 +133,7 @@ const authClassName = 'Auth'
 
 export class Auth implements AuthService, ConnectionManager {
     readonly #ssoCache = getCache()
+    readonly #ssoCacheWatcher = getCacheFileWatcher()
     readonly #validationErrors = new Map<Connection['id'], Error>()
     readonly #invalidCredentialsTimeouts = new Map<Connection['id'], Timeout>()
     readonly #onDidChangeActiveConnection = new vscode.EventEmitter<StatefulConnection | undefined>()
@@ -160,6 +162,10 @@ export class Auth implements AuthService, ConnectionManager {
         return this.store.listProfiles().length !== 0
     }
 
+    public get cacheWatcher() {
+        return this.#ssoCacheWatcher
+    }
+
     /**
      * Map startUrl -> declared connections
      */
@@ -878,8 +884,22 @@ export class Auth implements AuthService, ConnectionManager {
     @withTelemetryContext({ name: 'handleInvalidCredentials', class: authClassName })
     private async handleInvalidCredentials<T>(id: Connection['id'], refresh: () => Promise<T>): Promise<T> {
         getLogger().info(`auth: Handling invalid credentials of connection: ${id}`)
-        const profile = this.store.getProfile(id)
-        const previousState = profile?.metadata.connectionState
+
+        let profile: StoredProfile
+        try {
+            profile = this.store.getProfileOrThrow(id)
+        } catch (err) {
+            if (err instanceof ProfileNotFoundError) {
+                getLogger().info(
+                    `Auth: deleting connection '${id}' due to error encountered while fetching auth token: %s`,
+                    err
+                )
+                await this.deleteConnection({ id })
+            }
+            throw err
+        }
+
+        const previousState = profile.metadata.connectionState
         await this.updateConnectionState(id, 'invalid')
 
         if (previousState === 'invalid') {
@@ -895,8 +915,7 @@ export class Auth implements AuthService, ConnectionManager {
             const timeout = new Timeout(60000)
             this.#invalidCredentialsTimeouts.set(id, timeout)
 
-            const connLabel =
-                profile?.metadata.label ?? (profile?.type === 'sso' ? this.getSsoProfileLabel(profile) : id)
+            const connLabel = profile.metadata.label ?? (profile.type === 'sso' ? this.getSsoProfileLabel(profile) : id)
             const message = localize(
                 'aws.auth.invalidConnection',
                 'Connection "{0}" is invalid or expired, login again?',

packages/core/src/auth/connection.ts

@@ -201,6 +201,12 @@ export interface ProfileMetadata {
 
 export type StoredProfile<T extends Profile = Profile> = T & { readonly metadata: ProfileMetadata }
 
+export class ProfileNotFoundError extends Error {
+    public constructor(id: string) {
+        super(`Profile does not exist: ${id}`)
+    }
+}
+
 function getTelemetryForProfile(profile: StoredProfile<Profile> | undefined) {
     if (!profile) {
         return {}
@@ -245,7 +251,7 @@ export class ProfileStore {
         try {
             profile = this.getProfile(id)
             if (profile === undefined) {
-                throw new Error(`Profile does not exist: ${id}`)
+                throw new ProfileNotFoundError(id)
             }
         } catch (err) {
             // Always emit failures

packages/core/src/auth/secondaryAuth.ts

@@ -9,7 +9,7 @@ import * as vscode from 'vscode'
 import { getLogger } from '../shared/logger'
 import { cast, Optional } from '../shared/utilities/typeConstructors'
 import { Auth } from './auth'
-import { once, onceChanged } from '../shared/utilities/functionUtils'
+import { onceChanged } from '../shared/utilities/functionUtils'
 import { isNonNullable } from '../shared/utilities/tsUtils'
 import { ToolIdStateKey } from '../shared/globalState'
 import { Connection, getTelemetryMetadataForConn, SsoConnection, StatefulConnection } from './connection'
@@ -147,6 +147,40 @@ export class SecondaryAuth<T extends Connection = Connection> {
                 await this.clearSavedConnection()
             }
         })
+
+        const refreshConn = (event: string) => {
+            getLogger().debug(`secondaryAuth: detected ${event} event in sso cache, refreshing auth.`)
+            globals.clock.setTimeout(
+                telemetry.function_call.run(
+                    () => async () => {
+                        if (this.#savedConnection?.id === this.getStateConnectionId()) {
+                            // Someone updated our cache but the global state doesn't indicate anything new, so do nothing.
+                            // (it could have been us that updated the cache)
+                            getLogger().debug(
+                                `secondaryAuth: cache event did not update global state, no refresh is needed.`
+                            )
+                            return
+                        }
+                        await this.auth.restorePreviousSession()
+                        await this.restoreConnection(true)
+                    },
+                    {
+                        emit: false,
+                        functionId: { name: 'cacheWatchCallback', class: secondaryAuthClassName },
+                    }
+                ),
+                /**
+                 * The connection is first created and stored on disk, then it is stored in global state.
+                 * This creates a race condition for the callback, who listens to the disk event but
+                 * depends on the global state to make a connection decision. The time is arbitrary.
+                 *
+                 * TODO: fix this race condition.
+                 */
+                3000
+            )
+        }
+        this.auth.cacheWatcher.onDidCreate(() => refreshConn('create'))
+        this.auth.cacheWatcher.onDidDelete(() => refreshConn('delete'))
     }
 
     public get activeConnection(): T | undefined {
@@ -250,11 +284,17 @@ export class SecondaryAuth<T extends Connection = Connection> {
         return await addScopes(conn, extraScopes, this.auth)
     }
 
+    private hasRunRestoreConnection = false
+
     // Used to lazily restore persisted connections.
     // Kind of clunky. We need an async module loader layer to make things ergonomic.
-    public readonly restoreConnection: typeof this._restoreConnection = once(async () => this._restoreConnection())
     @withTelemetryContext({ name: 'restoreConnection', class: secondaryAuthClassName })
-    private async _restoreConnection(): Promise<T | undefined> {
+    public async restoreConnection(force: boolean = false): Promise<T | undefined> {
+        if (!force && (this.activeConnection !== undefined || this.hasRunRestoreConnection)) {
+            return
+        }
+        this.hasRunRestoreConnection = true
+
         try {
             return await telemetry.auth_modifyConnection.run(async (span) => {
                 span.record({
@@ -289,7 +329,7 @@ export class SecondaryAuth<T extends Connection = Connection> {
         // 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 = this.getStateConnectionId()
         if (id === undefined) {
             return
         }
@@ -335,6 +375,10 @@ export class SecondaryAuth<T extends Connection = Connection> {
             return conn
         }
     }
+
+    private getStateConnectionId() {
+        return cast(globals.globalState.get(this.key), Optional(String))
+    }
 }
 
 /**

packages/core/src/auth/sso/cache.ts

@@ -3,6 +3,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
+import * as vscode from 'vscode'
 import * as crypto from 'crypto'
 import * as path from 'path'
 import { getLogger } from '../../shared/logger/logger'
@@ -13,6 +14,7 @@ import { hasProps, selectFrom } from '../../shared/utilities/tsUtils'
 import { SsoToken, ClientRegistration } from './model'
 import { DevSettings } from '../../shared/settings'
 import { onceChanged } from '../../shared/utilities/functionUtils'
+import globals from '../../shared/extensionGlobals'
 
 interface RegistrationKey {
     readonly startUrl: string
@@ -42,6 +44,12 @@ export function getCache(directory = getCacheDir()): SsoCache {
     }
 }
 
+export function getCacheFileWatcher(directory = getCacheDir()) {
+    const watcher = vscode.workspace.createFileSystemWatcher(new vscode.RelativePattern(directory, '*.json'))
+    globals.context.subscriptions.push(watcher)
+    return watcher
+}
+
 export function getRegistrationCache(directory = getCacheDir()): KeyedCache<ClientRegistration, RegistrationKey> {
     // Compatability for older Toolkit versions (format on disk is unchanged)
     type StoredRegistration = Omit<ClientRegistration, 'expiresAt'> & { readonly expiresAt: string }

packages/core/src/codewhisperer/util/authUtil.ts

@@ -26,9 +26,10 @@ import {
     isIdcSsoConnection,
     hasExactScopes,
     getTelemetryMetadataForConn,
+    ProfileNotFoundError,
 } from '../../auth/connection'
 import { getLogger } from '../../shared/logger'
-import { Commands } from '../../shared/vscode/commands2'
+import { Commands, placeholder } from '../../shared/vscode/commands2'
 import { vsCodeState } from '../models/model'
 import { onceChanged, once } from '../../shared/utilities/functionUtils'
 import { indent } from '../../shared/utilities/textUtilities'
@@ -42,6 +43,7 @@ const localize = nls.loadMessageBundle()
 import { telemetry } from '../../shared/telemetry/telemetry'
 import { asStringifiedStack } from '../../shared/telemetry/spans'
 import { withTelemetryContext } from '../../shared/telemetry/util'
+import { focusAmazonQPanel } from '../../codewhispererChat/commands/registerCommands'
 
 /** Backwards compatibility for connections w pre-chat scopes */
 export const codeWhispererCoreScopes = [...scopesCodeWhispererCore]
@@ -134,14 +136,12 @@ export class AuthUtil {
                 Commands.tryExecute('aws.amazonq.updateReferenceLog'),
             ])
 
-            await setContext('aws.codewhisperer.connected', this.isConnected())
+            await this.setVscodeContextProps()
 
             // To check valid connection
             if (this.isValidEnterpriseSsoInUse() || (this.isBuilderIdInUse() && !this.isConnectionExpired())) {
-                // start the feature config polling job
                 await showAmazonQWalkthroughOnce()
             }
-            await this.setVscodeContextProps()
         })
     })
 
@@ -254,8 +254,16 @@ export class AuthUtil {
             throw new ToolkitError('Connection is not an SSO connection', { code: 'BadConnectionType' })
         }
 
-        const bearerToken = await this.conn.getToken()
-        return bearerToken.accessToken
+        try {
+            const bearerToken = await this.conn.getToken()
+            return bearerToken.accessToken
+        } catch (err) {
+            if (err instanceof ProfileNotFoundError) {
+                // Expected that connection would be deleted by conn.getToken()
+                void focusAmazonQPanel.execute(placeholder, 'profileNotFoundSignout')
+            }
+            throw err
+        }
     }
 
     @withTelemetryContext({ name: 'getCredentials', class: authClassName })

packages/core/src/login/webview/vue/amazonq/backend_amazonq.ts

@@ -3,7 +3,13 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 import * as vscode from 'vscode'
-import { AwsConnection, Connection, getTelemetryMetadataForConn, isSsoConnection } from '../../../../auth/connection'
+import {
+    AwsConnection,
+    Connection,
+    SsoConnection,
+    getTelemetryMetadataForConn,
+    isSsoConnection,
+} from '../../../../auth/connection'
 import { AuthUtil } from '../../../../codewhisperer/util/authUtil'
 import { CommonAuthWebview } from '../backend'
 import { awsIdSignIn } from '../../../../codewhisperer/util/showSsoPrompt'
@@ -178,6 +184,12 @@ export class AmazonQLoginWebview extends CommonAuthWebview {
         this.emitAuthMetric()
     }
 
+    async listSsoConnections(): Promise<SsoConnection[]> {
+        // Amazon Q only supports 1 connection at a time,
+        // so there isn't a need to de-duplicate connections.
+        return []
+    }
+
     override startIamCredentialSetup(
         profileName: string,
         accessKey: string,

packages/core/src/login/webview/vue/backend.ts

@@ -16,7 +16,6 @@ import {
     AwsConnection,
     Connection,
     hasScopes,
-    isSsoConnection,
     scopesCodeCatalyst,
     scopesCodeWhispererChat,
     scopesSsoAccountAccess,
@@ -189,9 +188,8 @@ export abstract class CommonAuthWebview extends VueWebview {
 
     abstract signout(): Promise<void>
 
-    async listSsoConnections(): Promise<SsoConnection[]> {
-        return (await Auth.instance.listConnections()).filter((conn) => isSsoConnection(conn)) as SsoConnection[]
-    }
+    /** List current connections known by the extension for the purpose of preventing duplicates. */
+    abstract listSsoConnections(): Promise<SsoConnection[]>
 
     /**
      * Emit stored metric metadata. Does not reset the stored metric metadata, because it

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

@@ -10,9 +10,11 @@ import { CommonAuthWebview } from '../backend'
 import {
     AwsConnection,
     Connection,
+    SsoConnection,
     TelemetryMetadata,
     createSsoProfile,
     getTelemetryMetadataForConn,
+    isSsoConnection,
 } from '../../../../auth/connection'
 import { Auth } from '../../../../auth/auth'
 import { CodeCatalystAuthenticationProvider } from '../../../../codecatalyst/auth'
@@ -145,6 +147,10 @@ export class ToolkitLoginWebview extends CommonAuthWebview {
         return connections
     }
 
+    async listSsoConnections(): Promise<SsoConnection[]> {
+        return (await Auth.instance.listConnections()).filter((conn) => isSsoConnection(conn)) as SsoConnection[]
+    }
+
     override reauthenticateConnection(): Promise<undefined> {
         throw new Error('Method not implemented.')
     }

packages/toolkit/.changes/next-release/Bug Fix-511c0b65-6f7b-4c1a-a0c5-616c3ec802f4.json

@@ -0,0 +1,4 @@
+{
+	"type": "Bug Fix",
+	"description": "Login state not updating across multiple VS Code windows."
+}