feat(telemetry): add telemetry for restoring and forgetting connections (#5454)

Includes a new telemetry event, 'auth_modifyConnection' which will begin to encompass changes to connections.

Author: hayemaxi

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() // method under test
+        await authUtil.clearExtraConnections('test') // method under test
 
         // Only the conn that AuthUtil is using is remaining
         assert.deepStrictEqual(

packages/core/src/auth/secondaryAuth.ts

@@ -12,8 +12,9 @@ import { Auth } from './auth'
 import { once, onceChanged } from '../shared/utilities/functionUtils'
 import { isNonNullable } from '../shared/utilities/tsUtils'
 import { ToolIdStateKey } from '../shared/globalState'
-import { Connection, SsoConnection, StatefulConnection } from './connection'
+import { Connection, getTelemetryMetadataForConn, SsoConnection, StatefulConnection } from './connection'
 import { indent } from '../shared/utilities/textUtilities'
+import { telemetry } from '../shared/telemetry/telemetry'
 
 export type ToolId = 'codecatalyst' | 'codewhisperer' | 'testId'
 
@@ -244,13 +245,28 @@ 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: () => Promise<T | undefined> = once(async () => {
+    public readonly restoreConnection: (source?: string) => Promise<T | undefined> = once(async (source?: string) => {
         try {
-            await this.auth.tryAutoConnect()
-            this.#savedConnection = await this.loadSavedConnection()
-            this.#onDidChangeActiveConnection.fire(this.activeConnection)
+            return await telemetry.auth_modifyConnection.run(async () => {
+                telemetry.record({
+                    source,
+                    action: 'restore',
+                    connectionState: 'undefined',
+                })
+                await this.auth.tryAutoConnect()
+                this.#savedConnection = await this.loadSavedConnection()
+                this.#onDidChangeActiveConnection.fire(this.activeConnection)
 
-            return this.#savedConnection
+                const conn = this.#activeConnection
+                if (conn) {
+                    telemetry.record({
+                        connectionState: this.auth.getConnectionState(conn),
+                        ...(await getTelemetryMetadataForConn(conn)),
+                    })
+                }
+
+                return this.#savedConnection
+            })
         } catch (err) {
             getLogger().warn(`auth (${this.toolId}): failed to restore connection: %s`, err)
         }

packages/core/src/codecatalyst/auth.ts

@@ -151,8 +151,8 @@ export class CodeCatalystAuthenticationProvider {
         }
     }
 
-    public async restore() {
-        await this.secondaryAuth.restoreConnection()
+    public async restore(source?: string) {
+        await this.secondaryAuth.restoreConnection(source)
     }
 
     private async accessDeniedExceptionHandler(showReauthPrompt: boolean = true) {

packages/core/src/codewhisperer/activation.ts

@@ -294,8 +294,8 @@ export async function activate(context: ExtContext): Promise<void> {
         vscode.commands.registerCommand('aws.amazonq.openEditorAtRange', openEditorAtRange)
     )
 
-    await auth.restore()
-    await auth.clearExtraConnections()
+    await auth.restore('startup')
+    await auth.clearExtraConnections('startup')
 
     if (auth.isConnectionExpired()) {
         auth.showReauthenticatePrompt().catch((e) => {

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

@@ -25,6 +25,7 @@ import {
     scopesGumby,
     isIdcSsoConnection,
     hasExactScopes,
+    getTelemetryMetadataForConn,
 } from '../../auth/connection'
 import { getLogger } from '../../shared/logger'
 import { Commands } from '../../shared/vscode/commands2'
@@ -35,6 +36,7 @@ import { showReauthenticateMessage } from '../../shared/utilities/messages'
 import { showAmazonQWalkthroughOnce } from '../../amazonq/onboardingPage/walkthrough'
 import { setContext } from '../../shared/vscode/setContext'
 import { isInDevEnv } from '../../shared/vscode/env'
+import { telemetry } from '../../shared/telemetry/telemetry'
 
 /** Backwards compatibility for connections w pre-chat scopes */
 export const codeWhispererCoreScopes = [...scopesCodeWhispererCore]
@@ -99,7 +101,7 @@ export class AuthUtil {
         'Amazon Q',
         isValidCodeWhispererCoreConnection
     )
-    public readonly restore = () => this.secondaryAuth.restoreConnection()
+    public readonly restore = (source?: string) => this.secondaryAuth.restoreConnection(source)
 
     public constructor(public readonly auth = Auth.instance) {}
 
@@ -402,15 +404,29 @@ export class AuthUtil {
      * auth connections that the Amazon Q extension has cached. We need to remove these
      * as they are irrelevant to the Q extension and can cause issues.
      */
-    public async clearExtraConnections(): Promise<void> {
+    public async clearExtraConnections(source: string): Promise<void> {
         const currentQConn = this.conn
         // Q currently only maintains 1 connection at a time, so we assume everything else is extra.
         // IMPORTANT: In the case Q starts to manage multiple connections, this implementation will need to be updated.
         const allOtherConnections = (await this.auth.listConnections()).filter((c) => c.id !== currentQConn?.id)
         for (const conn of allOtherConnections) {
             getLogger().warn(`forgetting extra amazon q connection: %O`, conn)
-            // in a Dev Env the connection may be used by code catalyst, so we forget instead of fully deleting
-            isInDevEnv() ? await this.auth.forgetConnection(conn) : await this.auth.deleteConnection(conn)
+            await telemetry.auth_modifyConnection.run(async () => {
+                telemetry.record({
+                    connectionState: Auth.instance.getConnectionState(conn) ?? 'undefined',
+                    source,
+                    ...(await getTelemetryMetadataForConn(conn)),
+                })
+
+                if (isInDevEnv()) {
+                    telemetry.record({ action: 'forget' })
+                    // in a Dev Env the connection may be used by code catalyst, so we forget instead of fully deleting
+                    await this.auth.forgetConnection(conn)
+                } else {
+                    telemetry.record({ action: 'delete' })
+                    await this.auth.deleteConnection(conn)
+                }
+            })
         }
     }
 }

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

@@ -270,6 +270,11 @@
             "name": "awsRegion",
             "type": "string",
             "description": "An AWS region."
+        },
+        {
+            "name": "connectionState",
+            "type": "string",
+            "description": "A detailed state of a specific auth connection. Use `authStatus` for a higher level view of an extension's general connection."
         }
     ],
     "metrics": [
@@ -1039,6 +1044,40 @@
                 }
             ],
             "passive": true
+        },
+        {
+            "name": "auth_modifyConnection",
+            "description": "An auth connection was modified in some way, e.g. deleted, updated",
+            "metadata": [
+                {
+                    "type": "action",
+                    "required": true
+                },
+                {
+                    "type": "authScopes",
+                    "required": false
+                },
+                {
+                    "type": "connectionState",
+                    "required": true
+                },
+                {
+                    "type": "credentialStartUrl",
+                    "required": false
+                },
+                {
+                    "type": "source",
+                    "required": true
+                },
+                {
+                    "type": "ssoRegistrationClientId",
+                    "required": false
+                },
+                {
+                    "type": "ssoRegistrationExpiresAt",
+                    "required": false
+                }
+            ]
         }
     ]
 }