feat(amazonq): Add telemetry for auth through LSP

opieter-aws · opieter-aws · commit b6a245336e13 · 2025-05-05T09:25:07.000-04:00
diff --git a/packages/amazonq/src/extensionNode.ts b/packages/amazonq/src/extensionNode.ts
@@ -7,22 +7,22 @@ import * as vscode from 'vscode'
 import { activateAmazonQCommon, amazonQContextPrefix, deactivateCommon } from './extension'
 import { DefaultAmazonQAppInitContext } from 'aws-core-vscode/amazonq'
 import { activate as activateQGumby } from 'aws-core-vscode/amazonqGumby'
-import { ExtContext, globals, CrashMonitoring /* getLogger, isSageMaker,*/, Experiments } from 'aws-core-vscode/shared'
+import { ExtContext, globals, CrashMonitoring, getLogger, isSageMaker, Experiments } from 'aws-core-vscode/shared'
 import { filetypes, SchemaService } from 'aws-core-vscode/sharedNode'
 import { updateDevMode } from 'aws-core-vscode/dev'
 import { CommonAuthViewProvider } from 'aws-core-vscode/login'
 import { isExtensionActive, VSCODE_EXTENSION_ID } from 'aws-core-vscode/utils'
 import { registerSubmitFeedback } from 'aws-core-vscode/feedback'
 import { DevOptions } from 'aws-core-vscode/dev'
-import { Auth /* , AuthUtils, getTelemetryMetadataForConn, isAnySsoConnection*/ } from 'aws-core-vscode/auth'
+import { Auth, authUtils } from 'aws-core-vscode/auth'
 import api from './api'
 import { activate as activateCWChat } from './app/chat/activation'
 import { activate as activateInlineChat } from './inlineChat/activation'
 import { beta } from 'aws-core-vscode/dev'
 import * as amazonq from 'aws-core-vscode/amazonq'
-import { /* activate as activateNotifications,*/ NotificationsController } from 'aws-core-vscode/notifications'
-// import { AuthState, AuthUtil } from 'aws-core-vscode/codewhisperer'
-// import { telemetry, AuthUserState } from 'aws-core-vscode/telemetry'
+import { activate as activateNotifications, NotificationsController } from 'aws-core-vscode/notifications'
+import { AuthUtil } from 'aws-core-vscode/codewhisperer'
+import { telemetry, AuthUserState, AuthStatus } from 'aws-core-vscode/telemetry'
 
 export async function activate(context: vscode.ExtensionContext) {
     // IMPORTANT: No other code should be added to this function. Place it in one of the following 2 functions where appropriate.
@@ -73,34 +73,32 @@ async function activateAmazonQNode(context: vscode.ExtensionContext) {
     await setupDevMode(context)
     await beta.activate(context)
 
-    // TODO: @opieter Fix telemetry
-    // Will the web metric look the same?
-    // telemetry.auth_userState.emit({
-    //     passive: true,
-    //     result: 'Succeeded',
-    //     source: AuthUtils.ExtensionUse.instance.sourceForTelemetry(),
-    //     ...(await getAuthState()),
-    // })
+    await getAuthState()
+    telemetry.auth_userState.emit({
+        passive: true,
+        result: 'Succeeded',
+        source: authUtils.ExtensionUse.instance.sourceForTelemetry(),
+        authStatus: AuthUtil.instance.getAuthState(),
+        authEnabledConnections: (await AuthUtil.instance.getAuthFormIds()).join(','),
+    })
 
-    // void activateNotifications(context, getAuthState)
+    void activateNotifications(context, getAuthState)
 }
 
-// async function getAuthState(): Promise<Omit<auth2.AuthState, 'source'>> {
-//     let authState: AuthState = 'disconnected'
-//     authState = AuthUtil.instance.getAuthState()
-
-//     if (AuthUtil.instance.isConnected() && !(AuthUtil.instance.isSsoSession() || isSageMaker())) {
-//         getLogger().error('Current Amazon Q connection is not SSO')
-//     }
-
-//     return {
-//         authStatus:
-//             authState === 'connected' || authState === 'expired'
-//                 ? authState
-//                 : 'notConnected',
-//         ...(await getTelemetryMetadataForConn(currConn)),
-//     }
-// }
+async function getAuthState(): Promise<Omit<AuthUserState, 'source'>> {
+    const state = AuthUtil.instance.getAuthState()
+
+    if (AuthUtil.instance.isConnected() && !(AuthUtil.instance.isSsoSession() || isSageMaker())) {
+        getLogger().error('Current Amazon Q connection is not SSO')
+    }
+
+    return {
+        // @ts-ignore
+        authStatus: (state ?? 'notConnected') as AuthStatus,
+        authEnabledConnections: (await AuthUtil.instance.getAuthFormIds()).join(','),
+        ...AuthUtil.instance.getTelemetryMetadata(),
+    }
+}
 
 /**
  * Some parts of this do not work in Web mode so we need to set Dev Mode up here.
diff --git a/packages/core/src/codewhisperer/util/authUtil.ts b/packages/core/src/codewhisperer/util/authUtil.ts
@@ -8,7 +8,15 @@ import * as localizedText from '../../shared/localizedText'
 import * as nls from 'vscode-nls'
 import { ToolkitError } from '../../shared/errors'
 import { AmazonQPromptSettings } from '../../shared/settings'
-import { scopesCodeWhispererCore, scopesCodeWhispererChat, scopesFeatureDev, scopesGumby } from '../../auth/connection'
+import {
+    scopesCodeWhispererCore,
+    scopesCodeWhispererChat,
+    scopesFeatureDev,
+    scopesGumby,
+    TelemetryMetadata,
+    scopesSsoAccountAccess,
+    hasScopes,
+} from '../../auth/connection'
 import { getLogger } from '../../shared/logger/logger'
 import { Commands } from '../../shared/vscode/commands2'
 import { vsCodeState } from '../models/model'
@@ -21,6 +29,7 @@ import { AuthStateEvent, LanguageClientAuth, LoginTypes, SsoLogin } from '../../
 import { builderIdStartUrl } from '../../auth/sso/constants'
 import { VSCODE_EXTENSION_ID } from '../../shared/extensions'
 import { RegionProfileManager } from '../region/regionProfileManager'
+import { AuthFormId } from '../../login/webview/vue/types'
 
 const localize = nls.loadMessageBundle()
 
@@ -260,4 +269,55 @@ export class AuthUtil implements IAuthProvider {
             void vscode.commands.executeCommand('aws.amazonq.notifyNewCustomizations')
         }
     }
+
+    async getTelemetryMetadata(): Promise<TelemetryMetadata> {
+        if (!this.isConnected()) {
+            return {
+                id: 'undefined',
+            }
+        }
+
+        if (this.isSsoSession()) {
+            const ssoSessionDetails = (await this.session.getProfile()).ssoSession?.settings
+            return {
+                authScopes: ssoSessionDetails?.sso_registration_scopes?.join(','),
+                credentialSourceId: AuthUtil.instance.isBuilderIdConnection() ? 'awsId' : 'iamIdentityCenter',
+                credentialStartUrl: AuthUtil.instance.connection?.startUrl,
+                awsRegion: AuthUtil.instance.connection?.region,
+            }
+        } else if (!AuthUtil.instance.isSsoSession) {
+            return {
+                credentialSourceId: 'sharedCredentials',
+            }
+        }
+
+        throw new Error('getTelemetryMetadataForConn() called with unknown connection type')
+    }
+
+    async getAuthFormIds(): Promise<AuthFormId[]> {
+        if (!this.isConnected()) {
+            return []
+        }
+
+        const authIds: AuthFormId[] = []
+        let connType: 'builderId' | 'identityCenter'
+
+        // TODO: update when there is IAM support
+        if (!this.isSsoSession()) {
+            return ['credentials']
+        } else if (this.isBuilderIdConnection()) {
+            connType = 'builderId'
+        } else if (this.isIdcConnection()) {
+            connType = 'identityCenter'
+            const ssoSessionDetails = (await this.session.getProfile()).ssoSession?.settings
+            if (hasScopes(ssoSessionDetails?.sso_registration_scopes ?? [], scopesSsoAccountAccess)) {
+                authIds.push('identityCenterExplorer')
+            }
+        } else {
+            return ['unknown']
+        }
+        authIds.push(`${connType}CodeWhisperer`)
+
+        return authIds
+    }
 }
diff --git a/packages/core/src/login/webview/commonAuthViewProvider.ts b/packages/core/src/login/webview/commonAuthViewProvider.ts
@@ -45,6 +45,7 @@ import { telemetry } from '../../shared/telemetry/telemetry'
 import { AuthSources } from './util'
 import { AuthFlowStates } from './vue/types'
 import { ExtensionUse } from '../../auth/utils'
+import { AuthUtil } from '../../codewhisperer/util/authUtil'
 
 export class CommonAuthViewProvider implements WebviewViewProvider {
     public readonly viewType: string
@@ -107,7 +108,7 @@ export class CommonAuthViewProvider implements WebviewViewProvider {
                 if (authState === AuthFlowStates.REAUTHNEEDED || authState === AuthFlowStates.REAUTHENTICATING) {
                     this.webView!.server.storeMetricMetadata({
                         isReAuth: true,
-                        // ...(await getTelemetryMetadataForConn(AuthUtil.instance.conn)), // TODO: @opieter Re-add telemetry
+                        ...(await AuthUtil.instance.getTelemetryMetadata()),
                     })
                 } else {
                     this.webView!.server.storeMetricMetadata({ isReAuth: false })
diff --git a/packages/core/src/login/webview/vue/amazonq/backend_amazonq.ts b/packages/core/src/login/webview/vue/amazonq/backend_amazonq.ts
@@ -65,8 +65,7 @@ export class AmazonQLoginWebview extends CommonAuthWebview {
                 isReAuth: false,
             })
             await awsIdSignIn()
-            // this.storeMetricMetadata(await getTelemetryMetadataForConn(conn))
-            // TODO: @opieter re-add metrics
+            this.storeMetricMetadata(await AuthUtil.instance.getTelemetryMetadata())
 
             void vscode.window.showInformationMessage('AmazonQ: Successfully connected to AWS Builder ID')
         })
@@ -83,8 +82,7 @@ export class AmazonQLoginWebview extends CommonAuthWebview {
             })
 
             await connectToEnterpriseSso(startUrl, region)
-            // this.storeMetricMetadata(await getTelemetryMetadataForConn(conn))
-            // TODO: @opieter re-add metrics
+            this.storeMetricMetadata(await AuthUtil.instance.getTelemetryMetadata())
 
             void vscode.window.showInformationMessage('AmazonQ: Successfully connected to AWS IAM Identity Center')
         })
@@ -110,16 +108,15 @@ export class AmazonQLoginWebview extends CommonAuthWebview {
              * causes the reauth page to refresh before the user is actually done the whole reauth flow.
              */
             this.reauthError = await this.ssoSetup('reauthenticateAmazonQ', async () => {
-                // TODO: @ opieter Re-add metrics
-                // this.storeMetricMetadata({
-                //     authEnabledFeatures: this.getAuthEnabledFeatures(conn),
-                //     isReAuth: true,
-                //     ...(await getTelemetryMetadataForConn(conn)),
-                // })
+                this.storeMetricMetadata({
+                    authEnabledFeatures: 'codewhisperer',
+                    isReAuth: true,
+                    ...(await AuthUtil.instance.getTelemetryMetadata()),
+                })
                 await AuthUtil.instance.reauthenticate()
-                // this.storeMetricMetadata({
-                //     ...(await getTelemetryMetadataForConn(conn)),
-                // })
+                this.storeMetricMetadata({
+                    ...(await AuthUtil.instance.getTelemetryMetadata()),
+                })
             })
         } finally {
             this.isReauthenticating = false
