fix(cwspr): add more log info for auth code path (#4526)

* add more log info for auth code path

* more logs

* reformat logs

* log prefix

* add log at event listener

* Update packages/core/src/auth/secondaryAuth.ts

Co-authored-by: Justin M. Keyes <[email protected]>

* Update packages/core/src/auth/sso/ssoAccessTokenProvider.ts

Co-authored-by: Justin M. Keyes <[email protected]>

* add one more log

---------

Co-authored-by: Justin M. Keyes <[email protected]>

Author: leigaol

packages/core/src/auth/auth.ts

@@ -188,7 +188,7 @@ export class Auth implements AuthService, ConnectionManager {
         if (profile === undefined) {
             throw new Error(`Connection does not exist: ${id}`)
         }
-
+        getLogger().info(`auth: validating connection before using it: ${id}`)
         const validated = await this.validateConnection(id, profile)
         const conn =
             validated.type === 'sso' ? this.getSsoConnection(id, validated) : this.getIamConnection(id, validated)
@@ -204,7 +204,7 @@ export class Auth implements AuthService, ConnectionManager {
         if (this.activeConnection === undefined) {
             return
         }
-
+        getLogger().info(`auth: logout`)
         await this.store.setCurrentProfileId(undefined)
         await this.invalidateConnection(this.activeConnection.id)
         this.#activeConnection = undefined
@@ -339,6 +339,7 @@ export class Auth implements AuthService, ConnectionManager {
      * Put the SSO connection in to an expired state
      */
     public async expireConnection(conn: Pick<SsoConnection, 'id'>): Promise<void> {
+        getLogger().info(`auth: Expiring connection ${conn.id}`)
         const profile = this.store.getProfileOrThrow(conn.id)
         if (profile.type === 'iam') {
             throw new ToolkitError('Auth: Cannot force expire an IAM connection')
@@ -374,6 +375,7 @@ export class Auth implements AuthService, ConnectionManager {
 
     public async updateConnection(connection: Pick<SsoConnection, 'id'>, profile: SsoProfile): Promise<SsoConnection>
     public async updateConnection(connection: Pick<Connection, 'id'>, profile: Profile): Promise<Connection> {
+        getLogger().info(`auth: Updating connection ${connection.id}`)
         if (profile.type === 'iam') {
             throw new Error('Updating IAM connections is not supported')
         }
@@ -454,6 +456,7 @@ export class Auth implements AuthService, ConnectionManager {
      * before the local token(s) are cleared as they are needed in the request.
      */
     private async invalidateConnection(id: Connection['id'], opt?: { skipGlobalLogout?: boolean }) {
+        getLogger().info(`auth: Invalidating connection: ${id}`)
         const profile = this.store.getProfileOrThrow(id)
 
         if (profile.type === 'sso') {
@@ -479,6 +482,7 @@ export class Auth implements AuthService, ConnectionManager {
     }
 
     private async updateConnectionState(id: Connection['id'], connectionState: ProfileMetadata['connectionState']) {
+        getLogger().info(`auth: Updating connection state of ${id} to ${connectionState}`)
         const oldProfile = this.store.getProfileOrThrow(id)
         if (oldProfile.metadata.connectionState === connectionState) {
             return oldProfile
@@ -504,8 +508,10 @@ export class Auth implements AuthService, ConnectionManager {
             if (profile.type === 'sso') {
                 const provider = this.getTokenProvider(id, profile)
                 if ((await provider.getToken()) === undefined) {
+                    getLogger().info(`auth: Connection is not valid: ${id} `)
                     return this.updateConnectionState(id, 'invalid')
                 } else {
+                    getLogger().info(`auth: Connection is valid: ${id} `)
                     return this.updateConnectionState(id, 'valid')
                 }
             } else {
@@ -539,7 +545,7 @@ export class Auth implements AuthService, ConnectionManager {
 
     private async handleValidationError(id: Connection['id'], err: unknown) {
         this.#validationErrors.set(id, UnknownError.cast(err))
-
+        getLogger().info(`auth: Handling validation error of connection: ${id}`)
         return this.updateConnectionState(id, 'invalid')
     }
 
@@ -736,6 +742,7 @@ export class Auth implements AuthService, ConnectionManager {
     }
 
     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
         await this.updateConnectionState(id, 'invalid')

packages/core/src/auth/secondaryAuth.ts

@@ -12,6 +12,7 @@ import { Auth } from './auth'
 import { once } from '../shared/utilities/functionUtils'
 import { isNonNullable } from '../shared/utilities/tsUtils'
 import { Connection, SsoConnection, StatefulConnection } from './connection'
+import { indent } from '../shared/utilities/textUtilities'
 
 let currentConn: Auth['activeConnection']
 const auths = new Map<string, SecondaryAuth>()
@@ -156,6 +157,13 @@ export class SecondaryAuth<T extends Connection = Connection> {
     }
 
     public get isConnectionExpired() {
+        if (this.activeConnection) {
+            getLogger().info(
+                indent(`secondaryAuth connection id = ${this.activeConnection.id}
+            secondaryAuth connection status = ${this.auth.getConnectionState(this.activeConnection)}`,
+            4, true)
+            )
+        }
         return !!this.activeConnection && this.auth.getConnectionState(this.activeConnection) === 'invalid'
     }
 

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

@@ -26,6 +26,7 @@ import {
 import { getLogger } from '../../shared/logger'
 import { AwsRefreshCredentials, telemetry } from '../../shared/telemetry/telemetry'
 import { getIdeProperties, isCloud9 } from '../../shared/extensionUtilities'
+import { indent } from '../../shared/utilities/textUtilities'
 
 const clientRegistrationType = 'public'
 const deviceGrantType = 'urn:ietf:params:oauth:grant-type:device_code'
@@ -76,6 +77,7 @@ export class SsoAccessTokenProvider {
 
     public async invalidate(): Promise<void> {
         // Use allSettled() instead of all() to ensure all clear() calls are resolved.
+        getLogger().info(`SsoAccessTokenProvider invalidate token and registration`)
         await Promise.allSettled([
             this.cache.token.clear(this.tokenCacheKey, 'SsoAccessTokenProvider.invalidate()'),
             this.cache.registration.clear(this.registrationCacheKey, 'SsoAccessTokenProvider.invalidate()'),
@@ -84,7 +86,11 @@ export class SsoAccessTokenProvider {
 
     public async getToken(): Promise<SsoToken | undefined> {
         const data = await this.cache.token.load(this.tokenCacheKey)
-
+        getLogger().info(
+            indent(`current client registration id=${data?.registration?.clientId}, 
+                            expires at ${data?.registration?.expiresAt}, 
+                            key = ${this.tokenCacheKey}`, 4, true)
+        )
         if (!data || !isExpired(data.token)) {
             return data?.token
         }

packages/core/src/codewhisperer/service/featureConfigProvider.ts

@@ -54,7 +54,7 @@ export class FeatureConfigProvider {
                 )
             })
         } catch (e) {
-            getLogger().debug('CodeWhisperer: Error when fetching feature configs', e)
+            getLogger().error(`CodeWhisperer: Error when fetching feature configs ${e}`, e)
         }
         getLogger().debug(`CodeWhisperer: Current feature configs: ${this.getFeatureConfigsTelemetry()}`)
     }

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

@@ -112,6 +112,7 @@ export class AuthUtil {
 
     public constructor(public readonly auth = Auth.instance) {
         this.auth.onDidChangeConnectionState(async e => {
+            getLogger().info(`codewhisperer: connection changed to ${e.state}: ${e.id}`)
             if (e.state !== 'authenticating') {
                 await this.refreshCodeWhisperer()
             }
@@ -120,6 +121,7 @@ export class AuthUtil {
         })
 
         this.secondaryAuth.onDidChangeActiveConnection(async () => {
+            getLogger().info(`codewhisperer: active connection changed`)
             if (this.isValidEnterpriseSsoInUse()) {
                 void vscode.commands.executeCommand('aws.codeWhisperer.notifyNewCustomizations')
             }
@@ -295,11 +297,11 @@ export class AuthUtil {
             this.secondaryAuth.isConnectionExpired &&
             this.conn !== undefined &&
             isValidCodeWhispererCoreConnection(this.conn)
-        getLogger().debug(`codewhisperer: Connection expired = ${connectionExpired},
+        getLogger().info(`codewhisperer: Connection expired = ${connectionExpired},
                            secondaryAuth connection expired = ${this.secondaryAuth.isConnectionExpired},
                            connection is undefined = ${this.conn === undefined}`)
         if (this.conn) {
-            getLogger().debug(
+            getLogger().info(
                 `codewhisperer: isValidCodeWhispererConnection = ${isValidCodeWhispererCoreConnection(this.conn)}`
             )
         }

packages/core/src/shared/utilities/cacheUtils.ts

@@ -112,7 +112,7 @@ export function createDiskCache<T, K>(
                     log('read failed (file not found)', key)
                     return
                 }
-
+                log(`read failed ${error}`, key)
                 throw ToolkitError.chain(error, `Failed to read from "${target}"`, {
                     code: 'FSReadFailed',
                     details: { key },