telemetry(auth): snooze techdebt and add telemetry for evaluation (aws#5516)

hayemaxi · web-flow · commit 01d166bcfc8d · 2024-08-29T16:29:56.000-04:00
We are still getting hits in telemetry for users seeing the split
session notification. Adding more fine-grained telemetry to determined
if this users are still migrating. This telemetry will also help us
track other auth issues.

---

&lt;!--- REMINDER: Ensure that your PR meets the guidelines in
CONTRIBUTING.md --&gt;

License: I confirm that my contribution is made under the terms of the
Apache 2.0 license.
diff --git a/packages/core/src/auth/auth.ts b/packages/core/src/auth/auth.ts
@@ -579,6 +579,7 @@ export class Auth implements AuthService, ConnectionManager {
                     : undefined
             span.record({
                 action: 'updateConnectionState',
+                id,
                 connectionState,
                 source: asStringifiedStack(telemetry.getFunctionStack()),
                 credentialStartUrl: oldProfile.type === 'sso' ? oldProfile.startUrl : undefined,
diff --git a/packages/core/src/auth/connection.ts b/packages/core/src/auth/connection.ts
@@ -13,6 +13,10 @@ import { getLogger } from '../shared/logger/logger'
 import { showMessageWithUrl } from '../shared/utilities/messages'
 import { onceChanged } from '../shared/utilities/functionUtils'
 import { AuthAddConnection, AwsLoginWithBrowser } from '../shared/telemetry/telemetry.gen'
+import { withTelemetryContext } from '../shared/telemetry/util'
+import { AuthModifyConnection, telemetry } from '../shared/telemetry/telemetry'
+import { asStringifiedStack } from '../shared/telemetry/spans'
+import { getTelemetryReason, getTelemetryReasonDesc } from '../shared/errors'
 
 /** Shows a warning message unless it is the same as the last one shown. */
 const warnOnce = onceChanged((s: string, url: string) => {
@@ -197,17 +201,74 @@ export interface ProfileMetadata {
 
 export type StoredProfile<T extends Profile = Profile> = T & { readonly metadata: ProfileMetadata }
 
+function getTelemetryForProfile(profile: StoredProfile<Profile> | undefined) {
+    if (!profile) {
+        return {}
+    }
+
+    let metadata: Partial<AuthModifyConnection> = {
+        connectionState: profile?.metadata.connectionState ?? 'undefined',
+    }
+
+    if (profile.type === 'sso') {
+        metadata = {
+            ...metadata,
+            authScopes: profile.scopes?.join(','),
+            credentialStartUrl: profile.startUrl,
+            awsRegion: profile.ssoRegion,
+        }
+    }
+
+    return metadata
+}
+
+const profileStoreClassName = 'ProfileStore'
 export class ProfileStore {
+    // To de-dupe telemetry
+    private _prevGetProfile: { id: string; connectionState: ProfileMetadata['connectionState'] } | undefined
+
     public constructor(private readonly memento: vscode.Memento) {}
 
     public getProfile(id: string): StoredProfile | undefined {
         return this.getData()[id]
     }
 
+    @withTelemetryContext({ name: 'getProfileOrThrow', class: profileStoreClassName })
     public getProfileOrThrow(id: string): StoredProfile {
-        const profile = this.getProfile(id)
-        if (profile === undefined) {
-            throw new Error(`Profile does not exist: ${id}`)
+        const metadata: AuthModifyConnection = {
+            action: 'getProfile',
+            id,
+            source: asStringifiedStack(telemetry.getFunctionStack()),
+        }
+
+        let profile: StoredProfile<Profile> | undefined
+        try {
+            profile = this.getProfile(id)
+            if (profile === undefined) {
+                throw new Error(`Profile does not exist: ${id}`)
+            }
+        } catch (err) {
+            // Always emit failures
+            telemetry.auth_modifyConnection.emit({
+                ...metadata,
+                result: 'Failed',
+                reason: getTelemetryReason(err),
+                reasonDesc: getTelemetryReasonDesc(err),
+            })
+            throw err
+        }
+
+        // De-dupe metric on last id and connection state
+        if (
+            this._prevGetProfile?.id !== id ||
+            this._prevGetProfile?.connectionState !== profile.metadata.connectionState
+        ) {
+            telemetry.auth_modifyConnection.emit({
+                ...metadata,
+                ...getTelemetryForProfile(profile),
+                result: 'Succeeded',
+            })
+            this._prevGetProfile = { id, connectionState: profile.metadata.connectionState }
         }
 
         return profile
@@ -219,17 +280,40 @@ export class ProfileStore {
 
     public async addProfile(id: string, profile: SsoProfile): Promise<StoredProfile<SsoProfile>>
     public async addProfile(id: string, profile: IamProfile): Promise<StoredProfile<IamProfile>>
+    @withTelemetryContext({ name: 'addProfile', class: profileStoreClassName })
     public async addProfile(id: string, profile: Profile): Promise<StoredProfile> {
-        return this.putProfile(id, this.initMetadata(profile))
+        return telemetry.auth_modifyConnection.run(async (span) => {
+            span.record({
+                action: 'addProfile',
+                id,
+                source: asStringifiedStack(telemetry.getFunctionStack()),
+            })
+
+            const newProfile = this.initMetadata(profile)
+            span.record(getTelemetryForProfile(newProfile))
+
+            return await this.putProfile(id, newProfile)
+        })
     }
 
+    @withTelemetryContext({ name: 'updateProfile', class: profileStoreClassName })
     public async updateProfile(id: string, profile: Profile): Promise<StoredProfile> {
-        const oldProfile = this.getProfileOrThrow(id)
-        if (oldProfile.type !== profile.type) {
-            throw new Error(`Cannot change profile type from "${oldProfile.type}" to "${profile.type}"`)
-        }
+        return telemetry.auth_modifyConnection.run(async (span) => {
+            span.record({
+                action: 'updateProfile',
+                id,
+                source: asStringifiedStack(telemetry.getFunctionStack()),
+            })
+
+            const oldProfile = this.getProfileOrThrow(id)
+            if (oldProfile.type !== profile.type) {
+                throw new Error(`Cannot change profile type from "${oldProfile.type}" to "${profile.type}"`)
+            }
 
-        return this.putProfile(id, { ...oldProfile, ...profile })
+            const newProfile = await this.putProfile(id, { ...oldProfile, ...profile })
+            span.record(getTelemetryForProfile(newProfile))
+            return newProfile
+        })
     }
 
     public async updateMetadata(id: string, metadata: ProfileMetadata): Promise<StoredProfile> {
@@ -238,11 +322,21 @@ export class ProfileStore {
         return this.putProfile(id, { ...profile, metadata: { ...profile.metadata, ...metadata } })
     }
 
+    @withTelemetryContext({ name: 'deleteProfile', class: profileStoreClassName })
     public async deleteProfile(id: string): Promise<void> {
-        const data = this.getData()
-        delete (data as Mutable<typeof data>)[id]
-
-        await this.updateData(data)
+        return telemetry.auth_modifyConnection.run(async (span) => {
+            span.record({
+                action: 'deleteProfile',
+                id,
+                source: asStringifiedStack(telemetry.getFunctionStack()),
+            })
+
+            const data = this.getData()
+            span.record(getTelemetryForProfile(data[id]))
+            delete (data as Mutable<typeof data>)[id]
+
+            await this.updateData(data)
+        })
     }
 
     public getCurrentProfileId(): string | undefined {
@@ -395,11 +489,13 @@ export interface AwsConnection {
 }
 
 type Writeable<T> = { -readonly [U in keyof T]: T[U] }
-export type TelemetryMetadata = Partial<Writeable<AuthAddConnection | AwsLoginWithBrowser>>
+export type TelemetryMetadata = Partial<Writeable<AuthAddConnection & AwsLoginWithBrowser & AuthModifyConnection>>
 
 export async function getTelemetryMetadataForConn(conn?: Connection): Promise<TelemetryMetadata> {
     if (conn === undefined) {
-        return {}
+        return {
+            id: 'undefined',
+        }
     }
 
     if (isSsoConnection(conn)) {
diff --git a/packages/core/src/auth/secondaryAuth.ts b/packages/core/src/auth/secondaryAuth.ts
@@ -256,17 +256,25 @@ export class SecondaryAuth<T extends Connection = Connection> {
     @withTelemetryContext({ name: 'restoreConnection', class: secondaryAuthClassName })
     private async _restoreConnection(): Promise<T | undefined> {
         try {
-            return await telemetry.auth_modifyConnection.run(async () => {
-                telemetry.record({
+            return await telemetry.auth_modifyConnection.run(async (span) => {
+                span.record({
                     source: asStringifiedStack(telemetry.getFunctionStack()),
                     action: 'restore',
+                    id: 'undefined',
                     connectionState: 'undefined',
                 })
                 await this.auth.tryAutoConnect()
                 this.#savedConnection = await this._loadSavedConnection()
                 this.#onDidChangeActiveConnection.fire(this.activeConnection)
 
-                telemetry.record(await getTelemetryMetadataForConn(this.#savedConnection))
+                const conn = this.#savedConnection
+                if (conn) {
+                    span.record({
+                        connectionState: this.auth.getConnectionState(conn),
+                        ...(await getTelemetryMetadataForConn(conn)),
+                    })
+                }
+
                 return this.#savedConnection
             })
         } catch (err) {
diff --git a/packages/core/src/codecatalyst/activation.ts b/packages/core/src/codecatalyst/activation.ts
@@ -25,8 +25,10 @@ import { getLogger } from '../shared/logger/logger'
 import { DevEnvActivityStarter } from './devEnv'
 import { learnMoreCommand, onboardCommand, reauth } from './explorer'
 import { isInDevEnv } from '../shared/vscode/env'
-import { hasScopes, scopesCodeWhispererCore } from '../auth/connection'
+import { hasScopes, scopesCodeWhispererCore, getTelemetryMetadataForConn } from '../auth/connection'
 import { SessionSeparationPrompt } from '../auth/auth'
+import { telemetry } from '../shared/telemetry/telemetry'
+import { asStringifiedStack } from '../shared/telemetry/spans'
 
 const localize = nls.loadMessageBundle()
 
@@ -50,8 +52,23 @@ export async function activate(ctx: ExtContext): Promise<void> {
     // Note: credentials on disk in the dev env cannot have Q scopes, so it will never be forgotten.
     // TODO: Remove after some time?
     if (authProvider.isConnected() && hasScopes(authProvider.activeConnection!, scopesCodeWhispererCore)) {
-        await authProvider.secondaryAuth.forgetConnection()
-        await SessionSeparationPrompt.instance.showForCommand('aws.codecatalyst.manageConnections')
+        await telemetry.function_call.run(
+            async () => {
+                await telemetry.auth_modifyConnection.run(async () => {
+                    const conn = authProvider.activeConnection
+                    telemetry.record({
+                        action: 'forget',
+                        source: asStringifiedStack(telemetry.getFunctionStack()),
+                        connectionState: conn ? authProvider.auth.getConnectionState(conn) : undefined,
+                        ...(await getTelemetryMetadataForConn(conn)),
+                    })
+
+                    await authProvider.secondaryAuth.forgetConnection()
+                    await SessionSeparationPrompt.instance.showForCommand('aws.codecatalyst.manageConnections')
+                })
+            },
+            { emit: false, functionId: { name: 'activate', class: 'CodeCatalyst' } }
+        )
     }
 
     ctx.extensionContext.subscriptions.push(
diff --git a/packages/core/src/shared/telemetry/vscodeTelemetry.json b/packages/core/src/shared/telemetry/vscodeTelemetry.json
@@ -1094,7 +1094,7 @@
                 },
                 {
                     "type": "connectionState",
-                    "required": true
+                    "required": false
                 },
                 {
                     "type": "credentialStartUrl",
@@ -1115,6 +1115,10 @@
                 {
                     "type": "ssoRegistrationExpiresAt",
                     "required": false
+                },
+                {
+                    "type": "id",
+                    "required": false
                 }
             ],
             "passive": true
diff --git a/packages/core/src/test/credentials/auth.test.ts b/packages/core/src/test/credentials/auth.test.ts
@@ -55,12 +55,27 @@ describe('Auth', function () {
         await auth.deleteConnection({ id: conn.id })
         assert.strictEqual((await auth.listConnections()).length, 0)
         assertTelemetry('auth_modifyConnection', [
+            {
+                action: 'addProfile',
+                connectionState: 'unauthenticated',
+                source: 'Auth#createConnection:ProfileStore#addProfile',
+            },
+            {
+                action: 'getProfile',
+                connectionState: 'unauthenticated',
+                source: 'Auth#createConnection,updateConnectionState:ProfileStore#getProfileOrThrow',
+            },
             {
                 action: 'updateConnectionState',
                 connectionState: 'valid',
                 source: 'Auth#createConnection,updateConnectionState',
                 sessionDuration: undefined,
             },
+            {
+                action: 'getProfile',
+                connectionState: 'valid',
+                source: 'Auth#deleteConnection,invalidateConnection:ProfileStore#getProfileOrThrow',
+            },
             {
                 action: 'updateConnectionState',
                 connectionState: 'invalid',
@@ -78,12 +93,27 @@ describe('Auth', function () {
         assert.strictEqual((await auth.listConnections()).length, 0)
         assert.strictEqual(auth.activeConnection, undefined)
         assertTelemetry('auth_modifyConnection', [
+            {
+                action: 'addProfile',
+                connectionState: 'unauthenticated',
+                source: 'Auth#createConnection:ProfileStore#addProfile',
+            },
+            {
+                action: 'getProfile',
+                connectionState: 'unauthenticated',
+                source: 'Auth#createConnection,updateConnectionState:ProfileStore#getProfileOrThrow',
+            },
             {
                 action: 'updateConnectionState',
                 connectionState: 'valid',
                 source: 'Auth#createConnection,updateConnectionState',
                 sessionDuration: undefined,
             },
+            {
+                action: 'getProfile',
+                connectionState: 'valid',
+                source: 'Auth#useConnection,refreshConnectionState,validateConnection,updateConnectionState:ProfileStore#getProfileOrThrow',
+            },
             {
                 action: 'updateConnectionState',
                 connectionState: 'invalid',
@@ -113,12 +143,27 @@ describe('Auth', function () {
         assert.strictEqual(auth.activeConnection, undefined)
         assert.strictEqual(auth.activeConnectionEvents.last, undefined)
         assertTelemetry('auth_modifyConnection', [
+            {
+                action: 'addProfile',
+                connectionState: 'unauthenticated',
+                source: 'Auth#createConnection:ProfileStore#addProfile',
+            },
+            {
+                action: 'getProfile',
+                connectionState: 'unauthenticated',
+                source: 'Auth#createConnection,updateConnectionState:ProfileStore#getProfileOrThrow',
+            },
             {
                 action: 'updateConnectionState',
                 connectionState: 'valid',
                 source: 'Auth#createConnection,updateConnectionState',
                 sessionDuration: undefined,
             },
+            {
+                action: 'getProfile',
+                connectionState: 'valid',
+                source: 'Auth#useConnection,refreshConnectionState,validateConnection,updateConnectionState:ProfileStore#getProfileOrThrow',
+            },
             {
                 action: 'updateConnectionState',
                 connectionState: 'invalid',
diff --git a/packages/core/src/test/techdebt.test.ts b/packages/core/src/test/techdebt.test.ts
@@ -49,7 +49,11 @@ describe('tech debt', function () {
     it('remove separate sessions login edge cases', async function () {
         // src/auth/auth.ts:SessionSeparationPrompt
         // forgetConnection() function and calls
-        fixByDate('2024-08-30', 'Remove the edge case code from the commit that this test is a part of.')
+
+        // Monitor telemtry to determine removal or snooze
+        // toolkit_showNotification.id = sessionSeparation
+        // auth_modifyConnection.action = deleteProfile OR auth_modifyConnection.source contains CodeCatalyst
+        fixByDate('2024-9-30', 'Remove the edge case code from the commit that this test is a part of.')
     })
 
     it('remove deprecated code transform metrics', async function () {

Original file line number	Diff line number	Diff line change
`@@ -1094,7 +1094,7 @@`
`1094`	`1094`	`},`
`1095`	`1095`	`{`
`1096`	`1096`	`"type": "connectionState",`
`1097`		`- "required": true`
	`1097`	`+ "required": false`
`1098`	`1098`	`},`
`1099`	`1099`	`{`
`1100`	`1100`	`"type": "credentialStartUrl",`
`@@ -1115,6 +1115,10 @@`
`1115`	`1115`	`{`
`1116`	`1116`	`"type": "ssoRegistrationExpiresAt",`
`1117`	`1117`	`"required": false`
	`1118`	`+ },`
	`1119`	`+ {`
	`1120`	`+ "type": "id",`
	`1121`	`+ "required": false`
`1118`	`1122`	`}`
`1119`	`1123`	`],`
`1120`	`1124`	`"passive": true`