Merge master into feature/cwltail

aws-toolkit-automation · web-flow · commit 5032ce6d083e · 2024-09-25T07:17:15.000-07:00
diff --git a/packages/core/src/codewhisperer/activation.ts b/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) {
diff --git a/packages/core/src/codewhisperer/commands/onInlineAcceptance.ts b/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')
diff --git a/packages/core/src/codewhisperer/service/keyStrokeHandler.ts b/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,
diff --git a/packages/core/src/codewhisperer/service/recommendationHandler.ts b/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')
     }
diff --git a/packages/core/src/codewhisperer/service/recommendationService.ts b/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)
     }
 }
diff --git a/packages/core/src/codewhisperer/util/telemetryHelper.ts b/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 = ''
diff --git a/packages/core/src/codewhisperer/views/activeStateController.ts b/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()
diff --git a/packages/core/src/codewhisperer/views/lineAnnotationController.ts b/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)

Original file line number	Diff line number	Diff line change
`@@ -164,6 +164,7 @@ export class KeyStrokeHandler {`
`164`	`164`	`if (!editor) {`
`165`	`165`	`return`
`166`	`166`	`}`
	`167`	`+`
`167`	`168`	`// RecommendationHandler.instance.reportUserDecisionOfRecommendation(editor, -1)`
`168`	`169`	`await RecommendationService.instance.generateRecommendation(`
`169`	`170`	`client,`