Adding toolkit telemetry for build and execute in UTG

laileni-aws · laileni-aws · commit afbf326efdf2 · 2025-02-21T15:56:43.000-08:00
diff --git a/packages/core/src/amazonqTest/chat/controller/controller.ts b/packages/core/src/amazonqTest/chat/controller/controller.ts
@@ -285,19 +285,35 @@ export class TestController {
         if (session.stopIteration) {
             telemetryErrorMessage = getTelemetryReasonDesc(data.error.uiMessage.replaceAll('```', ''))
         }
-        TelemetryHelper.instance.sendTestGenerationToolkitEvent(
-            session,
-            true,
-            true,
-            isCancel ? 'Cancelled' : 'Failed',
-            session.startTestGenerationRequestId,
-            performance.now() - session.testGenerationStartTime,
-            telemetryErrorMessage,
-            session.isCodeBlockSelected,
-            session.artifactsUploadDuration,
-            session.srcPayloadSize,
-            session.srcZipFileSize
-        )
+        if (session.listOfTestGenerationJobId.length > 1) {
+            TelemetryHelper.instance.sendUnitTestGenerationEvent(
+                session,
+                isCancel ? 'Cancelled' : 'Failed',
+                session.artifactsUploadDuration,
+                session.srcZipFileSize,
+                session.charsOfCodeAccepted,
+                session.numberOfTestsGenerated,
+                session.linesOfCodeGenerated,
+                session.charsOfCodeGenerated,
+                session.numberOfTestsGenerated,
+                session.linesOfCodeGenerated,
+                undefined
+            )
+        } else {
+            TelemetryHelper.instance.sendTestGenerationToolkitEvent(
+                session,
+                session.isSupportedLanguage,
+                true,
+                isCancel ? 'Cancelled' : 'Failed',
+                session.startTestGenerationRequestId,
+                performance.now() - session.testGenerationStartTime,
+                telemetryErrorMessage,
+                session.isCodeBlockSelected,
+                session.artifactsUploadDuration,
+                session.srcPayloadSize,
+                session.srcZipFileSize
+            )
+        }
         if (session.stopIteration) {
             // Error from Science
             this.messenger.sendMessage(data.error.uiMessage.replaceAll('```', ''), data.tabID, 'answer')
@@ -503,6 +519,7 @@ export class TestController {
                     unsupportedMessage = `<span style="color: #EE9D28;">&#9888;<b>I'm sorry, but /test only supports Python and Java</b><br></span> I will still generate a suggestion below.`
                 }
                 this.messenger.sendMessage(unsupportedMessage, tabID, 'answer')
+                session.isSupportedLanguage = false
                 await this.onCodeGeneration(
                     session,
                     userPrompt,
@@ -530,6 +547,7 @@ export class TestController {
                     )
                 }
                 session.isCodeBlockSelected = selectionRange !== undefined
+                session.isSupportedLanguage = true
 
                 /**
                  * Zip the project
@@ -613,7 +631,7 @@ export class TestController {
         // return early if references are disabled and there are references
         if (!CodeWhispererSettings.instance.isSuggestionsWithCodeReferencesEnabled() && session.references.length > 0) {
             void vscode.window.showInformationMessage('Your settings do not allow code generation with references.')
-            await this.endSession(data, FollowUpTypes.SkipBuildAndFinish)
+            await this.endSession(data)
             await this.sessionCleanUp()
             return
         }
@@ -690,7 +708,12 @@ export class TestController {
         const rightUri = vscode.Uri.file(path.join(this.tempResultDirPath, 'resultArtifacts', filePath))
         const fileName = path.basename(absolutePath)
         await vscode.commands.executeCommand('vscode.diff', leftUri, rightUri, `${fileName} ${amazonQTabSuffix}`)
-        telemetry.ui_click.emit({ elementId: 'unitTestGeneration_viewDiff' })
+        const elementId =
+            session.listOfTestGenerationJobId.length > 1
+                ? 'unitTestGeneration_viewDiff_Iteration'
+                : 'unitTestGeneration_viewDiff'
+
+        telemetry.ui_click.emit({ elementId })
         session.latencyOfTestGeneration = performance.now() - session.testGenerationStartTime
         this.messenger.sendUpdatePlaceholder(message.tabID, `Please select an action to proceed (Accept or Reject)`)
     }
@@ -774,9 +797,12 @@ export class TestController {
         }
         const document = await vscode.workspace.openTextDocument(absolutePath)
         await vscode.window.showTextDocument(document)
-        // TODO: send the message once again once build is enabled
-        // this.messenger.sendMessage('Accepted', message.tabID, 'prompt')
-        telemetry.ui_click.emit({ elementId: 'unitTestGeneration_acceptDiff' })
+        const elementId =
+            session.listOfTestGenerationJobId.length > 1
+                ? 'unitTestGeneration_acceptDiff_Iteration'
+                : 'unitTestGeneration_acceptDiff'
+
+        telemetry.ui_click.emit({ elementId })
 
         getLogger().info(
             `Generated unit tests are accepted for ${session.fileLanguage ?? 'plaintext'} language with jobId: ${session.listOfTestGenerationJobId[0]}, jobGroupName: ${session.testGenerationJobGroupName}, result: Succeeded`
@@ -802,7 +828,7 @@ export class TestController {
         )
 
         if (!Auth.instance.isInternalAmazonUser()) {
-            await this.endSession(message, FollowUpTypes.SkipBuildAndFinish)
+            await this.endSession(message)
             return
         }
 
@@ -844,6 +870,19 @@ export class TestController {
             })
             this.messenger.sendChatInputEnabled(message.tabID, false)
         } else {
+            TelemetryHelper.instance.sendUnitTestGenerationEvent(
+                session,
+                undefined,
+                session.artifactsUploadDuration,
+                session.srcZipFileSize,
+                session.charsOfCodeAccepted,
+                session.numberOfTestsGenerated,
+                session.linesOfCodeGenerated,
+                session.charsOfCodeGenerated,
+                session.numberOfTestsGenerated,
+                session.linesOfCodeGenerated,
+                undefined
+            )
             this.sessionStorage.getSession().listOfTestGenerationJobId = []
             this.messenger.sendMessage(
                 'You have gone through both iterations and this unit test generation workflow is complete.',
@@ -904,31 +943,65 @@ export class TestController {
     }
 
     // TODO: Check if there are more cases to endSession if yes create a enum or type for step
-    private async endSession(data: any, step: FollowUpTypes) {
+    private async endSession(data: any, step?: FollowUpTypes) {
         this.messenger.sendMessage('Unit test generation completed.', data.tabID, 'answer')
 
         const session = this.sessionStorage.getSession()
         if (step === FollowUpTypes.RejectCode) {
-            TelemetryHelper.instance.sendTestGenerationToolkitEvent(
+            if (session.listOfTestGenerationJobId.length > 1) {
+                // TODO
+                // telemetry.amazonq_unitTestGeneration.emit()
+                TelemetryHelper.instance.sendUnitTestGenerationEvent(
+                    session,
+                    undefined,
+                    session.artifactsUploadDuration,
+                    session.srcZipFileSize,
+                    0,
+                    0,
+                    0,
+                    session.charsOfCodeGenerated,
+                    session.numberOfTestsGenerated,
+                    session.linesOfCodeGenerated,
+                    undefined
+                )
+                telemetry.ui_click.emit({ elementId: 'unitTestGeneration_rejectDiff_Iteration' })
+            } else {
+                TelemetryHelper.instance.sendTestGenerationToolkitEvent(
+                    session,
+                    true,
+                    true,
+                    'Succeeded',
+                    session.startTestGenerationRequestId,
+                    session.latencyOfTestGeneration,
+                    undefined,
+                    session.isCodeBlockSelected,
+                    session.artifactsUploadDuration,
+                    session.srcPayloadSize,
+                    session.srcZipFileSize,
+                    0,
+                    0,
+                    0,
+                    session.charsOfCodeGenerated,
+                    session.numberOfTestsGenerated,
+                    session.linesOfCodeGenerated
+                )
+                telemetry.ui_click.emit({ elementId: 'unitTestGeneration_rejectDiff' })
+            }
+        } else if (step === FollowUpTypes.SkipBuildAndFinish) {
+            TelemetryHelper.instance.sendUnitTestGenerationEvent(
                 session,
-                true,
-                true,
-                'Succeeded',
-                session.startTestGenerationRequestId,
-                session.latencyOfTestGeneration,
                 undefined,
-                session.isCodeBlockSelected,
                 session.artifactsUploadDuration,
-                session.srcPayloadSize,
                 session.srcZipFileSize,
-                0,
-                0,
-                0,
                 session.charsOfCodeGenerated,
                 session.numberOfTestsGenerated,
-                session.linesOfCodeGenerated
+                session.linesOfCodeGenerated,
+                session.charsOfCodeGenerated,
+                session.numberOfTestsGenerated,
+                session.linesOfCodeGenerated,
+                undefined
             )
-            telemetry.ui_click.emit({ elementId: 'unitTestGeneration_rejectDiff' })
+            telemetry.ui_click.emit({ elementId: 'unitTestGeneration_SkipAndFinish' })
         }
 
         await this.sessionCleanUp()
@@ -944,7 +1017,7 @@ export class TestController {
 
     private startInitialBuild(data: any) {
         // TODO: Remove the fallback build command after stable version of backend build command.
-        const userMessage = `Would you like me to help build and execute the test? I’ll run following commands.\n\`\`\`sh\n${this.sessionStorage.getSession().shortAnswer?.buildCommand}\n`
+        const userMessage = `Would you like me to help build and execute the test? I’ll run following commands.\n\`\`\`sh\n${this.getBuildCommands()}\n`
         const followUps: FollowUps = {
             text: '',
             options: [
@@ -1264,6 +1337,19 @@ export class TestController {
 
         if (codeDiffLength === 0 || session.buildStatus === BuildStatus.SUCCESS) {
             this.messenger.sendMessage('Unit test generation workflow is complete.', data.tabID, 'answer')
+            TelemetryHelper.instance.sendUnitTestGenerationEvent(
+                session,
+                undefined,
+                session.artifactsUploadDuration,
+                session.srcZipFileSize,
+                session.charsOfCodeGenerated,
+                session.numberOfTestsGenerated,
+                session.linesOfCodeGenerated,
+                session.charsOfCodeGenerated,
+                session.numberOfTestsGenerated,
+                session.linesOfCodeGenerated,
+                undefined
+            )
             await this.sessionCleanUp()
         }
     }
diff --git a/packages/core/src/amazonqTest/chat/session/session.ts b/packages/core/src/amazonqTest/chat/session/session.ts
@@ -35,6 +35,7 @@ export class Session {
     public testGenerationJob: TestGenerationJob | undefined
 
     // Start Test generation
+    public isSupportedLanguage: boolean = false
     public conversationState: ConversationState = ConversationState.IDLE
     public shortAnswer: ShortAnswer | undefined
     public sourceFilePath: string = ''
diff --git a/packages/core/src/codewhisperer/util/telemetryHelper.ts b/packages/core/src/codewhisperer/util/telemetryHelper.ts
@@ -13,6 +13,7 @@ import {
     CodewhispererPreviousSuggestionState,
     CodewhispererUserDecision,
     CodewhispererUserTriggerDecision,
+    Result,
     telemetry,
 } from '../../shared/telemetry/telemetry'
 import { CodewhispererCompletionType, CodewhispererSuggestionState } from '../../shared/telemetry/telemetry'
@@ -27,7 +28,7 @@ import { CodeWhispererSupplementalContext } from '../models/model'
 import { FeatureConfigProvider } from '../../shared/featureConfig'
 import { CodeScanRemediationsEventType } from '../client/codewhispereruserclient'
 import { CodeAnalysisScope as CodeAnalysisScopeClientSide } from '../models/constants'
-import { Session } from '../../amazonqTest/chat/session/session'
+import { BuildStatus, Session } from '../../amazonqTest/chat/session/session'
 
 export class TelemetryHelper {
     // Some variables for client component latency
@@ -90,7 +91,7 @@ export class TelemetryHelper {
         telemetry.amazonq_utgGenerateTests.emit({
             cwsprChatProgrammingLanguage: session.fileLanguage ?? 'plaintext',
             hasUserPromptSupplied: session.hasUserPromptSupplied,
-            isSupportedLanguage: isSupportedLanguage,
+            isSupportedLanguage: session.isSupportedLanguage,
             isFileInWorkspace: isFileInWorkspace,
             result: result,
             artifactsUploadDuration: artifactsUploadDuration,
@@ -113,6 +114,61 @@ export class TelemetryHelper {
         })
     }
 
+    /**
+     * amazonq_unitTestGeneration event is emitted in 6 cases for Unit Test Generation
+     * 1) If user close chat window
+     * 2) If build succeeds in build and execute
+     * 3) If user reject generated code
+     * 4) If user click on skip and finish button
+     * 5) If user completes 3 iterations in build and execute
+     * 6) If there is any error/exception
+     */
+    public sendUnitTestGenerationEvent(
+        session: Session,
+        reasonDesc?: string,
+        artifactsUploadDuration?: number,
+        buildZipFileBytes?: number,
+        acceptedCharactersCount?: number,
+        acceptedCount?: number,
+        acceptedLinesCount?: number,
+        generatedCharactersCount?: number,
+        generatedCount?: number,
+        generatedLinesCount?: number,
+        reason?: string
+    ) {
+        telemetry.amazonq_unitTestGeneration.emit({
+            cwsprChatProgrammingLanguage: session.fileLanguage ?? 'plaintext',
+            hasUserPromptSupplied: session.hasUserPromptSupplied,
+            isSupportedLanguage: session.isSupportedLanguage,
+            isFileInWorkspace: true, // Always true TODO: Confirm
+            result:
+                session.buildStatus === BuildStatus.SUCCESS
+                    ? 'Succeeded'
+                    : session.buildStatus === BuildStatus.FAILURE
+                      ? 'Failed'
+                      : ('Cancelled' as Result),
+
+            artifactsUploadDuration: artifactsUploadDuration,
+            buildZipFileBytes: buildZipFileBytes,
+            credentialStartUrl: AuthUtil.instance.startUrl,
+            acceptedCharactersCount: acceptedCharactersCount,
+            acceptedCount: acceptedCount,
+            acceptedLinesCount: acceptedLinesCount,
+            generatedCharactersCount: generatedCharactersCount,
+            generatedCount: generatedCount,
+            generatedLinesCount: generatedLinesCount,
+            isCodeBlockSelected: session.isCodeBlockSelected,
+            jobGroup: session.testGenerationJobGroupName,
+            jobId: session.listOfTestGenerationJobId.at(-1) ?? undefined,
+            perfClientLatency: 0, // TODO for V2 version: Calculate the client side latency between iterations
+            requestId: session.startTestGenerationRequestId,
+            reasonDesc: reasonDesc,
+            reason: reason,
+            update: (session.updatedBuildCommands?.length ?? 0) > 0, // If user modifies command return true else false
+            count: session.listOfTestGenerationJobId.length - 1, // Number of build and execute iterations if count = -1 then user did not even start the build and cycle loop and failed/terminated at the vanilla UTG.
+        })
+    }
+
     public recordServiceInvocationTelemetry(
         requestId: string,
         sessionId: string,
diff --git a/packages/core/src/shared/telemetry/vscodeTelemetry.json b/packages/core/src/shared/telemetry/vscodeTelemetry.json
