fix(auth): handle network errors during restore/getChatAuthState (aws#5476)

## Problem

During extension startup, the auth system tries to restore the
connection. This may result in a `refreshToken` call. If a network error
happens during this, the connection restore fails and no connection is
set to active. Then, Amazon Q will clean up any "leftover" connections,
including this one, and the user will be completely logged out. The
connection is technically valid despite having an expired token, and
another call may clear up the network error and allow us to get better
token.

## Solution
The auth system will restore this connection even if it runs into a
network error. Then, Amazon Q will display as if its logged in normally.
If the network error was one-off, then there is no indication that there
was an issue to the user. If there are still network issues, using chat
or inline-suggestions will make an error display until `refreshToken` is
successful (or the registration expires).

Also, some telemetry updates were included to help make this visible.

Side effects:
- On session startup (with network error):
  - No notification for new CW customizations
  - "featureConfig" AKA A/B testing? will not initialize.

<!---
    REMINDER:
    - Read CONTRIBUTING.md first.
    - Add test coverage for your changes.
    - Update the changelog using `npm run newChange`.
    - Link to related issues/commits.
    - Testing: how did you test your changes?
- Screenshots (if the pull request is related to UI/UX then please
include light and dark theme screenshots)
-->

## License

By submitting this pull request, 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-910f825b-6832-42ac-80b4-ffcac58d15c5.json

@@ -0,0 +1,4 @@
+{
+	"type": "Bug Fix",
+	"description": "Auth: Users may be silently logged out due to network issues when starting the extension."
+}

packages/amazonq/src/extension.ts

@@ -12,6 +12,7 @@ import {
     isAnySsoConnection,
 } from 'aws-core-vscode/auth'
 import {
+    AuthState,
     AuthUtil,
     activate as activateCodeWhisperer,
     shutdown as shutdownCodeWhisperer,
@@ -34,6 +35,7 @@ import {
     globals,
     initialize,
     initializeComputeRegion,
+    isNetworkError,
     messages,
     placeholder,
     setContext,
@@ -141,32 +143,52 @@ export async function activateAmazonQCommon(context: vscode.ExtensionContext, is
         }, 1000)
     }
 
-    await telemetry.auth_userState.run(async () => {
-        telemetry.record({ passive: true })
-
-        const firstUse = AuthUtils.ExtensionUse.instance.isFirstUse()
-        const wasUpdated = AuthUtils.ExtensionUse.instance.wasUpdated()
-
-        if (firstUse) {
-            telemetry.record({ source: ExtStartUpSources.firstStartUp })
-        } else if (wasUpdated) {
-            telemetry.record({ source: ExtStartUpSources.update })
-        } else {
-            telemetry.record({ source: ExtStartUpSources.reload })
-        }
-
-        const authState = (await AuthUtil.instance.getChatAuthState()).codewhispererChat
-        const currConn = AuthUtil.instance.conn
-        if (currConn !== undefined && !isAnySsoConnection(currConn)) {
-            getLogger().error(`Current Amazon Q connection is not SSO, type is: %s`, currConn?.type)
-        }
-
-        telemetry.record({
-            authStatus: authState === 'connected' || authState === 'expired' ? authState : 'notConnected',
-            authEnabledConnections: AuthUtils.getAuthFormIdsFromConnection(currConn).join(','),
-            ...(await getTelemetryMetadataForConn(currConn)),
+    await telemetry.auth_userState
+        .run(async () => {
+            telemetry.record({ passive: true })
+
+            const firstUse = AuthUtils.ExtensionUse.instance.isFirstUse()
+            const wasUpdated = AuthUtils.ExtensionUse.instance.wasUpdated()
+
+            if (firstUse) {
+                telemetry.record({ source: ExtStartUpSources.firstStartUp })
+            } else if (wasUpdated) {
+                telemetry.record({ source: ExtStartUpSources.update })
+            } else {
+                telemetry.record({ source: ExtStartUpSources.reload })
+            }
+
+            let authState: AuthState = 'disconnected'
+            try {
+                // May call connection validate functions that try to refresh the token.
+                // This could result in network errors.
+                authState = (await AuthUtil.instance.getChatAuthState(false)).codewhispererChat
+            } catch (err) {
+                if (
+                    isNetworkError(err) &&
+                    AuthUtil.instance.conn &&
+                    AuthUtil.instance.auth.getConnectionState(AuthUtil.instance.conn) === 'valid'
+                ) {
+                    authState = 'connectedWithNetworkError'
+                } else {
+                    throw err
+                }
+            }
+            const currConn = AuthUtil.instance.conn
+            if (currConn !== undefined && !isAnySsoConnection(currConn)) {
+                getLogger().error(`Current Amazon Q connection is not SSO, type is: %s`, currConn?.type)
+            }
+
+            telemetry.record({
+                authStatus:
+                    authState === 'connected' || authState === 'expired' || authState === 'connectedWithNetworkError'
+                        ? authState
+                        : 'notConnected',
+                authEnabledConnections: AuthUtils.getAuthFormIdsFromConnection(currConn).join(','),
+                ...(await getTelemetryMetadataForConn(currConn)),
+            })
         })
-    })
+        .catch((err) => getLogger().error('Error collecting telemetry for auth_userState: %s', err))
 }
 
 export async function deactivateCommon() {

packages/core/src/auth/auth.ts

@@ -601,6 +601,11 @@ export class Auth implements AuthService, ConnectionManager {
         })
     }
 
+    /**
+     * Updates the state of the connection based on a series of validations.
+     * Will throw for network errors encoutered during connection checks, but the state will not be
+     * invalidated for these errors.
+     */
     @withTelemetryContext({ name: 'validateConnection', class: authClassName })
     private async validateConnection<T extends Profile>(id: Connection['id'], profile: StoredProfile<T>) {
         const runCheck = async () => {
@@ -843,6 +848,7 @@ export class Auth implements AuthService, ConnectionManager {
         if (isNetworkError(e)) {
             throw new ToolkitError('Failed to update connection due to networking issues', {
                 cause: e,
+                code: e.code,
             })
         }
     }

packages/core/src/auth/secondaryAuth.ts

@@ -14,9 +14,10 @@ import { isNonNullable } from '../shared/utilities/tsUtils'
 import { ToolIdStateKey } from '../shared/globalState'
 import { Connection, getTelemetryMetadataForConn, SsoConnection, StatefulConnection } from './connection'
 import { indent } from '../shared/utilities/textUtilities'
-import { telemetry } from '../shared/telemetry/telemetry'
+import { AuthStatus, telemetry } from '../shared/telemetry/telemetry'
 import { asStringifiedStack } from '../shared/telemetry/spans'
 import { withTelemetryContext } from '../shared/telemetry/util'
+import { isNetworkError } from '../shared/errors'
 
 export type ToolId = 'codecatalyst' | 'codewhisperer' | 'testId'
 
@@ -262,25 +263,21 @@ export class SecondaryAuth<T extends Connection = Connection> {
                     connectionState: 'undefined',
                 })
                 await this.auth.tryAutoConnect()
-                this.#savedConnection = await this.loadSavedConnection()
+                this.#savedConnection = await this._loadSavedConnection()
                 this.#onDidChangeActiveConnection.fire(this.activeConnection)
 
-                const conn = this.#activeConnection
-                if (conn) {
-                    telemetry.record({
-                        connectionState: this.auth.getConnectionState(conn),
-                        ...(await getTelemetryMetadataForConn(conn)),
-                    })
-                }
-
+                telemetry.record(await getTelemetryMetadataForConn(this.#savedConnection))
                 return this.#savedConnection
             })
         } catch (err) {
             getLogger().warn(`auth (${this.toolId}): failed to restore connection: %s`, err)
         }
     }
 
