telemetry(amazonq): send metric data in onCodeGeneration #6226

siakmun-aws · web-flow · commit c449b0d08e38 · 2024-12-19T15:12:33.000-08:00
## Problem
This is a part of the task to implement client side alarms in order to
track success rate for the client.

## Solution
- Emit metric data telemetry on success/failure.
diff --git a/packages/core/src/amazonqFeatureDev/client/featureDev.ts b/packages/core/src/amazonqFeatureDev/client/featureDev.ts
@@ -25,7 +25,12 @@ import { createCodeWhispererChatStreamingClient } from '../../shared/clients/cod
 import { getClientId, getOptOutPreference, getOperatingSystem } from '../../shared/telemetry/util'
 import { extensionVersion } from '../../shared/vscode/env'
 import apiConfig = require('./codewhispererruntime-2022-11-11.json')
-import { FeatureDevCodeAcceptanceEvent, FeatureDevCodeGenerationEvent, TelemetryEvent } from './featuredevproxyclient'
+import {
+    FeatureDevCodeAcceptanceEvent,
+    FeatureDevCodeGenerationEvent,
+    MetricData,
+    TelemetryEvent,
+} from './featuredevproxyclient'
 
 // Re-enable once BE is able to handle retries.
 const writeAPIRetryOptions = {
@@ -299,6 +304,11 @@ export class FeatureDevClient {
         await this.sendFeatureDevEvent('featureDevCodeAcceptanceEvent', event)
     }
 
+    public async sendMetricData(event: MetricData) {
+        getLogger().debug(`featureDevCodeGenerationMetricData: dimensions: ${event.dimensions}`)
+        await this.sendFeatureDevEvent('metricData', event)
+    }
+
     public async sendFeatureDevEvent<T extends keyof TelemetryEvent>(
         eventName: T,
         event: NonNullable<TelemetryEvent[T]>
diff --git a/packages/core/src/amazonqFeatureDev/controllers/chat/controller.ts b/packages/core/src/amazonqFeatureDev/controllers/chat/controller.ts
@@ -30,7 +30,7 @@ import {
 import { codeGenRetryLimit, defaultRetryLimit } from '../../limits'
 import { Session } from '../../session/session'
 import { featureDevScheme, featureName } from '../../constants'
-import { DeletedFileInfo, DevPhase, type NewFileInfo } from '../../types'
+import { DeletedFileInfo, DevPhase, MetricDataOperationName, MetricDataResult, type NewFileInfo } from '../../types'
 import { AuthUtil } from '../../../codewhisperer/util/authUtil'
 import { AuthController } from '../../../amazonq/auth/controller'
 import { getLogger } from '../../../shared/logger'
@@ -413,6 +413,7 @@ export class FeatureDevController {
                 canBeVoted: true,
             })
             this.messenger.sendUpdatePlaceholder(tabID, i18n('AWS.amazonq.featureDev.pillText.generatingCode'))
+            await session.sendMetricDataTelemetry(MetricDataOperationName.StartCodeGeneration, MetricDataResult.Success)
             await session.send(message)
             const filePaths = session.state.filePaths ?? []
             const deletedFiles = session.state.deletedFiles ?? []
@@ -486,6 +487,31 @@ export class FeatureDevController {
                 await session.sendLinesOfCodeGeneratedTelemetry()
             }
             this.messenger.sendUpdatePlaceholder(tabID, i18n('AWS.amazonq.featureDev.pillText.selectOption'))
+        } catch (err: any) {
+            getLogger().error(`${featureName}: Error during code generation: ${err}`)
+
+            let result: string
+            switch (err.constructor.name) {
+                case FeatureDevServiceError.name:
+                    if (err.code === 'EmptyPatchException') {
+                        result = MetricDataResult.LlmFailure
+                    } else if (err.code === 'GuardrailsException' || err.code === 'ThrottlingException') {
+                        result = MetricDataResult.Error
+                    } else {
+                        result = MetricDataResult.Fault
+                    }
+                    break
+                case PromptRefusalException.name:
+                case NoChangeRequiredException.name:
+                    result = MetricDataResult.Error
+                    break
+                default:
+                    result = MetricDataResult.Fault
+                    break
+            }
+
+            await session.sendMetricDataTelemetry(MetricDataOperationName.EndCodeGeneration, result)
+            throw err
         } finally {
             // Finish processing the event
 
@@ -517,6 +543,7 @@ export class FeatureDevController {
                 }
             }
         }
