Merge master into feature/dynamodb

Author: aws-toolkit-automation

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' } }
+    )
 }