telemetry(amazonq): Add telemetry tracing for amazon q inline (aws#5655)

## Problem
- The flow from when a message is shown to the user and
rejected/accepted by the user isn't instrumented with any way to trace
events

## Solution
- Create a trace from when the event is shown (either by manual
invocation or auto trigger) and then connect all the telemetry events
through a traceId.

### Implementation details
- The traceId is initially created when manual invocation/auto trigger
occurs
- For every telemetry event thats in the function hierarchy of the
manual invocation/auto trigger, they automatically assume the same
traceId
- When an event listener is called it uses the traceId we've set in
TelemetryHelper to setup the next traceId. This is because the event no
longer belongs to the same context
- Inside of commands we use `telemetry.record` because it allows us to
override the default telemetry id that gets automatically created and
instead associate it with the known trace

Author: jpinkney-aws

packages/core/src/codewhisperer/activation.ts

@@ -200,7 +200,11 @@ export async function activate(context: ExtContext): Promise<void> {
                 await openSettings('amazonQ')
             }
         }),
-        Commands.register('aws.amazonq.refreshAnnotation', async (forceProceed: boolean = false) => {
+        Commands.register('aws.amazonq.refreshAnnotation', async (forceProceed: boolean) => {
+            telemetry.record({
+                traceId: TelemetryHelper.instance.traceId,
+            })
+
             const editor = vscode.window.activeTextEditor
             if (editor) {
                 if (forceProceed) {

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

@@ -30,6 +30,8 @@ import { session } from '../util/codeWhispererSession'
 import path from 'path'
 import { RecommendationService } from '../service/recommendationService'
 import { Container } from '../service/serviceContainer'
+import { telemetry } from '../../shared/telemetry'
+import { TelemetryHelper } from '../util/telemetryHelper'
 
 export const acceptSuggestion = Commands.declare(
     'aws.amazonq.accept',
@@ -46,6 +48,10 @@ export const acceptSuggestion = Commands.declare(
             language: CodewhispererLanguage,
             references: codewhispererClient.References
         ) => {
+            telemetry.record({
+                traceId: TelemetryHelper.instance.traceId,
+            })
+
             RecommendationService.instance.incrementAcceptedCount()
             const editor = vscode.window.activeTextEditor
             await Container.instance.lineAnnotationController.refresh(editor, 'codewhisperer')

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

@@ -164,6 +164,7 @@ export class KeyStrokeHandler {
         if (!editor) {
             return
         }
+
         // RecommendationHandler.instance.reportUserDecisionOfRecommendation(editor, -1)
         await RecommendationService.instance.generateRecommendation(
             client,

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

@@ -59,6 +59,10 @@ const nextCommand = Commands.declare('editor.action.inlineSuggest.showNext', ()
 })
 
 const rejectCommand = Commands.declare('aws.amazonq.rejectCodeSuggestion', () => async () => {
+    telemetry.record({
+        traceId: TelemetryHelper.instance.traceId,
+    })
+
     if (!isCloud9('any')) {
         await vscode.commands.executeCommand('editor.action.inlineSuggest.hide')
     }

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

@@ -6,13 +6,19 @@ import * as vscode from 'vscode'
 import { ConfigurationEntry, GetRecommendationsResponse, vsCodeState } from '../models/model'
 import { isCloud9 } from '../../shared/extensionUtilities'
 import { isInlineCompletionEnabled } from '../util/commonUtil'
-import { CodewhispererAutomatedTriggerType, CodewhispererTriggerType } from '../../shared/telemetry/telemetry'
+import {
+    CodewhispererAutomatedTriggerType,
+    CodewhispererTriggerType,
+    telemetry,
+} from '../../shared/telemetry/telemetry'
 import { AuthUtil } from '../util/authUtil'
 import { isIamConnection } from '../../auth/connection'
 import { RecommendationHandler } from '../service/recommendationHandler'
 import { InlineCompletionService } from '../service/inlineCompletionService'
 import { ClassifierTrigger } from './classifierTrigger'
 import { DefaultCodeWhispererClient } from '../client/codewhisperer'
+import { randomUUID } from '../../shared/crypto'
+import { TelemetryHelper } from '../util/telemetryHelper'
 
 export interface SuggestionActionEvent {
     readonly editor: vscode.TextEditor | undefined
@@ -68,99 +74,107 @@ export class RecommendationService {
             return
         }
 
-        if (isCloud9('any')) {
-            // C9 manual trigger key alt/option + C is ALWAYS enabled because the VSC version C9 is on doesn't support setContextKey which is used for CODEWHISPERER_ENABLED
-            // therefore we need a connection check if there is ANY connection(regardless of the connection's state) connected to CodeWhisperer on C9
-            if (triggerType === 'OnDemand' && !AuthUtil.instance.isConnected()) {
-                return
-            }
+        /**
+         * Use an existing trace ID if invoked through a command (e.g., manual invocation),
+         * otherwise generate a new trace ID
+         */
+        const traceId = telemetry.attributes?.traceId ?? randomUUID()
+        TelemetryHelper.instance.setTraceId(traceId)
+        await telemetry.withTraceId(async () => {
+            if (isCloud9('any')) {
+                // C9 manual trigger key alt/option + C is ALWAYS enabled because the VSC version C9 is on doesn't support setContextKey which is used for CODEWHISPERER_ENABLED
+                // therefore we need a connection check if there is ANY connection(regardless of the connection's state) connected to CodeWhisperer on C9
+                if (triggerType === 'OnDemand' && !AuthUtil.instance.isConnected()) {
+                    return
+                }
 
-            RecommendationHandler.instance.checkAndResetCancellationTokens()
-            vsCodeState.isIntelliSenseActive = false
-            this._isRunning = true
-            let response: GetRecommendationsResponse = {
-                result: 'Failed',
-                errorMessage: undefined,
-                recommendationCount: 0,
-            }
+                RecommendationHandler.instance.checkAndResetCancellationTokens()
+                vsCodeState.isIntelliSenseActive = false
+                this._isRunning = true
+                let response: GetRecommendationsResponse = {
+                    result: 'Failed',
+                    errorMessage: undefined,
+                    recommendationCount: 0,
+                }
 
-            try {
-                this._onSuggestionActionEvent.fire({
-                    editor: editor,
-                    isRunning: true,
-                    triggerType: triggerType,
-                    response: undefined,
-                })
+                try {
+                    this._onSuggestionActionEvent.fire({
+                        editor: editor,
+                        isRunning: true,
+                        triggerType: triggerType,
+                        response: undefined,
+                    })
 
-                if (isCloud9('classic') || isIamConnection(AuthUtil.instance.conn)) {
-                    response = await RecommendationHandler.instance.getRecommendations(
-                        client,
-                        editor,
-                        triggerType,
-                        config,
-                        autoTriggerType,
-                        false
-                    )
-                } else {
-                    if (AuthUtil.instance.isConnectionExpired()) {
-                        await AuthUtil.instance.showReauthenticatePrompt()
+                    if (isCloud9('classic') || isIamConnection(AuthUtil.instance.conn)) {
+                        response = await RecommendationHandler.instance.getRecommendations(
+                            client,
+                            editor,
+                            triggerType,
+                            config,
+                            autoTriggerType,
+                            false
+                        )
+                    } else {
+                        if (AuthUtil.instance.isConnectionExpired()) {
+                            await AuthUtil.instance.showReauthenticatePrompt()
+                        }
+                        response = await RecommendationHandler.instance.getRecommendations(
+                            client,
+                            editor,
+                            triggerType,
+                            config,
+                            autoTriggerType,
+                            true
+                        )
+                    }
+                    if (RecommendationHandler.instance.canShowRecommendationInIntelliSense(editor, true, response)) {
+                        await vscode.commands.executeCommand('editor.action.triggerSuggest').then(() => {
+                            vsCodeState.isIntelliSenseActive = true
+                        })
                     }
-                    response = await RecommendationHandler.instance.getRecommendations(
+                } finally {
+                    this._isRunning = false
+                    this._onSuggestionActionEvent.fire({
+                        editor: editor,
+                        isRunning: false,
+                        triggerType: triggerType,
+                        response: response,
+                    })
+                }
+            } else if (isInlineCompletionEnabled()) {
+                if (triggerType === 'OnDemand') {
+                    ClassifierTrigger.instance.recordClassifierResultForManualTrigger(editor)
+                }
+
+                this._isRunning = true
+                let response: GetRecommendationsResponse | undefined = undefined
+
+                try {
+                    this._onSuggestionActionEvent.fire({
+                        editor: editor,
+                        isRunning: true,
+                        triggerType: triggerType,
+                        response: undefined,
+                    })
+
+                    response = await InlineCompletionService.instance.getPaginatedRecommendation(
                         client,
                         editor,
                         triggerType,
                         config,
                         autoTriggerType,
-                        true
+                        event
                     )
-                }
-                if (RecommendationHandler.instance.canShowRecommendationInIntelliSense(editor, true, response)) {
-                    await vscode.commands.executeCommand('editor.action.triggerSuggest').then(() => {
-                        vsCodeState.isIntelliSenseActive = true
+                } finally {
+                    this._isRunning = false
+                    this._onSuggestionActionEvent.fire({
+                        editor: editor,
+                        isRunning: false,
+                        triggerType: triggerType,
+                        response: response,
                     })
                 }
-            } finally {
-                this._isRunning = false
-                this._onSuggestionActionEvent.fire({
-                    editor: editor,
-                    isRunning: false,
-                    triggerType: triggerType,
-                    response: response,
-                })
-            }
-        } else if (isInlineCompletionEnabled()) {
-            if (triggerType === 'OnDemand') {
-                ClassifierTrigger.instance.recordClassifierResultForManualTrigger(editor)
             }
-
-            this._isRunning = true
-            let response: GetRecommendationsResponse | undefined = undefined
-
-            try {
-                this._onSuggestionActionEvent.fire({
-                    editor: editor,
-                    isRunning: true,
-                    triggerType: triggerType,
-                    response: undefined,
-                })
-
-                response = await InlineCompletionService.instance.getPaginatedRecommendation(
-                    client,
-                    editor,
-                    triggerType,
-                    config,
-                    autoTriggerType,
-                    event
-                )
-            } finally {
-                this._isRunning = false
-                this._onSuggestionActionEvent.fire({
-                    editor: editor,
-                    isRunning: false,
-                    triggerType: triggerType,
-                    response: response,
-                })
-            }
-        }
+        }, traceId)
     }
 }

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

@@ -46,6 +46,8 @@ export class TelemetryHelper {
     private timeToFirstRecommendation = 0
     private classifierResult?: number = undefined
     private classifierThreshold?: number = undefined
+    // variables for tracking end to end sessions
+    public traceId: string = 'notSet'
 
     // use this to distinguish DocumentChangeEvent from CWSPR or from other sources
     public lastSuggestionInDisplay = ''
@@ -90,6 +92,7 @@ export class TelemetryHelper {
             codewhispererSupplementalContextLength: supplementalContextMetadata?.contentsLength,
             codewhispererUserGroup: CodeWhispererUserGroupSettings.getUserGroup().toString(),
             codewhispererCustomizationArn: getSelectedCustomization().arn,
+            traceId: this.traceId,
         }
         telemetry.codewhisperer_serviceInvocation.emit(event)
     }
@@ -118,6 +121,7 @@ export class TelemetryHelper {
             codewhispererSupplementalContextTimeout: supplementalContextMetadata?.isProcessTimeout,
             codewhispererSupplementalContextIsUtg: supplementalContextMetadata?.isUtg,
             codewhispererSupplementalContextLength: supplementalContextMetadata?.contentsLength,
+            traceId: this.traceId,
         })
     }
 
@@ -171,6 +175,7 @@ export class TelemetryHelper {
                 codewhispererSupplementalContextTimeout: supplementalContextMetadata?.isProcessTimeout,
                 codewhispererSupplementalContextIsUtg: supplementalContextMetadata?.isUtg,
                 codewhispererSupplementalContextLength: supplementalContextMetadata?.contentsLength,
+                traceId: this.traceId,
             }
             telemetry.codewhisperer_userDecision.emit(event)
             events.push(event)
@@ -234,6 +239,7 @@ export class TelemetryHelper {
             codewhispererSupplementalContextTimeout: supplementalContextMetadata?.isProcessTimeout,
             codewhispererSupplementalContextIsUtg: supplementalContextMetadata?.isUtg,
             codewhispererSupplementalContextLength: supplementalContextMetadata?.contentsLength,
+            traceId: this.traceId,
         }
         return aggregated
     }
@@ -299,6 +305,7 @@ export class TelemetryHelper {
             codewhispererSupplementalContextStrategyId: supplementalContextMetadata?.strategy,
             codewhispererCharactersAccepted: acceptedRecommendationContent.length,
             codewhispererFeatureEvaluations: FeatureConfigProvider.instance.getFeatureConfigsTelemetry(),
+            traceId: this.traceId,
         }
         telemetry.codewhisperer_userTriggerDecision.emit(aggregated)
         this.prevTriggerDecision = this.getAggregatedSuggestionState(this.sessionDecisions)
@@ -386,6 +393,10 @@ export class TelemetryHelper {
         }
     }
 
+    public setTraceId(traceId: string) {
+        this.traceId = traceId
+    }
+
     private resetUserTriggerDecisionTelemetry() {
         this.sessionDecisions = []
         this.triggerChar = ''

packages/core/src/codewhisperer/views/activeStateController.ts

@@ -11,6 +11,8 @@ import { subscribeOnce } from '../../shared/utilities/vsCodeUtils'
 import { Container } from '../service/serviceContainer'
 import { RecommendationHandler } from '../service/recommendationHandler'
 import { cancellableDebounce } from '../../shared/utilities/functionUtils'
+import { telemetry } from '../../shared/telemetry'
+import { TelemetryHelper } from '../util/telemetryHelper'
 
 export class ActiveStateController implements vscode.Disposable {
     private readonly _disposable: vscode.Disposable
@@ -33,7 +35,9 @@ export class ActiveStateController implements vscode.Disposable {
     constructor(private readonly container: Container) {
         this._disposable = vscode.Disposable.from(
             RecommendationService.instance.suggestionActionEvent(async (e) => {
-                await this.onSuggestionActionEvent(e)
+                await telemetry.withTraceId(async () => {
+                    await this.onSuggestionActionEvent(e)
+                }, TelemetryHelper.instance.traceId)
             }),
             RecommendationHandler.instance.onDidReceiveRecommendation(async (_) => {
                 await this.onDidReceiveRecommendation()

packages/core/src/codewhisperer/views/lineAnnotationController.ts

@@ -20,6 +20,7 @@ import { session } from '../util/codeWhispererSession'
 import { RecommendationHandler } from '../service/recommendationHandler'
 import { runtimeLanguageContext } from '../util/runtimeLanguageContext'
 import { setContext } from '../../shared'
+import { TelemetryHelper } from '../util/telemetryHelper'
 
 const case3TimeWindow = 30000 // 30 seconds
 
@@ -246,24 +247,26 @@ export class LineAnnotationController implements vscode.Disposable {
                 await this.onReady()
             }),
             RecommendationService.instance.suggestionActionEvent(async (e) => {
-                if (!this._isReady) {
-                    return
-                }
-
-                if (this._currentState instanceof ManualtriggerState) {
-                    if (e.triggerType === 'OnDemand' && this._currentState.hasManualTrigger === false) {
-                        this._currentState.hasManualTrigger = true
+                await telemetry.withTraceId(async () => {
+                    if (!this._isReady) {
+                        return
                     }
-                    if (
-                        e.response?.recommendationCount !== undefined &&
-                        e.response?.recommendationCount > 0 &&
-                        this._currentState.hasValidResponse === false
-                    ) {
-                        this._currentState.hasValidResponse = true
+
+                    if (this._currentState instanceof ManualtriggerState) {
+                        if (e.triggerType === 'OnDemand' && this._currentState.hasManualTrigger === false) {
+                            this._currentState.hasManualTrigger = true
+                        }
+                        if (
+                            e.response?.recommendationCount !== undefined &&
+                            e.response?.recommendationCount > 0 &&
+                            this._currentState.hasValidResponse === false
+                        ) {
+                            this._currentState.hasValidResponse = true
+                        }
                     }
-                }
 
-                await this.refresh(e.editor, 'codewhisperer')
+                    await this.refresh(e.editor, 'codewhisperer')
+                }, TelemetryHelper.instance.traceId)
             }),
             this.container.lineTracker.onDidChangeActiveLines(async (e) => {
                 await this.onActiveLinesChanged(e)