+        await session.sendMetricDataTelemetry(MetricDataOperationName.EndCodeGeneration, MetricDataResult.Success)
     }
 
     private sendUpdateCodeMessage(tabID: string) {
diff --git a/packages/core/src/amazonqFeatureDev/session/session.ts b/packages/core/src/amazonqFeatureDev/session/session.ts
@@ -285,6 +285,25 @@ export class Session {
         return { leftPath, rightPath, ...diff }
     }
 
+    public async sendMetricDataTelemetry(operationName: string, result: string) {
+        await this.proxyClient.sendMetricData({
+            metricName: 'Operation',
+            metricValue: 1,
+            timestamp: new Date(),
+            product: 'FeatureDev',
+            dimensions: [
+                {
+                    name: 'operationName',
+                    value: operationName,
+                },
+                {
+                    name: 'result',
+                    value: result,
+                },
+            ],
+        })
+    }
+
     public async sendLinesOfCodeGeneratedTelemetry() {
         let charactersOfCodeGenerated = 0
         let linesOfCodeGenerated = 0
diff --git a/packages/core/src/amazonqFeatureDev/types.ts b/packages/core/src/amazonqFeatureDev/types.ts
@@ -115,3 +115,15 @@ export interface UpdateFilesPathsParams {
     messageId: string
     disableFileActions?: boolean
 }
+
+export enum MetricDataOperationName {
+    StartCodeGeneration = 'StartCodeGeneration',
+    EndCodeGeneration = 'EndCodeGeneration',
+}
+
+export enum MetricDataResult {
+    Success = 'Success',
+    Fault = 'Fault',
+    Error = 'Error',
+    LlmFailure = 'LLMFailure',
+}
diff --git a/packages/core/src/test/amazonqFeatureDev/controllers/chat/controller.test.ts b/packages/core/src/test/amazonqFeatureDev/controllers/chat/controller.test.ts
@@ -9,7 +9,13 @@ import * as path from 'path'
 import sinon from 'sinon'
 import { waitUntil } from '../../../../shared/utilities/timeoutUtils'
 import { ControllerSetup, createController, createSession, generateVirtualMemoryUri } from '../../utils'
-import { CurrentWsFolders, DeletedFileInfo, NewFileInfo } from '../../../../amazonqFeatureDev/types'
+import {
+    CurrentWsFolders,
+    DeletedFileInfo,
+    MetricDataOperationName,
+    MetricDataResult,
+    NewFileInfo,
+} from '../../../../amazonqFeatureDev/types'
 import { Session } from '../../../../amazonqFeatureDev/session/session'
 import { Prompter } from '../../../../shared/ui/prompter'
 import { assertTelemetry, toFile } from '../../../testUtil'
@@ -36,6 +42,7 @@ import { AuthUtil } from '../../../../codewhisperer'
 import { featureDevScheme, featureName, messageWithConversationId } from '../../../../amazonqFeatureDev'
 import { i18n } from '../../../../shared/i18n-helper'
 import { FollowUpTypes } from '../../../../amazonq/commons/types'
+import { ToolkitError } from '../../../../shared'
 
 let mockGetCodeGeneration: sinon.SinonStub
 describe('Controller', () => {
@@ -395,7 +402,47 @@ describe('Controller', () => {
     })
 
     describe('processUserChatMessage', function () {
-        async function fireChatMessage() {
+        // TODO: fix disablePreviousFileList error
+        const runs = [
+            { name: 'ContentLengthError', error: new ContentLengthError() },
+            {
+                name: 'MonthlyConversationLimitError',
+                error: new MonthlyConversationLimitError('Service Quota Exceeded'),
+            },
+            {
+                name: 'FeatureDevServiceErrorGuardrailsException',
+                error: new FeatureDevServiceError(
+                    i18n('AWS.amazonq.featureDev.error.codeGen.default'),
+                    'GuardrailsException'
+                ),
+            },
+            {
+                name: 'FeatureDevServiceErrorEmptyPatchException',
+                error: new FeatureDevServiceError(
+                    i18n('AWS.amazonq.featureDev.error.throttling'),
+                    'EmptyPatchException'
+                ),
+            },
+            {
+                name: 'FeatureDevServiceErrorThrottlingException',
+                error: new FeatureDevServiceError(
+                    i18n('AWS.amazonq.featureDev.error.codeGen.default'),
+                    'ThrottlingException'
+                ),
+            },
+            { name: 'UploadCodeError', error: new UploadCodeError('403: Forbiden') },
+            { name: 'UserMessageNotFoundError', error: new UserMessageNotFoundError() },
+            { name: 'TabIdNotFoundError', error: new TabIdNotFoundError() },
+            { name: 'PrepareRepoFailedError', error: new PrepareRepoFailedError() },
+            { name: 'PromptRefusalException', error: new PromptRefusalException() },
+            { name: 'ZipFileError', error: new ZipFileError() },
+            { name: 'CodeIterationLimitError', error: new CodeIterationLimitError() },
+            { name: 'UploadURLExpired', error: new UploadURLExpired() },
+            { name: 'NoChangeRequiredException', error: new NoChangeRequiredException() },
+            { name: 'default', error: new ToolkitError('Default', { code: 'Default' }) },
+        ]
+
+        async function fireChatMessage(session: Session) {
             const getSessionStub = sinon.stub(controllerSetup.sessionStorage, 'getSession').resolves(session)
 
             controllerSetup.emitters.processHumanChatMessage.fire({
@@ -410,44 +457,121 @@ describe('Controller', () => {
             }, {})
         }
 
-        describe('processErrorChatMessage', function () {
-            // TODO: fix disablePreviousFileList error
-            const runs = [
-                { name: 'ContentLengthError', error: new ContentLengthError() },
-                {
-                    name: 'MonthlyConversationLimitError',
-                    error: new MonthlyConversationLimitError('Service Quota Exceeded'),
-                },
-                {
-                    name: 'FeatureDevServiceError',
-                    error: new FeatureDevServiceError(
-                        i18n('AWS.amazonq.featureDev.error.codeGen.default'),
-                        'GuardrailsException'
-                    ),
-                },
-                { name: 'UploadCodeError', error: new UploadCodeError('403: Forbiden') },
-                { name: 'UserMessageNotFoundError', error: new UserMessageNotFoundError() },
-                { name: 'TabIdNotFoundError', error: new TabIdNotFoundError() },
-                { name: 'PrepareRepoFailedError', error: new PrepareRepoFailedError() },
-                { name: 'PromptRefusalException', error: new PromptRefusalException() },
-                { name: 'ZipFileError', error: new ZipFileError() },
-                { name: 'CodeIterationLimitError', error: new CodeIterationLimitError() },
-                { name: 'UploadURLExpired', error: new UploadURLExpired() },
-                { name: 'NoChangeRequiredException', error: new NoChangeRequiredException() },
-                { name: 'default', error: new Error() },
-            ]
+        describe('onCodeGeneration', function () {
+            let session: any
+            let sendMetricDataTelemetrySpy: sinon.SinonStub
+
+            const errorResultMapping = new Map([
+                ['EmptyPatchException', MetricDataResult.LlmFailure],
+                [PromptRefusalException.name, MetricDataResult.Error],
+                [NoChangeRequiredException.name, MetricDataResult.Error],
+            ])
+
+            function getMetricResult(error: ToolkitError): MetricDataResult {
+                if (error instanceof FeatureDevServiceError && error.code) {
+                    return errorResultMapping.get(error.code) ?? MetricDataResult.Error
+                }
+                return errorResultMapping.get(error.constructor.name) ?? MetricDataResult.Fault
+            }
+
+            async function createCodeGenState() {
+                mockGetCodeGeneration = sinon.stub().resolves({ codeGenerationStatus: { status: 'Complete' } })
+
+                const workspaceFolders = [controllerSetup.workspaceFolder] as CurrentWsFolders
+                const testConfig = {
+                    conversationId: conversationID,
+                    proxyClient: {
+                        createConversation: () => sinon.stub(),
+                        createUploadUrl: () => sinon.stub(),
+                        generatePlan: () => sinon.stub(),
+                        startCodeGeneration: () => sinon.stub(),
+                        getCodeGeneration: () => mockGetCodeGeneration(),
+                        exportResultArchive: () => sinon.stub(),
+                    } as unknown as FeatureDevClient,
+                    workspaceRoots: [''],
+                    uploadId: uploadID,
+                    workspaceFolders,
+                }
+
+                const codeGenState = new CodeGenState(testConfig, getFilePaths(controllerSetup), [], [], tabID, 0, {})
+                const newSession = await createSession({
+                    messenger: controllerSetup.messenger,
+                    sessionState: codeGenState,
+                    conversationID,
+                    tabID,
+                    uploadID,
+                    scheme: featureDevScheme,
+                })
+                return newSession
+            }
+
+            async function verifyException(error: ToolkitError) {
+                sinon.stub(session, 'send').throws(error)
+
+                await fireChatMessage(session)
+                await verifyMetricsCalled()
+                assert.ok(
+                    sendMetricDataTelemetrySpy.calledWith(
+                        MetricDataOperationName.StartCodeGeneration,
+                        MetricDataResult.Success
+                    )
+                )
+                const metricResult = getMetricResult(error)
+                assert.ok(
+                    sendMetricDataTelemetrySpy.calledWith(MetricDataOperationName.EndCodeGeneration, metricResult)
+                )
+            }
+
+            async function verifyMetricsCalled() {
+                await waitUntil(() => Promise.resolve(sendMetricDataTelemetrySpy.callCount >= 2), {})
+            }
+
+            beforeEach(async () => {
+                session = await createCodeGenState()
+                sinon.stub(session, 'preloader').resolves()
+                sendMetricDataTelemetrySpy = sinon.stub(session, 'sendMetricDataTelemetry')
+            })
+
+            it('sends success operation telemetry', async () => {
+                sinon.stub(session, 'send').resolves()
+                sinon.stub(session, 'sendLinesOfCodeGeneratedTelemetry').resolves() // Avoid sending extra telemetry
 
+                await fireChatMessage(session)
+                await verifyMetricsCalled()
+
+                assert.ok(
+                    sendMetricDataTelemetrySpy.calledWith(
+                        MetricDataOperationName.StartCodeGeneration,
+                        MetricDataResult.Success
+                    )
+                )
+                assert.ok(
+                    sendMetricDataTelemetrySpy.calledWith(
+                        MetricDataOperationName.EndCodeGeneration,
+                        MetricDataResult.Success
+                    )
+                )
+            })
+
+            runs.forEach(({ name, error }) => {
+                it(`sends failure operation telemetry on ${name}`, async () => {
+                    await verifyException(error)
+                })
+            })
+        })
+
+        describe('processErrorChatMessage', function () {
             function createTestErrorMessage(message: string) {
                 return createUserFacingErrorMessage(`${featureName} request failed: ${message}`)
             }
 
-            async function verifyException(error: Error) {
+            async function verifyException(error: ToolkitError) {
                 sinon.stub(session, 'preloader').throws(error)
                 const sendAnswerSpy = sinon.stub(controllerSetup.messenger, 'sendAnswer')
                 const sendErrorMessageSpy = sinon.stub(controllerSetup.messenger, 'sendErrorMessage')
                 const sendMonthlyLimitErrorSpy = sinon.stub(controllerSetup.messenger, 'sendMonthlyLimitError')
 
-                await fireChatMessage()
+                await fireChatMessage(session)
 
                 switch (error.constructor.name) {
                     case ContentLengthError.name:
