Various telemetry adjustments (aws#4871)

* fix(telemetry): get proper authEnabledConnections for amazonq auth_userState

* refactor: circular dependency

* feat(telemetry): set source for aws_loginWithBrowser for internal connections

Only covers direct connections via command cases. We will have to follow up with additional cases later.

Author: hayemaxi

packages/amazonq/src/auth/util.ts

@@ -4,15 +4,16 @@
  */
 
 import { AuthUtil } from 'aws-core-vscode/codewhisperer'
-import { AuthStatus } from '../../../core/dist/src/shared/telemetry/telemetry'
-import { AwsConnection } from 'aws-core-vscode/auth'
+import { AuthStatus } from 'aws-core-vscode/telemetry'
+import { AwsConnection, Connection, AuthUtils } from 'aws-core-vscode/auth'
 import { activateExtension, getLogger } from 'aws-core-vscode/shared'
 import { VSCODE_EXTENSION_ID } from 'aws-core-vscode/utils'
 
 /** Provides the status of the Auth connection for Amazon Q, specifically for telemetry purposes. */
-export async function getAuthStatus(): Promise<AuthStatus> {
+export async function getAuthStatus() {
     // Get auth state from the Amazon Q extension
     const authState = (await AuthUtil.instance.getChatAuthState()).codewhispererChat
+    let authEnabledConnections = AuthUtils.getAuthFormIdsFromConnection(AuthUtil.instance.conn)
     let authStatus: AuthStatus = authState === 'connected' || authState === 'expired' ? authState : 'notConnected'
 
     // If the Q extension does not have its own connection, it will fallback and check
@@ -28,10 +29,14 @@ export async function getAuthStatus(): Promise<AuthStatus> {
         // Determine the status of the Toolkit connection we will autoconnect to
         if (autoConnectConn) {
             authStatus = autoConnectConn.state === 'valid' ? 'connected' : 'expired'
+
+            // Though TS won't say it, AwsConnection sufficiently overlaps with Connection for the purposes
+            // of `getAuthFormIdsFromConnection`
+            authEnabledConnections = AuthUtils.getAuthFormIdsFromConnection(autoConnectConn as unknown as Connection)
         }
     }
 
-    return authStatus
+    return { authStatus, authEnabledConnections: authEnabledConnections.join(',') }
 }
 
 /**

packages/amazonq/src/extensionShared.ts

@@ -156,19 +156,10 @@ export async function activateShared(context: vscode.ExtensionContext, isWeb: bo
             telemetry.record({ source: ExtStartUpSources.reload })
         }
 
-        const authKinds: AuthUtils.AuthSimpleId[] = []
-        if (await AuthUtils.hasBuilderId('codewhisperer')) {
-            authKinds.push('builderIdCodeWhisperer')
-        }
-        if (await AuthUtils.hasSso('codewhisperer')) {
-            authKinds.push('identityCenterCodeWhisperer')
-        }
-
-        const authStatus = await getAuthStatus()
-
+        const { authStatus, authEnabledConnections } = await getAuthStatus()
         telemetry.record({
             authStatus,
-            authEnabledConnections: authKinds.join(','),
+            authEnabledConnections,
         })
     })
 }

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

@@ -29,7 +29,7 @@ import {
     isNetworkError,
 } from '../../shared/errors'
 import { getLogger } from '../../shared/logger'
-import { AwsRefreshCredentials, telemetry } from '../../shared/telemetry/telemetry'
+import { AwsLoginWithBrowser, AwsRefreshCredentials, Metric, telemetry } from '../../shared/telemetry/telemetry'
 import { indent } from '../../shared/utilities/textUtilities'
 import { AuthSSOServer } from './server'
 import { CancellationError, sleep } from '../../shared/utilities/timeoutUtils'
@@ -84,6 +84,16 @@ const refreshGrantType = 'refresh_token'
  *         - RefreshToken (optional)
  */
 export abstract class SsoAccessTokenProvider {
+    /**
+     * Source to pass to aws_loginWithBrowser metric. Due to the complexity of how auth can be called,
+     * there is no other easy way to pass this in without signficant refactors.
+     */
+    private static _authSource: string = 'unknown'
+
+    public static set authSource(val: string) {
+        SsoAccessTokenProvider._authSource = val
+    }
+
     public constructor(
         protected readonly profile: Pick<SsoProfile, 'startUrl' | 'region' | 'scopes' | 'identifier'>,
         protected readonly cache = getCache(),
@@ -195,15 +205,28 @@ export abstract class SsoAccessTokenProvider {
      * Wraps the given function with telemetry related to the browser login.
      */
     protected withBrowserLoginTelemetry<T extends (...args: any[]) => any>(func: T): ReturnType<T> {
-        if (telemetry.spans.some(s => s.name === 'aws_loginWithBrowser')) {
-            // During certain flows, eg reauthentication, we are already running within a span (run())
-            // so we don't need to create a new one.
+        const run = (span: Metric<AwsLoginWithBrowser>) => {
+            span.record({
+                credentialStartUrl: this.profile.startUrl,
+                source: SsoAccessTokenProvider._authSource,
+            })
+
+            // Reset source in case there is a case where browser login was called but we forgot to set the source.
+            // We don't want to attribute the wrong source.
+            SsoAccessTokenProvider.authSource = 'unknown'
+
             return func()
         }
 
+        // During certain flows, eg reauthentication, we are already running within a span (run())
+        // so we don't need to create a new one.
+        const span = telemetry.spans.find(s => s.name === 'aws_loginWithBrowser')
+        if (span !== undefined) {
+            return run(span as unknown as Metric<AwsLoginWithBrowser>)
+        }
+
         return telemetry.aws_loginWithBrowser.run(span => {
-            span.record({ credentialStartUrl: this.profile.startUrl })
-            return func()
+            return run(span)
         })
     }
 

packages/core/src/auth/utils.ts

@@ -38,16 +38,19 @@ import {
     isIamConnection,
     isValidCodeCatalystConnection,
     createSsoProfile,
+    hasScopes,
+    scopesSsoAccountAccess,
 } from './connection'
 import { Commands, placeholder, RegisteredCommand, vscodeComponent } from '../shared/vscode/commands2'
 import { Auth } from './auth'
 import { validateIsNewSsoUrl, validateSsoUrlFormat } from './sso/validation'
 import { getLogger } from '../shared/logger'
-import { isValidCodeWhispererCoreConnection } from '../codewhisperer/util/authUtil'
+import { isValidAmazonQConnection, isValidCodeWhispererCoreConnection } from '../codewhisperer/util/authUtil'
 import { authHelpUrl } from '../shared/constants'
 import { getResourceFromTreeNode } from '../shared/treeview/utils'
 import { Instance } from '../shared/utilities/typeConstructors'
 import { openUrl } from '../shared/utilities/vsCodeUtils'
+import { AuthFormId } from './ui/vue/authForms/types'
 import { extensionVersion } from '../shared/vscode/env'
 import { ExtStartUpSources } from '../shared/telemetry'
 import { CommonAuthWebview } from '../login/webview/vue/backend'
@@ -684,19 +687,6 @@ export class ExtensionUse {
     }
 }
 
-/**
- * @deprecated
- * Remove in favor of AuthFormId
- * Simple readable ID for telemetry reporting
- */
-export type AuthSimpleId =
-    | 'sharedCredentials'
-    | 'builderIdCodeWhisperer'
-    | 'builderIdCodeCatalyst'
-    | 'identityCenterCodeWhisperer'
-    | 'identityCenterCodeCatalyst'
-    | 'identityCenterExplorer'
-
 type AuthCommands = {
     switchConnections: RegisteredCommand<(auth: Auth | TreeNode | unknown) => Promise<void>>
     help: RegisteredCommand<() => Promise<void>>
@@ -832,3 +822,37 @@ export function initAuthCommands(prefix: string) {
         }),
     }
 }
+
+/**
+ * Returns readable auth Ids for a connection, which are used by telemetry.
+ */
+export function getAuthFormIdsFromConnection(conn?: Connection): AuthFormId[] {
+    if (!conn) {
+        return []
+    }
+
+    const authIds: AuthFormId[] = []
+    let connType: 'builderId' | 'identityCenter'
+
+    if (isIamConnection(conn)) {
+        return ['credentials']
+    } else if (isBuilderIdConnection(conn)) {
+        connType = 'builderId'
+    } else if (isIdcSsoConnection(conn)) {
+        connType = 'identityCenter'
+        if (hasScopes(conn, scopesSsoAccountAccess)) {
+            authIds.push('identityCenterExplorer')
+        }
+    } else {
+        return ['unknown']
+    }
+
+    if (isValidCodeCatalystConnection(conn)) {
+        authIds.push(`${connType}CodeCatalyst`)
+    }
+    if (isValidAmazonQConnection(conn)) {
+        authIds.push(`${connType}CodeWhisperer`)
+    }
+
+    return authIds
+}

packages/core/src/codewhisperer/commands/basicCommands.ts

@@ -38,6 +38,7 @@ import { focusAmazonQPanel } from '../../codewhispererChat/commands/registerComm
 import { removeDiagnostic } from '../service/diagnosticsProvider'
 import { SecurityIssueHoverProvider } from '../service/securityIssueHoverProvider'
 import { SecurityIssueCodeActionProvider } from '../service/securityIssueCodeActionProvider'
+import { SsoAccessTokenProvider } from '../../auth/sso/ssoAccessTokenProvider'
 
 export const toggleCodeSuggestions = Commands.declare(
     { id: 'aws.amazonq.toggleCodeSuggestion', compositeKey: { 1: 'source' } },
@@ -183,23 +184,13 @@ export const connectWithCustomization = Commands.declare(
             customizationArn?: string,
             customizationNamePrefix?: string
         ) => {
+            SsoAccessTokenProvider.authSource = source
             if (startUrl && region) {
                 await connectToEnterpriseSso(startUrl, region)
             } else {
                 await getStartUrl()
             }
 
-            // This shortcut is unusual, and currently would only be used if another extension
-            // triggered a connection to Amazon Q. We should still capture and emit the event.
-            telemetry.auth_addConnection.emit({
-                source,
-                isReAuth: false,
-                credentialStartUrl: startUrl,
-                region,
-                authEnabledFeatures: 'codewhisperer',
-                credentialSourceId: 'iamIdentityCenter',
-            })
-
             // No customization match information given, exit early.
             if (!customizationArn && !customizationNamePrefix) {
                 return

packages/core/src/codewhisperer/commands/types.ts

@@ -3,7 +3,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { ExtStartUpSources } from '../../shared/telemetry'
+import { ExtStartUpSources } from '../../shared/telemetry/util'
 import { CompositeKey, Commands, vscodeComponent } from '../../shared/vscode/commands2'
 
 /** Indicates a CodeWhisperer command was executed through a tree node */

packages/core/src/extensionShared.ts

@@ -42,8 +42,8 @@ import { UriHandler } from './shared/vscode/uriHandler'
 import { disableAwsSdkWarning } from './shared/awsClientBuilder'
 import { FileResourceFetcher } from './shared/resourcefetcher/fileResourceFetcher'
 import { ResourceFetcher } from './shared/resourcefetcher/resourcefetcher'
-import { ExtStartUpSources, getAuthFormIdsFromConnection } from './shared/telemetry/util'
-import { ExtensionUse } from './auth/utils'
+import { ExtStartUpSources } from './shared/telemetry/util'
+import { ExtensionUse, getAuthFormIdsFromConnection } from './auth/utils'
 import { Auth } from './auth'
 import { AuthFormId } from './auth/ui/vue/authForms/types'
 

packages/core/src/login/webview/util.ts

@@ -14,13 +14,14 @@ import { vscodeComponent } from '../../shared/vscode/commands2'
  */
 export const AuthSources = {
     addConnectionQuickPick: 'addConnectionQuickPick',
-    firstStartup: ExtStartUpSources.firstStartUp,
+    firstStartUp: ExtStartUpSources.firstStartUp,
     codecatalystDeveloperTools: 'codecatalystDeveloperTools',
     vscodeComponent: vscodeComponent,
     cwQuickPick: cwQuickPickSource,
     cwTreeNode: cwTreeNodeSource,
     amazonQChat: amazonQChatSource,
-    authNode: 'authNode',
+    authNode: 'authNode', // deprecated?
+    unknown: 'unknown',
 } as const
 
 export type AuthSource = (typeof AuthSources)[keyof typeof AuthSources]

packages/core/src/shared/telemetry/index.ts

@@ -4,5 +4,5 @@
  */
 
 export { activate } from './activation'
-export { telemetry } from './telemetry'
+export { telemetry, AuthStatus } from './telemetry'
 export { ExtStartUpSources } from './util'

packages/core/src/shared/telemetry/util.ts

@@ -20,17 +20,6 @@ import { isExtensionInstalled, VSCODE_EXTENSION_ID } from '../utilities'
 import { randomUUID } from '../../common/crypto'
 import { activateExtension } from '../utilities/vsCodeUtils'
 import { ClassToInterfaceType } from '../utilities/tsUtils'
-import {
-    Connection,
-    hasScopes,
-    isBuilderIdConnection,
-    isIamConnection,
-    isIdcSsoConnection,
-    isValidCodeCatalystConnection,
-    scopesSsoAccountAccess,
-} from '../../auth/connection'
-import { AuthFormId } from '../../auth/ui/vue/authForms/types'
-import { isValidAmazonQConnection } from '../../codewhisperer/util/authUtil'
 
 const legacySettingsTelemetryValueDisable = 'Disable'
 const legacySettingsTelemetryValueEnable = 'Enable'
@@ -268,33 +257,3 @@ export const ExtStartUpSources = {
 } as const
 
 export type ExtStartUpSource = (typeof ExtStartUpSources)[keyof typeof ExtStartUpSources]
-
-/**
- * Returns readable auth Ids for a connection, which are used by telemetry.
- */
-export function getAuthFormIdsFromConnection(conn: Connection): AuthFormId[] {
-    const authIds: AuthFormId[] = []
-    let connType: 'builderId' | 'identityCenter'
-
-    if (isIamConnection(conn)) {
-        return ['credentials']
-    } else if (isBuilderIdConnection(conn)) {
-        connType = 'builderId'
-    } else if (isIdcSsoConnection(conn)) {
-        connType = 'identityCenter'
-        if (hasScopes(conn, scopesSsoAccountAccess)) {
-            authIds.push('identityCenterExplorer')
-        }
-    } else {
-        return ['unknown']
-    }
-
-    if (isValidCodeCatalystConnection(conn)) {
-        authIds.push(`${connType}CodeCatalyst`)
-    }
-    if (isValidAmazonQConnection(conn)) {
-        authIds.push(`${connType}CodeWhisperer`)
-    }
-
-    return authIds
-}