Skip to content

Commit c10819e

Browse files
zixlin7zixlin7
authored
feat: add userWrittenCodeTracker (#1308)
* feat: add userWrittenCodeTracker * fix: add unit test
1 parent c8a9044 commit c10819e

File tree

9 files changed

+324
-2
lines changed

9 files changed

+324
-2
lines changed

‎server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.ts‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -122,6 +122,7 @@ import {
122122
import { URI } from 'vscode-uri'
123123
import { AgenticChatError, customerFacingErrorCodes, unactionableErrorCodes } from './errors'
124124
import { CommandCategory } from './tools/executeBash'
125+
import { UserWrittenCodeTracker } from '../../shared/userWrittenCodeTracker'
125126

126127
type ChatHandlers = Omit<
127128
LspHandlers<Chat>,
@@ -148,6 +149,7 @@ export class AgenticChatController implements ChatHandlers {
148149
#additionalContextProvider: AdditionalContextProvider
149150
#contextCommandsProvider: ContextCommandsProvider
150151
#stoppedToolUses = new Set<string>()
152+
#userWrittenCodeTracker: UserWrittenCodeTracker | undefined
151153

152154
/**
153155
* Determines the appropriate message ID for a tool use based on tool type and name
@@ -341,6 +343,7 @@ export class AgenticChatController implements ChatHandlers {
341343
this.#telemetryController.dispose()
342344
this.#chatHistoryDb.close()
343345
this.#contextCommandsProvider?.dispose()
346+
this.#userWrittenCodeTracker?.dispose()
344347
}
345348

346349
async onListConversations(params: ListConversationsParams) {
@@ -387,6 +390,9 @@ export class AgenticChatController implements ChatHandlers {
387390

388391
try {
389392
const triggerContext = await this.#getTriggerContext(params, metric)
393+
if (triggerContext.programmingLanguage?.languageName) {
394+
this.#userWrittenCodeTracker?.recordUsageCount(triggerContext.programmingLanguage.languageName)
395+
}
390396
const isNewConversation = !session.conversationId
391397
session.contextListSent = false
392398
if (isNewConversation) {
@@ -1798,7 +1804,10 @@ export class AgenticChatController implements ChatHandlers {
17981804
],
17991805
},
18001806
}
1807+
1808+
this.#userWrittenCodeTracker?.onQStartsMakingEdits()
18011809
const applyResult = await this.#features.lsp.workspace.applyWorkspaceEdit(workspaceEdit)
1810+
this.#userWrittenCodeTracker?.onQFinishesEdits()
18021811

18031812
if (applyResult.applied) {
18041813
this.#log(`Q Chat server inserted code successfully`)
@@ -2213,6 +2222,9 @@ export class AgenticChatController implements ChatHandlers {
22132222

22142223
updateConfiguration = (newConfig: AmazonQWorkspaceConfig) => {
22152224
this.#customizationArn = newConfig.customizationArn
2225+
if (newConfig.sendUserWrittenCodeMetrics) {
2226+
this.#userWrittenCodeTracker = UserWrittenCodeTracker.getInstance(this.#telemetryService)
2227+
}
22162228
this.#log(`Chat configuration updated customizationArn to ${this.#customizationArn}`)
22172229
/*
22182230
The flag enableTelemetryEventsToDestination is set to true temporarily. It's value will be determined through destination

‎server/aws-lsp-codewhisperer/src/language-server/inline-completion/codeWhispererServer.ts‎

Lines changed: 27 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,7 @@ import { getOrThrowBaseIAMServiceManager } from '../../shared/amazonQServiceMana
5252
// import { WorkspaceFolderManager } from '../workspaceContext/workspaceFolderManager'
5353
import path = require('path')
5454
import { getRelativePath } from '../workspaceContext/util'
55+
import { UserWrittenCodeTracker } from '../../shared/userWrittenCodeTracker'
5556

5657
const EMPTY_RESULT = { sessionId: '', items: [] }
5758
export const CONTEXT_CHARACTERS_LIMIT = 10240
@@ -276,6 +277,7 @@ export const CodewhispererServerFactory =
276277

277278
// CodePercentage and codeDiff tracker have a dependency on TelemetryService, so initialization is also delayed to `onInitialized` handler
278279
let codePercentageTracker: CodePercentageTracker
280+
let userWrittenCodeTracker: UserWrittenCodeTracker | undefined
279281
let codeDiffTracker: CodeDiffTracker<AcceptedInlineSuggestionEntry>
280282

281283
const onInlineCompletionHandler = async (
@@ -458,6 +460,8 @@ export const CodewhispererServerFactory =
458460
): Promise<InlineCompletionListWithReferences> => {
459461
codePercentageTracker.countInvocation(session.language)
460462

463+
userWrittenCodeTracker?.recordUsageCount(session.language)
464+
461465
if (isNewSession) {
462466
// Populate the session with information from codewhisperer response
463467
session.suggestions = suggestionResponse.suggestions
@@ -491,7 +495,7 @@ export const CodewhispererServerFactory =
491495
sessionManager.activateSession(session)
492496

493497
// Process suggestions to apply Empty or Filter filters
494-
const filteredSuggestions = suggestionResponse.suggestions
498+
const filteredSuggestions = session.suggestions
495499
// Empty suggestion filter
496500
.filter(suggestion => {
497501
if (suggestion.content === '') {
@@ -644,9 +648,15 @@ export const CodewhispererServerFactory =
644648
const updateConfiguration = (updatedConfig: AmazonQWorkspaceConfig) => {
645649
logging.debug('Updating configuration of inline complete server.')
646650

647-
const { customizationArn, optOutTelemetryPreference } = updatedConfig
651+
const { customizationArn, optOutTelemetryPreference, sendUserWrittenCodeMetrics } = updatedConfig
648652

649653
codePercentageTracker.customizationArn = customizationArn
654+
if (sendUserWrittenCodeMetrics) {
655+
userWrittenCodeTracker = UserWrittenCodeTracker.getInstance(telemetryService)
656+
}
657+
if (userWrittenCodeTracker) {
658+
userWrittenCodeTracker.customizationArn = customizationArn
659+
}
650660
logging.debug(`CodePercentageTracker customizationArn updated to ${customizationArn}`)
651661
/*
652662
The flag enableTelemetryEventsToDestination is set to true temporarily. It's value will be determined through destination
@@ -706,6 +716,20 @@ export const CodewhispererServerFactory =
706716

707717
p.contentChanges.forEach(change => {
708718
codePercentageTracker.countTotalTokens(languageId, change.text, false)
719+
720+
const { sendUserWrittenCodeMetrics } = amazonQServiceManager.getConfiguration()
721+
if (!sendUserWrittenCodeMetrics) {
722+
return
723+
}
724+
// exclude cases that the document change is from Q suggestions
725+
const currentSession = sessionManager.getCurrentSession()
726+
if (
727+
!currentSession?.suggestions.some(
728+
suggestion => suggestion?.insertText && suggestion.insertText === change.text
729+
)
730+
) {
731+
userWrittenCodeTracker?.countUserWrittenTokens(languageId, change.text)
732+
}
709733
})
710734

711735
// Record last user modification time for any document
@@ -719,6 +743,7 @@ export const CodewhispererServerFactory =
719743

720744
return async () => {
721745
codePercentageTracker?.dispose()
746+
userWrittenCodeTracker?.dispose()
722747
await codeDiffTracker?.shutdown()
723748
}
724749
}

‎server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/configurationUtils.test.ts‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@ describe('getAmazonQRelatedWorkspaceConfigs', () => {
3434
includeSuggestionsWithCodeReferences: true,
3535
includeImportsWithSuggestions: true,
3636
shareCodeWhispererContentWithAWS: true,
37+
sendUserWrittenCodeMetrics: false,
3738
}
3839

3940
beforeEach(() => {
@@ -53,6 +54,7 @@ describe('getAmazonQRelatedWorkspaceConfigs', () => {
5354
includeSuggestionsWithCodeReferences: MOCKED_AWS_CODEWHISPERER_SECTION.includeSuggestionsWithCodeReferences,
5455
includeImportsWithSuggestions: MOCKED_AWS_CODEWHISPERER_SECTION.includeImportsWithSuggestions,
5556
shareCodeWhispererContentWithAWS: MOCKED_AWS_CODEWHISPERER_SECTION.shareCodeWhispererContentWithAWS,
57+
sendUserWrittenCodeMetrics: MOCKED_AWS_CODEWHISPERER_SECTION.sendUserWrittenCodeMetrics,
5658
projectContext: {
5759
enableLocalIndexing: MOCKED_AWS_Q_SECTION.projectContext.enableLocalIndexing,
5860
enableGpuAcceleration: MOCKED_AWS_Q_SECTION.projectContext?.enableGpuAcceleration,
@@ -101,6 +103,7 @@ describe('AmazonQConfigurationCache', () => {
101103
includeSuggestionsWithCodeReferences: false,
102104
includeImportsWithSuggestions: false,
103105
shareCodeWhispererContentWithAWS: true,
106+
sendUserWrittenCodeMetrics: false,
104107
projectContext: {
105108
enableLocalIndexing: true,
106109
enableGpuAcceleration: true,

‎server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/configurationUtils.ts‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -92,6 +92,7 @@ interface CodeWhispererConfigSection {
9292
includeSuggestionsWithCodeReferences: boolean // aws.codeWhisperer.includeSuggestionsWithCodeReferences - return suggestions with code references
9393
includeImportsWithSuggestions: boolean // aws.codeWhisperer.includeImportsWithSuggestions - return imports with suggestions
9494
shareCodeWhispererContentWithAWS: boolean // aws.codeWhisperer.shareCodeWhispererContentWithAWS - share content with AWS
95+
sendUserWrittenCodeMetrics: boolean
9596
}
9697

9798
export type AmazonQWorkspaceConfig = QConfigSection & CodeWhispererConfigSection
@@ -150,6 +151,7 @@ export async function getAmazonQRelatedWorkspaceConfigs(
150151
newCodeWhispererConfig['includeSuggestionsWithCodeReferences'] === true,
151152
includeImportsWithSuggestions: newCodeWhispererConfig['includeImportsWithSuggestions'] === true,
152153
shareCodeWhispererContentWithAWS: newCodeWhispererConfig['shareCodeWhispererContentWithAWS'] === true,
154+
sendUserWrittenCodeMetrics: newCodeWhispererConfig['sendUserWrittenCodeMetrics'] === true,
153155
}
154156

155157
logging.log(
@@ -179,6 +181,7 @@ export const defaultAmazonQWorkspaceConfigFactory = (): AmazonQWorkspaceConfig =
179181
includeSuggestionsWithCodeReferences: false,
180182
includeImportsWithSuggestions: false,
181183
shareCodeWhispererContentWithAWS: false,
184+
sendUserWrittenCodeMetrics: false,
182185
projectContext: {
183186
enableLocalIndexing: false,
184187
enableGpuAcceleration: false,

‎server/aws-lsp-codewhisperer/src/shared/telemetry/telemetryService.test.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -529,6 +529,8 @@ describe('TelemetryService', () => {
529529
acceptedCharacterCount: 123,
530530
totalCharacterCount: 456,
531531
timestamp: new Date(Date.now()),
532+
userWrittenCodeCharacterCount: undefined,
533+
userWrittenCodeLineCount: undefined,
532534
},
533535
},
534536
optOutPreference: 'OPTIN',

‎server/aws-lsp-codewhisperer/src/shared/telemetry/telemetryService.ts‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -341,6 +341,8 @@ export class TelemetryService {
341341
acceptedCharacterCount: number
342342
totalCharacterCount: number
343343
customizationArn?: string
344+
userWrittenCodeCharacterCount?: number
345+
userWrittenCodeLineCount?: number
344346
},
345347
additionalParams: Partial<{
346348
percentage: number
@@ -366,6 +368,8 @@ export class TelemetryService {
366368
acceptedCharacterCount: params.acceptedCharacterCount,
367369
totalCharacterCount: params.totalCharacterCount,
368370
timestamp: new Date(Date.now()),
371+
userWrittenCodeCharacterCount: params.userWrittenCodeCharacterCount,
372+
userWrittenCodeLineCount: params.userWrittenCodeLineCount,
369373
}
370374
if (params.customizationArn) event.customizationArn = params.customizationArn
371375

‎server/aws-lsp-codewhisperer/src/shared/telemetry/types.ts‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -75,6 +75,12 @@ export interface CodeWhispererCodePercentageEvent {
7575
successCount: number
7676
}
7777

78+
export interface UserWrittenPercentageEvent {
79+
codewhispererLanguage: string
80+
userWrittenCodeCharacterCount: number
81+
userWrittenCodeLineCount: number
82+
}
83+
7884
export interface CodeWhispererUserDecisionEvent {
7985
codewhispererRequestId?: string
8086
codewhispererSessionId?: string

‎server/aws-lsp-codewhisperer/src/shared/userWrittenCodeTracker.test.ts‎

Lines changed: 114 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,114 @@
1+
import sinon, { StubbedInstance, stubInterface } from 'ts-sinon'
2+
import { TelemetryService } from './telemetry/telemetryService'
3+
import { UserWrittenCodeTracker } from './userWrittenCodeTracker'
4+
5+
describe('UserWrittenCodeTracker', () => {
6+
const LANGUAGE_ID = 'python'
7+
const OTHER_LANGUAGE_ID = 'typescript'
8+
const SOME_CONTENT = 'some text\n'
9+
10+
let tracker: UserWrittenCodeTracker
11+
let telemetryService: StubbedInstance<TelemetryService> = stubInterface<TelemetryService>()
12+
let clock: sinon.SinonFakeTimers
13+
14+
beforeEach(() => {
15+
clock = sinon.useFakeTimers()
16+
telemetryService = stubInterface<TelemetryService>()
17+
tracker = UserWrittenCodeTracker.getInstance(telemetryService)
18+
})
19+
20+
afterEach(() => {
21+
tracker.dispose()
22+
clock.reset()
23+
clock.restore()
24+
})
25+
26+
it('does not send telemetry without edits', () => {
27+
clock.tick(5000 * 60)
28+
sinon.assert.notCalled(telemetryService.emitCodeCoverageEvent)
29+
})
30+
31+
it('emits metrics every 5 minutes while editing', () => {
32+
tracker.countUserWrittenTokens(LANGUAGE_ID, SOME_CONTENT)
33+
tracker.recordUsageCount(LANGUAGE_ID)
34+
35+
clock.tick(5000 * 60)
36+
37+
sinon.assert.calledWith(
38+
telemetryService.emitCodeCoverageEvent,
39+
{
40+
languageId: LANGUAGE_ID,
41+
totalCharacterCount: 0,
42+
acceptedCharacterCount: 0,
43+
customizationArn: undefined,
44+
userWrittenCodeCharacterCount: 10,
45+
userWrittenCodeLineCount: 1,
46+
},
47+
{}
48+
)
49+
})
50+
51+
it('emits no metrics without invocations', () => {
52+
tracker.countUserWrittenTokens(LANGUAGE_ID, SOME_CONTENT)
53+
54+
clock.tick(5000 * 60 + 1)
55+
56+
sinon.assert.notCalled(telemetryService.emitCodeCoverageEvent)
57+
})
58+
59+
it('emits separate metrics for different languages', () => {
60+
tracker.recordUsageCount(LANGUAGE_ID)
61+
tracker.countUserWrittenTokens(LANGUAGE_ID, SOME_CONTENT)
62+
tracker.recordUsageCount(OTHER_LANGUAGE_ID)
63+
tracker.countUserWrittenTokens(OTHER_LANGUAGE_ID, SOME_CONTENT)
64+
65+
clock.tick(5000 * 60)
66+
67+
sinon.assert.calledWith(
68+
telemetryService.emitCodeCoverageEvent,
69+
{
70+
languageId: LANGUAGE_ID,
71+
totalCharacterCount: 0,
72+
acceptedCharacterCount: 0,
73+
customizationArn: undefined,
74+
userWrittenCodeCharacterCount: 10,
75+
userWrittenCodeLineCount: 1,
76+
},
77+
{}
78+
)
79+
80+
sinon.assert.calledWith(
81+
telemetryService.emitCodeCoverageEvent,
82+
{
83+
languageId: OTHER_LANGUAGE_ID,
84+
totalCharacterCount: 0,
85+
acceptedCharacterCount: 0,
86+
customizationArn: undefined,
87+
userWrittenCodeCharacterCount: 10,
88+
userWrittenCodeLineCount: 1,
89+
},
90+
{}
91+
)
92+
})
93+
94+
it('emits metrics with customizationArn value', () => {
95+
tracker.recordUsageCount(LANGUAGE_ID)
96+
tracker.customizationArn = 'test-arn'
97+
tracker.countUserWrittenTokens(LANGUAGE_ID, SOME_CONTENT)
98+
99+
clock.tick(5000 * 60)
100+
101+
sinon.assert.calledWith(
102+
telemetryService.emitCodeCoverageEvent,
103+
{
104+
languageId: LANGUAGE_ID,
105+
totalCharacterCount: 0,
106+
acceptedCharacterCount: 0,
107+
customizationArn: 'test-arn',
108+
userWrittenCodeCharacterCount: 10,
109+
userWrittenCodeLineCount: 1,
110+
},
111+
{}
112+
)
113+
})
114+
})

0 commit comments

Comments
 (0)