telemetry: add function stack to auth connection change (#5465)

* telemetry: add function stack to auth connection change

Problem:

When an auth connection changes, eg goes from valid to invalid, there were
cases for some users where it happened unexpectedly. It was difficult to
answer why the users connection became invalid

Solution:

Wrap many auth related functions with their "function entry" context.
This will then allow the final telemetry event in `Auth.updateConnectionState()`
to report the entire call stack in its `source` field for the metric `auth_modifyConnection`.
What this results in is a string like: `AmazonQ#signoutButton:AuthUtil#signout:Auth#signout,deleteConnection,updateConnectionState`
which tells us the path taken to get to updating the connection state.

In telemetry we can then follow the change in a users connection state and hopefully determine why they
may have expired prematurely since the functions which can expire a connection will now emit telemetry.

Implementation Details:

- A decorator, `withTelemetryContext` was created to wrap a Class Method with its context for telemetry
  - Every method that uses it adds their function name and class as context
  - Under the hood it wraps the Class Method in a `telemetry.function_call.run()` with some predefined values
  - It does not emit a telemetry event by default, as we only need it to contribute to the function stack
  - This decorator significantly reduces the diff compared to using the explicit `run()` since all method code
    would need to be explicitly wrapped in it.
  - **Note that decorators only work with class methods**
- For non class methods we just use `telemetry.function_call.run()`

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

* fix: failing unit test

Unit test was failing here in CI with a "Reference error", and once
I removed the decorator it was working fine again.

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

* update tests for decorator

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

* add context for dev mode expire connections

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

* add more tests for asserting expected telemetry in auth

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

* add sessionDuration to connection state change

When a connection becomes invalid we will add the sessionDuration
to the `auth_modifyConnection` metric emitted.

This will allow us to filter for SSO connections that prematurely became invalid.

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

---------

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

Author: nkomonen-amazon

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

@@ -290,7 +290,7 @@ describe('AuthUtil', async function () {
         )
         await authUtil.secondaryAuth.useNewConnection(conn3)
 
-        await authUtil.clearExtraConnections('test') // method under test
+        await authUtil.clearExtraConnections() // method under test
 
         // Only the conn that AuthUtil is using is remaining
         assert.deepStrictEqual(

packages/core/src/auth/auth.ts

@@ -62,6 +62,8 @@ import {
 import { isSageMaker, isCloud9, isAmazonQ } from '../shared/extensionUtilities'
 import { telemetry } from '../shared/telemetry/telemetry'
 import { randomUUID } from '../shared/crypto'
+import { asStringifiedStack } from '../shared/telemetry/spans'
+import { withTelemetryContext } from '../shared/telemetry/util'
 
 interface AuthService {
     /**
@@ -124,6 +126,9 @@ export type AuthType = Auth
 export type DeletedConnection = { connId: Connection['id']; storedProfile?: StoredProfile }
 type DeclaredConnection = Pick<SsoProfile, 'ssoRegion' | 'startUrl'> & { source: string }
 
+// Must be declared outside of class for use by decorator
+const authClassName = 'Auth'
+
 export class Auth implements AuthService, ConnectionManager {
     readonly #ssoCache = getCache()
     readonly #validationErrors = new Map<Connection['id'], Error>()
@@ -162,6 +167,7 @@ export class Auth implements AuthService, ConnectionManager {
         return Object.values(this._declaredConnections)
     }
 
+    @withTelemetryContext({ name: 'restorePreviousSession', class: authClassName })
     public async restorePreviousSession(): Promise<Connection | undefined> {
         const id = this.store.getCurrentProfileId()
         if (id === undefined) {
@@ -177,6 +183,7 @@ export class Auth implements AuthService, ConnectionManager {
 
     public async reauthenticate({ id }: Pick<SsoConnection, 'id'>, invalidate?: boolean): Promise<SsoConnection>
     public async reauthenticate({ id }: Pick<IamConnection, 'id'>, invalidate?: boolean): Promise<IamConnection>
+    @withTelemetryContext({ name: 'reauthenticate', class: authClassName })
     public async reauthenticate({ id }: Pick<Connection, 'id'>, invalidate?: boolean): Promise<Connection> {
         const shouldInvalidate = invalidate ?? true
         const profile = this.store.getProfileOrThrow(id)
@@ -195,6 +202,7 @@ export class Auth implements AuthService, ConnectionManager {
 
     public async useConnection({ id }: Pick<SsoConnection, 'id'>): Promise<SsoConnection>
     public async useConnection({ id }: Pick<IamConnection, 'id'>): Promise<IamConnection>
+    @withTelemetryContext({ name: 'useConnection', class: authClassName })
     public async useConnection({ id }: Pick<Connection, 'id'>): Promise<Connection> {
         await this.refreshConnectionState({ id })
 
@@ -211,6 +219,7 @@ export class Auth implements AuthService, ConnectionManager {
         return conn
     }
 
+    @withTelemetryContext({ name: 'logout', class: authClassName })
     public async logout(): Promise<void> {
         if (this.activeConnection === undefined) {
             return
@@ -222,6 +231,7 @@ export class Auth implements AuthService, ConnectionManager {
         this.#onDidChangeActiveConnection.fire(undefined)
     }
 
+    @withTelemetryContext({ name: 'listConnections', class: authClassName })
     public async listConnections(): Promise<Connection[]> {
         await loadIamProfilesIntoStore(this.store, this.iamProfileProvider)
 
@@ -281,6 +291,7 @@ export class Auth implements AuthService, ConnectionManager {
     }
 
     public async createConnection(profile: SsoProfile): Promise<SsoConnection>
+    @withTelemetryContext({ name: 'createConnection', class: authClassName })
     public async createConnection(profile: Profile): Promise<Connection> {
         if (profile.type === 'iam') {
             throw new Error('Creating IAM connections is not supported')
@@ -311,6 +322,7 @@ export class Auth implements AuthService, ConnectionManager {
         }
     }
 
+    @withTelemetryContext({ name: 'deleteConnection', class: authClassName })
     public async deleteConnection(connection: Pick<Connection, 'id'>): Promise<void> {
         const connId = connection.id
         const profile = this.store.getProfile(connId)
@@ -343,6 +355,7 @@ export class Auth implements AuthService, ConnectionManager {
      *
      * Forget about a connection without logging out, invalidating it, or deleting it from disk.
      */
+    @withTelemetryContext({ name: 'forgetConnection', class: authClassName })
     public async forgetConnection(connection: Pick<Connection, 'id'>): Promise<void> {
         const connId = connection.id
         const profile = this.store.getProfile(connId)
@@ -355,6 +368,7 @@ export class Auth implements AuthService, ConnectionManager {
         this.#onDidDeleteConnection.fire({ connId, storedProfile: profile })
     }
 
+    @withTelemetryContext({ name: 'clearStaleLinkedIamConnections', class: authClassName })
     private async clearStaleLinkedIamConnections() {
         // Updates our store, evicting stale IAM credential profiles if the
         // SSO they are linked to was removed.
@@ -373,6 +387,7 @@ export class Auth implements AuthService, ConnectionManager {
      *
      * Put the SSO connection in to an expired state
      */
+    @withTelemetryContext({ name: 'expireConnection', class: authClassName })
     public async expireConnection(conn: Pick<SsoConnection, 'id'>): Promise<void> {
         getLogger().info(`auth: Expiring connection ${conn.id}`)
         const profile = this.store.getProfileOrThrow(conn.id)
@@ -398,6 +413,7 @@ export class Auth implements AuthService, ConnectionManager {
      * Alternatively you can use the `getToken()` call on an SSO connection to do the same thing,
      * but it will additionally prompt for reauthentication if the connection is invalid.
      */
+    @withTelemetryContext({ name: 'refreshConnectionState', class: authClassName })
     public async refreshConnectionState(connection?: Pick<Connection, 'id'>): Promise<undefined> {
         if (connection === undefined) {
             return
@@ -416,6 +432,7 @@ export class Auth implements AuthService, ConnectionManager {
         profile: SsoProfile,
         invalidate?: boolean
     ): Promise<SsoConnection>
+    @withTelemetryContext({ name: 'updateConnection', class: authClassName })
     public async updateConnection(
         connection: Pick<Connection, 'id'>,
         profile: Profile,
@@ -458,6 +475,7 @@ export class Auth implements AuthService, ConnectionManager {
      *
      * @returns undefined if authentication succeeds, otherwise object with error info
      */
+    @withTelemetryContext({ name: 'authenticateData', class: authClassName })
     public async authenticateData(data: StaticProfile): Promise<StaticProfileKeyErrorMessage | undefined> {
         const tempId = await this.addTempCredential(data)
         const tempIdString = asString(tempId)
@@ -507,6 +525,7 @@ export class Auth implements AuthService, ConnectionManager {
      * For SSO, this involves an API call to clear server-side state. The call happens
      * before the local token(s) are cleared as they are needed in the request.
      */
+    @withTelemetryContext({ name: 'invalidateConnection', class: authClassName })
     private async invalidateConnection(id: Connection['id'], opt?: { skipGlobalLogout?: boolean }) {
         getLogger().info(`auth: Invalidating connection: ${id}`)
         const profile = this.store.getProfileOrThrow(id)
@@ -534,6 +553,7 @@ export class Auth implements AuthService, ConnectionManager {
         await this.updateConnectionState(id, 'invalid')
     }
 
+    @withTelemetryContext({ name: 'updateConnectionState', class: authClassName })
     private async updateConnectionState(id: Connection['id'], connectionState: ProfileMetadata['connectionState']) {
         getLogger().info(`auth: Updating connection state of ${id} to ${connectionState}`)
 
@@ -544,24 +564,44 @@ export class Auth implements AuthService, ConnectionManager {
 
         const oldProfile = this.store.getProfileOrThrow(id)
         if (oldProfile.metadata.connectionState === connectionState) {
+            // new state is same as old state, no need to do anything
             return oldProfile
         }
 
-        const profile = await this.store.updateMetadata(id, { connectionState })
-        if (connectionState !== 'invalid') {
-            this.#validationErrors.delete(id)
-            this.#invalidCredentialsTimeouts.get(id)?.dispose()
-        }
+        // Emit an event for when the connection changes state. This the the root of the
+        // state change. We rely on functions higher in the stack to have added their contexts
+        // so this event has useful information in the `source` field.
+        return telemetry.auth_modifyConnection.run(async (span) => {
+            // if we have an Sso session that became invalid, we will add its session duration to the event
+            const ssoSessionDuration =
+                connectionState === 'invalid' && oldProfile.type === 'sso'
+                    ? this.getSsoTokenProvider(id, oldProfile).getSessionDuration()
+                    : undefined
+            span.record({
+                action: 'updateConnectionState',
+                connectionState,
+                source: asStringifiedStack(telemetry.getFunctionStack()),
+                credentialStartUrl: oldProfile.type === 'sso' ? oldProfile.startUrl : undefined,
+                sessionDuration: ssoSessionDuration,
+            })
 
-        if (this.#activeConnection?.id === id) {
-            this.#activeConnection.state = connectionState
-            this.#onDidChangeActiveConnection.fire(this.#activeConnection)
-        }
-        this.#onDidChangeConnectionState.fire({ id, state: connectionState })
+            const profile = await this.store.updateMetadata(id, { connectionState })
+            if (connectionState !== 'invalid') {
+                this.#validationErrors.delete(id)
+                this.#invalidCredentialsTimeouts.get(id)?.dispose()
+            }
 
-        return profile
+            if (this.#activeConnection?.id === id) {
+                this.#activeConnection.state = connectionState
+                this.#onDidChangeActiveConnection.fire(this.#activeConnection)
+            }
+            this.#onDidChangeConnectionState.fire({ id, state: connectionState })
+
+            return profile
+        })
     }
 
+    @withTelemetryContext({ name: 'validateConnection', class: authClassName })
     private async validateConnection<T extends Profile>(id: Connection['id'], profile: StoredProfile<T>) {
         const runCheck = async () => {
             if (profile.type === 'sso') {
@@ -747,6 +787,7 @@ export class Auth implements AuthService, ConnectionManager {
     }
 
     private readonly authenticate = keyedDebounce(this._authenticate.bind(this))
+    @withTelemetryContext({ name: 'authenticate', class: authClassName })
     private async _authenticate<T>(id: Connection['id'], callback: () => Promise<T>, invalidate?: boolean): Promise<T> {
         const originalState = this.getConnectionState({ id }) ?? 'unauthenticated'
         await this.updateConnectionState(id, 'authenticating')
@@ -780,6 +821,7 @@ export class Auth implements AuthService, ConnectionManager {
     }
 
     private readonly getToken = keyedDebounce(this._getToken.bind(this))
+    @withTelemetryContext({ name: '_getToken', class: authClassName })
     private async _getToken(id: Connection['id'], provider: SsoAccessTokenProvider): Promise<SsoToken> {
         const token = await provider.getToken().catch((err) => {
             this.throwOnNetworkError(err)
@@ -817,6 +859,7 @@ 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)
@@ -861,7 +904,9 @@ export class Auth implements AuthService, ConnectionManager {
         return this.authenticate(id, refresh)
     }
 
-    public readonly tryAutoConnect = once(async () => {
+    public readonly tryAutoConnect = once(async () => this._tryAutoConnect())
+    @withTelemetryContext({ name: 'tryAutoConnect', class: authClassName })
+    private async _tryAutoConnect() {
         if (this.activeConnection !== undefined) {
             return
         }
@@ -940,7 +985,7 @@ export class Auth implements AuthService, ConnectionManager {
                 }
             }
         }
-    })
+    }
 
     static #instance: Auth | undefined
     public static get instance() {

packages/core/src/auth/secondaryAuth.ts

@@ -15,6 +15,8 @@ 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 { asStringifiedStack } from '../shared/telemetry/spans'
+import { withTelemetryContext } from '../shared/telemetry/util'
 
 export type ToolId = 'codecatalyst' | 'codewhisperer' | 'testId'
 
@@ -78,6 +80,8 @@ export function getAllConnectionsInUse(auth: Auth): StatefulConnection[] {
 const onDidChangeConnectionsEmitter = new vscode.EventEmitter<void>()
 export const onDidChangeConnections = onDidChangeConnectionsEmitter.event
 
+// Variable must be declared outside of class to work with decorators
+const secondaryAuthClassName = 'SecondaryAuth'
 /**
  * Enables a tool to bind to a connection independently from the global {@link Auth} service.
  *
@@ -187,6 +191,7 @@ export class SecondaryAuth<T extends Connection = Connection> {
      * Globally deletes the connection that this secondary auth is using,
      * effectively doing a signout.
      */
+    @withTelemetryContext({ name: 'deleteConnection', class: secondaryAuthClassName })
     public async deleteConnection() {
         if (this.activeConnection) {
             await this.auth.deleteConnection(this.activeConnection)
@@ -228,6 +233,7 @@ export class SecondaryAuth<T extends Connection = Connection> {
         this.#onDidChangeActiveConnection.fire(undefined)
     }
 
+    @withTelemetryContext({ name: 'useNewConnection', class: secondaryAuthClassName })
     public async useNewConnection(conn: T): Promise<T> {
         await this.saveConnection(conn)
         if (this.auth.activeConnection === undefined) {
@@ -245,11 +251,13 @@ export class SecondaryAuth<T extends Connection = Connection> {
 
     // Used to lazily restore persisted connections.
     // Kind of clunky. We need an async module loader layer to make things ergonomic.
-    public readonly restoreConnection: (source?: string) => Promise<T | undefined> = once(async (source?: string) => {
+    public readonly restoreConnection: typeof this._restoreConnection = once(async () => this._restoreConnection())
+    @withTelemetryContext({ name: 'restoreConnection', class: secondaryAuthClassName })
+    private async _restoreConnection(): Promise<T | undefined> {
         try {
             return await telemetry.auth_modifyConnection.run(async () => {
                 telemetry.record({
-                    source,
+                    source: asStringifiedStack(telemetry.getFunctionStack()),
                     action: 'restore',
                     connectionState: 'undefined',
                 })
@@ -270,7 +278,7 @@ export class SecondaryAuth<T extends Connection = Connection> {
         } catch (err) {
             getLogger().warn(`auth (${this.toolId}): failed to restore connection: %s`, err)
         }
-    })
+    }
 
     private async loadSavedConnection() {
         // TODO: fix this
@@ -309,28 +317,38 @@ export class SecondaryAuth<T extends Connection = Connection> {
  * Note: This should exist in connection.ts or utils.ts, but due to circular dependencies, it must go here.
  */
 export async function addScopes(conn: SsoConnection, extraScopes: string[], auth = Auth.instance) {
-    const oldScopes = conn.scopes ?? []
-    const newScopes = Array.from(new Set([...oldScopes, ...extraScopes]))
-
-    const updatedConn = await setScopes(conn, newScopes, auth)
-
-    try {
-        return await auth.reauthenticate(updatedConn, false)
-    } catch (e) {
-        // We updated the connection scopes pre-emptively, but if there is some issue (e.g. user cancels,
-        // InvalidGrantException, etc), then we need to revert to the old connection scopes. Otherwise,
-        // this could soft-lock users into a broken connection that cannot be re-authenticated without
-        // first deleting the connection.
-        await setScopes(conn, oldScopes, auth)
-        throw e
-    }
+    return telemetry.function_call.run(
+        async () => {
+            const oldScopes = conn.scopes ?? []
+            const newScopes = Array.from(new Set([...oldScopes, ...extraScopes]))
+
+            const updatedConn = await setScopes(conn, newScopes, auth)
+
+            try {
+                return await auth.reauthenticate(updatedConn, false)
+            } catch (e) {
+                // We updated the connection scopes pre-emptively, but if there is some issue (e.g. user cancels,
+                // InvalidGrantException, etc), then we need to revert to the old connection scopes. Otherwise,
+                // this could soft-lock users into a broken connection that cannot be re-authenticated without
+                // first deleting the connection.
+                await setScopes(conn, oldScopes, auth)
+                throw e
+            }
+        },
+        { emit: false, functionId: { name: 'addScopesSecondaryAuth' } }
+    )
 }
 
 export function setScopes(conn: SsoConnection, scopes: string[], auth = Auth.instance): Promise<SsoConnection> {
-    return auth.updateConnection(conn, {
-        type: 'sso',
-        scopes,
-        startUrl: conn.startUrl,
-        ssoRegion: conn.ssoRegion,
-    })
+    return telemetry.function_call.run(
+        () => {
+            return auth.updateConnection(conn, {
+                type: 'sso',
+                scopes,
+                startUrl: conn.startUrl,
+                ssoRegion: conn.ssoRegion,
+            })
+        },
+        { emit: false, functionId: { name: 'setScopesSecondaryAuth' } }
+    )
 }