-    private async loadSavedConnection() {
+    /**
+     * 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
@@ -297,7 +294,36 @@ export class SecondaryAuth<T extends Connection = Connection> {
             getLogger().warn(`auth (${this.toolId}): saved connection "${this.key}" is not valid`)
             await globalState.update(this.key, undefined)
         } else {
-            await this.auth.refreshConnectionState(conn)
+            const getAuthStatus = (state: ReturnType<typeof this.auth.getConnectionState>): AuthStatus => {
+                return state === 'invalid' ? 'expired' : 'connected'
+            }
+
+            let connectionState = this.auth.getConnectionState(conn)
+
+            // This function is expected to be called in the context of restoreConnection()
+            telemetry.auth_modifyConnection.record({
+                connectionState,
+                authStatus: getAuthStatus(connectionState),
+            })
+
+            try {
+                await this.auth.refreshConnectionState(conn)
+
+                connectionState = this.auth.getConnectionState(conn)
+                telemetry.auth_modifyConnection.record({
+                    connectionState,
+                    authStatus: getAuthStatus(connectionState),
+                })
+            } catch (err) {
+                // The purpose of this function is load a saved connection into memory, not manage the state.
+                // If updating the state fails, then we should delegate downstream to handle getting the proper state.
+                getLogger().error('loadSavedConnection: Failed to refresh connection state: %s', err)
+                if (isNetworkError(err) && connectionState === 'valid') {
+                    telemetry.auth_modifyConnection.record({
+                        authStatus: 'connectedWithNetworkError',
+                    })
+                }
+            }
             return conn
         }
     }

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

@@ -6,7 +6,7 @@
 import * as vscode from 'vscode'
 import * as localizedText from '../../shared/localizedText'
 import { Auth } from '../../auth/auth'
-import { ToolkitError } from '../../shared/errors'
+import { ToolkitError, isNetworkError, tryRun } from '../../shared/errors'
 import { getSecondaryAuth, setScopes } from '../../auth/secondaryAuth'
 import { isCloud9, isSageMaker } from '../../shared/extensionUtilities'
 import { AmazonQPromptSettings } from '../../shared/settings'
@@ -403,12 +403,24 @@ export class AuthUtil {
      * Asynchronously returns a snapshot of the overall auth state of CodeWhisperer + Chat features.
      * It guarantees the latest state is correct at the risk of modifying connection state.
      * If this guarantee is not required, use sync method getChatAuthStateSync()
+     *
+     * By default, network errors are ignored when determining auth state since they may be silently
+     * recoverable later.
      */
     @withTelemetryContext({ name: 'getChatAuthState', class: authClassName })
-    public async getChatAuthState(): Promise<FeatureAuthState> {
+    public async getChatAuthState(ignoreNetErr: boolean = true): Promise<FeatureAuthState> {
         // The state of the connection may not have been properly validated
         // and the current state we see may be stale, so refresh for latest state.
-        await this.auth.refreshConnectionState(this.conn)
+        if (ignoreNetErr) {
+            await tryRun(
+                () => this.auth.refreshConnectionState(this.conn),
+                (err) => !isNetworkError(err),
+                'getChatAuthState: Cannot refresh connection state due to network error: %s'
+            )
+        } else {
+            await this.auth.refreshConnectionState(this.conn)
+        }
+
         return this.getChatAuthStateSync(this.conn)
     }
 
@@ -525,6 +537,11 @@ export const AuthStates = {
      * Eg: We are currently using Builder ID, but must use Identity Center.
      */
     unsupported: 'unsupported',
+    /**
+     * The current connection exists and isn't expired,
+     * but fetching/refreshing the token resulted in a network error.
+     */
+    connectedWithNetworkError: 'connectedWithNetworkError',
 } as const
 const Features = {
     codewhispererCore: 'codewhispererCore',

packages/core/src/shared/errors.ts

@@ -14,6 +14,7 @@ import type * as nodefs from 'fs'
 import type * as os from 'os'
 import { CodeWhispererStreamingServiceException } from '@amzn/codewhisperer-streaming'
 import { driveLetterRegex } from './utilities/pathUtils'
+import { getLogger } from './logger/logger'
 
 let _username = 'unknown-user'
 let _isAutomation = false
@@ -911,3 +912,38 @@ export function getReasonFromSyntaxError(err: Error): string | undefined {
 
     return undefined
 }
+
+/**
+ * Run a function and swallow any errors that are not specified by `shouldThrow`
+ */
+export function tryRun<T>(fn: () => T, shouldThrow: (err: Error) => boolean, logMsg?: string): T | undefined
+export function tryRun<T>(
+    fn: () => Promise<T>,
+    shouldThrow: (err: Error) => boolean,
+    logMsg?: string
+): Promise<T> | undefined
+export function tryRun<T>(
+    fn: () => T | Promise<T>,
+    shouldThrow: (err: Error) => boolean,
+    logMsg?: string
+): T | Promise<T | void> | undefined {
+    // The function signature pattern will accept both async and non-async types.
+
+    const catchErr = (err: Error) => {
+        if (shouldThrow(err)) {
+            throw err
+        }
+
+        getLogger().error(logMsg ?? 'unknown caller: Error ignored: %s', err)
+    }
+
+    try {
+        const result = fn()
+        if (result instanceof Promise) {
+            return result.catch(catchErr)
+        }
+        return result
+    } catch (error: any) {
+        catchErr(error)
+    }
+}

packages/core/src/shared/telemetry/vscodeTelemetry.json

@@ -248,7 +248,7 @@
         {
             "name": "authStatus",
             "type": "string",
-            "allowedValues": ["connected", "notConnected", "expired"],
+            "allowedValues": ["connected", "notConnected", "expired", "connectedWithNetworkError"],
             "description": "Status of the an auth connection."
         },
         {
@@ -1068,6 +1068,10 @@
                     "type": "authScopes",
                     "required": false
                 },
+                {
+                    "type": "authStatus",
+                    "required": false
+                },
                 {
                     "type": "connectionState",
                     "required": true

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

@@ -289,6 +289,7 @@ describe('Auth', function () {
             const networkError = new ToolkitError('test', { code: 'ETIMEDOUT' })
             const expectedError = new ToolkitError('Failed to update connection due to networking issues', {
                 cause: networkError,
+                code: 'ETIMEDOUT',
             })
             const conn = await auth.createConnection(ssoProfile)
             auth.getTestTokenProvider(conn)?.getToken.rejects(networkError)

packages/core/src/test/shared/errors.test.ts

@@ -16,6 +16,7 @@ import {
     resolveErrorMessageToDisplay,
     scrubNames,
     ToolkitError,
+    tryRun,
     UnknownError,
 } from '../../shared/errors'
 import { CancellationError } from '../../shared/utilities/timeoutUtils'
@@ -558,3 +559,49 @@ describe('util', function () {
         assert.deepStrictEqual(scrubNames('unix ~/.aws/config failed', fakeUser), 'unix ~/.aws/config failed')
     })
 })
+
+describe('errors.tryRun()', function () {
+    it('swallows error from sync fn', function () {
+        const err = new Error('err')
+        tryRun(
+            () => {
+                throw err
+            },
+            () => false
+        )
+    })
+
+    it('swallows error from async fn', async function () {
+        const err = new Error('err')
+        await tryRun(
+            async () => {
+                throw err
+            },
+            () => false
+        )
+    })
+
+    it('throws error from sync fn', function () {
+        const err = new Error('err')
+        assert.throws(() => {
+            tryRun(
+                () => {
+                    throw err
+                },
+                () => true
+            )
+        }, err)
+    })
+
+    it('throws error from async fn', async function () {
+        const err = new Error('err')
+        await assert.rejects(async () => {
+            await tryRun(
+                async () => {
+                    throw err
+                },
+                () => true
+            )
+        }, err)
+    })
+})

packages/toolkit/.changes/next-release/Bug Fix-a9d6d04d-f861-49ce-8d08-ff50b3e0bb2b.json

@@ -0,0 +1,4 @@
+{
+	"type": "Bug Fix",
+	"description": "Auth: Users may be silently logged out due to network issues when starting the extension."
+}