From 50c91210736c1226e8bcb4c1af10fe29fe202ee7 Mon Sep 17 00:00:00 2001 From: Zoe Lin <60411978+zixlin7@users.noreply.github.com> Date: Tue, 4 Feb 2025 12:40:52 -0800 Subject: [PATCH] feat(amazonq): pre-fetch next recommendations for inline completions (#6419) ## Problem To support complex recommendations broken into multiple chunks, we need the ability to immediately show user next recommendation when they accept previous one ## Solution pre-fetch next recommendation predicting user accept first recommendation when showing them to user, and if user did accept first recommendation, immediately show next recommendation --- - Treat all work as PUBLIC. Private `feature/x` branches will not be squash-merged at release time. - Your code changes must meet the guidelines in [CONTRIBUTING.md](https://github.com/aws/aws-toolkit-vscode/blob/master/CONTRIBUTING.md#guidelines). - License: I confirm that my contribution is made under the terms of the Apache 2.0 license. --------- Co-authored-by: Will Lo --- ...-b63d765a-24f9-44e6-b826-b905f35e67d1.json | 4 + .../commands/onAcceptance.test.ts | 6 +- .../commands/onInlineAcceptance.test.ts | 11 +- .../service/completionProvider.test.ts | 5 +- .../service/inlineCompletionService.test.ts | 4 +- .../service/recommendationHandler.test.ts | 23 ++- .../codewhisperer/service/telemetry.test.ts | 4 +- .../util/telemetryHelper.test.ts | 4 +- .../src/codewhisperer/client/codewhisperer.ts | 3 +- .../commands/invokeRecommendation.ts | 3 +- .../commands/onInlineAcceptance.ts | 15 +- packages/core/src/codewhisperer/index.ts | 2 +- .../service/completionProvider.ts | 4 +- .../service/inlineCompletionItemProvider.ts | 21 +- .../service/inlineCompletionService.ts | 4 +- .../service/recommendationHandler.ts | 190 +++++++++++++----- .../service/referenceLogViewProvider.ts | 3 +- .../util/codeWhispererSession.ts | 76 ++++--- .../src/codewhisperer/util/telemetryHelper.ts | 14 +- .../views/lineAnnotationController.ts | 3 +- .../core/src/test/codewhisperer/testUtil.ts | 3 +- .../codewhisperer/referenceTracker.test.ts | 3 +- .../codewhisperer/serviceInvocations.test.ts | 3 +- 23 files changed, 292 insertions(+), 116 deletions(-) create mode 100644 packages/amazonq/.changes/next-release/Feature-b63d765a-24f9-44e6-b826-b905f35e67d1.json diff --git a/packages/amazonq/.changes/next-release/Feature-b63d765a-24f9-44e6-b826-b905f35e67d1.json b/packages/amazonq/.changes/next-release/Feature-b63d765a-24f9-44e6-b826-b905f35e67d1.json new file mode 100644 index 00000000000..595e48f0026 --- /dev/null +++ b/packages/amazonq/.changes/next-release/Feature-b63d765a-24f9-44e6-b826-b905f35e67d1.json @@ -0,0 +1,4 @@ +{ + "type": "Feature", + "description": "Inline suggestions: Pre-fetch recommendations to reduce suggestion latency." +} diff --git a/packages/amazonq/test/unit/codewhisperer/commands/onAcceptance.test.ts b/packages/amazonq/test/unit/codewhisperer/commands/onAcceptance.test.ts index 5af3252ec82..a94bbd5a3fe 100644 --- a/packages/amazonq/test/unit/codewhisperer/commands/onAcceptance.test.ts +++ b/packages/amazonq/test/unit/codewhisperer/commands/onAcceptance.test.ts @@ -9,19 +9,21 @@ import * as sinon from 'sinon' import { onAcceptance, AcceptedSuggestionEntry, - session, + CodeWhispererSessionState, CodeWhispererTracker, RecommendationHandler, AuthUtil, + CodeWhispererSession, } from 'aws-core-vscode/codewhisperer' import { resetCodeWhispererGlobalVariables, createMockTextEditor } from 'aws-core-vscode/test' import { assertTelemetryCurried } from 'aws-core-vscode/test' describe('onAcceptance', function () { + let session: CodeWhispererSession describe('onAcceptance', function () { beforeEach(async function () { + session = CodeWhispererSessionState.instance.getSession() await resetCodeWhispererGlobalVariables() - session.reset() }) afterEach(function () { diff --git a/packages/amazonq/test/unit/codewhisperer/commands/onInlineAcceptance.test.ts b/packages/amazonq/test/unit/codewhisperer/commands/onInlineAcceptance.test.ts index ae02e7bd7c3..6cff08c2ded 100644 --- a/packages/amazonq/test/unit/codewhisperer/commands/onInlineAcceptance.test.ts +++ b/packages/amazonq/test/unit/codewhisperer/commands/onInlineAcceptance.test.ts @@ -8,15 +8,22 @@ import * as vscode from 'vscode' import * as sinon from 'sinon' import { resetCodeWhispererGlobalVariables, createMockTextEditor } from 'aws-core-vscode/test' import { assertTelemetryCurried } from 'aws-core-vscode/test' -import { onInlineAcceptance, RecommendationHandler, AuthUtil, session } from 'aws-core-vscode/codewhisperer' +import { + onInlineAcceptance, + RecommendationHandler, + AuthUtil, + CodeWhispererSessionState, + CodeWhispererSession, +} from 'aws-core-vscode/codewhisperer' import { globals } from 'aws-core-vscode/shared' import { extensionVersion } from 'aws-core-vscode/shared' describe('onInlineAcceptance', function () { + let session: CodeWhispererSession describe('onInlineAcceptance', function () { beforeEach(async function () { + session = CodeWhispererSessionState.instance.getSession() await resetCodeWhispererGlobalVariables() - session.reset() }) afterEach(function () { diff --git a/packages/amazonq/test/unit/codewhisperer/service/completionProvider.test.ts b/packages/amazonq/test/unit/codewhisperer/service/completionProvider.test.ts index 956999d64ad..a7f87a346ad 100644 --- a/packages/amazonq/test/unit/codewhisperer/service/completionProvider.test.ts +++ b/packages/amazonq/test/unit/codewhisperer/service/completionProvider.test.ts @@ -12,7 +12,7 @@ import { getLabel, Recommendation, RecommendationHandler, - session, + CodeWhispererSessionState, } from 'aws-core-vscode/codewhisperer' import { createMockDocument, resetCodeWhispererGlobalVariables } from 'aws-core-vscode/test' @@ -39,6 +39,7 @@ describe('completionProviderService', function () { describe('getCompletionItem', function () { it('should return targetCompletionItem given input', function () { + const session = CodeWhispererSessionState.instance.getSession() session.startPos = new vscode.Position(0, 0) RecommendationHandler.instance.requestId = 'mock_requestId_getCompletionItem' session.sessionId = 'mock_sessionId_getCompletionItem' @@ -95,6 +96,7 @@ describe('completionProviderService', function () { describe('getCompletionItems', function () { it('should return completion items for each non-empty recommendation', async function () { + const session = CodeWhispererSessionState.instance.getSession() session.recommendations = [ { content: "\n\t\tconsole.log('Hello world!');\n\t}" }, { content: '\nvar a = 10' }, @@ -106,6 +108,7 @@ describe('completionProviderService', function () { }) it('should return empty completion items when recommendation is empty', async function () { + const session = CodeWhispererSessionState.instance.getSession() session.recommendations = [] const mockPosition = new vscode.Position(14, 83) const mockDocument = createMockDocument() diff --git a/packages/amazonq/test/unit/codewhisperer/service/inlineCompletionService.test.ts b/packages/amazonq/test/unit/codewhisperer/service/inlineCompletionService.test.ts index 18fd7d2f21b..f1618c83dac 100644 --- a/packages/amazonq/test/unit/codewhisperer/service/inlineCompletionService.test.ts +++ b/packages/amazonq/test/unit/codewhisperer/service/inlineCompletionService.test.ts @@ -14,7 +14,7 @@ import { CodeSuggestionsState, ConfigurationEntry, CWInlineCompletionItemProvider, - session, + CodeWhispererSessionState, AuthUtil, listCodeWhispererCommandsId, DefaultCodeWhispererClient, @@ -46,6 +46,7 @@ describe('inlineCompletionService', function () { }) it('should call checkAndResetCancellationTokens before showing inline and next token to be null', async function () { + const session = CodeWhispererSessionState.instance.getSession() const mockEditor = createMockTextEditor() sinon.stub(RecommendationHandler.instance, 'getRecommendations').resolves({ result: 'Succeeded', @@ -70,6 +71,7 @@ describe('inlineCompletionService', function () { describe('clearInlineCompletionStates', function () { it('should remove inline reference and recommendations', async function () { + const session = CodeWhispererSessionState.instance.getSession() const fakeReferences = [ { message: '', diff --git a/packages/amazonq/test/unit/codewhisperer/service/recommendationHandler.test.ts b/packages/amazonq/test/unit/codewhisperer/service/recommendationHandler.test.ts index d8855796df0..9668fa4c5a4 100644 --- a/packages/amazonq/test/unit/codewhisperer/service/recommendationHandler.test.ts +++ b/packages/amazonq/test/unit/codewhisperer/service/recommendationHandler.test.ts @@ -8,7 +8,7 @@ import * as vscode from 'vscode' import * as sinon from 'sinon' import { ReferenceInlineProvider, - session, + CodeWhispererSessionState, AuthUtil, DefaultCodeWhispererClient, RecommendationsList, @@ -55,6 +55,7 @@ describe('recommendationHandler', function () { }) it('should assign correct recommendations given input', async function () { + const session = CodeWhispererSessionState.instance.getSession() assert.strictEqual(CodeWhispererCodeCoverageTracker.instances.size, 0) assert.strictEqual( CodeWhispererCodeCoverageTracker.getTracker(mockEditor.document.languageId)?.serviceInvocationCount, @@ -74,7 +75,7 @@ describe('recommendationHandler', function () { } const handler = new RecommendationHandler() sinon.stub(handler, 'getServerResponse').resolves(mockServerResult) - await handler.getRecommendations(mockClient, mockEditor, 'AutoTrigger', config, 'Enter', false) + await handler.getRecommendations(mockClient, mockEditor, 'AutoTrigger', config, session, 'Enter', false) const actual = session.recommendations const expected: RecommendationsList = [{ content: "print('Hello World!')" }, { content: '' }] assert.deepStrictEqual(actual, expected) @@ -85,6 +86,7 @@ describe('recommendationHandler', function () { }) it('should assign request id correctly', async function () { + const session = CodeWhispererSessionState.instance.getSession() const mockServerResult = { recommendations: [{ content: "print('Hello World!')" }, { content: '' }], $response: { @@ -99,7 +101,7 @@ describe('recommendationHandler', function () { const handler = new RecommendationHandler() sinon.stub(handler, 'getServerResponse').resolves(mockServerResult) sinon.stub(handler, 'isCancellationRequested').returns(false) - await handler.getRecommendations(mockClient, mockEditor, 'AutoTrigger', config, 'Enter', false) + await handler.getRecommendations(mockClient, mockEditor, 'AutoTrigger', config, session, 'Enter', false) assert.strictEqual(handler.requestId, 'test_request') assert.strictEqual(session.sessionId, 'test_request') assert.strictEqual(session.triggerType, 'AutoTrigger') @@ -128,9 +130,10 @@ describe('recommendationHandler', function () { strategy: 'empty', }) sinon.stub(performance, 'now').returns(0.0) + const session = CodeWhispererSessionState.instance.getSession() session.startPos = new vscode.Position(1, 0) session.startCursorOffset = 2 - await handler.getRecommendations(mockClient, mockEditor, 'AutoTrigger', config, 'Enter') + await handler.getRecommendations(mockClient, mockEditor, 'AutoTrigger', config, session, 'Enter') const assertTelemetry = assertTelemetryCurried('codewhisperer_serviceInvocation') assertTelemetry({ codewhispererRequestId: 'test_request', @@ -167,10 +170,11 @@ describe('recommendationHandler', function () { const handler = new RecommendationHandler() sinon.stub(handler, 'getServerResponse').resolves(mockServerResult) sinon.stub(performance, 'now').returns(0.0) + const session = CodeWhispererSessionState.instance.getSession() session.startPos = new vscode.Position(1, 0) session.requestIdList = ['test_request_empty'] session.startCursorOffset = 2 - await handler.getRecommendations(mockClient, mockEditor, 'AutoTrigger', config, 'Enter') + await handler.getRecommendations(mockClient, mockEditor, 'AutoTrigger', config, session, 'Enter') const assertTelemetry = assertTelemetryCurried('codewhisperer_userDecision') assertTelemetry({ codewhispererRequestId: 'test_request_empty', @@ -192,6 +196,7 @@ describe('recommendationHandler', function () { sinon.restore() }) it('should return true if any response is not empty', function () { + const session = CodeWhispererSessionState.instance.getSession() const handler = new RecommendationHandler() session.recommendations = [ { @@ -204,12 +209,14 @@ describe('recommendationHandler', function () { }) it('should return false if response is empty', function () { + const session = CodeWhispererSessionState.instance.getSession() const handler = new RecommendationHandler() session.recommendations = [] assert.ok(!handler.isValidResponse()) }) it('should return false if all response has no string length', function () { + const session = CodeWhispererSessionState.instance.getSession() const handler = new RecommendationHandler() session.recommendations = [{ content: '' }, { content: '' }] assert.ok(!handler.isValidResponse()) @@ -222,6 +229,7 @@ describe('recommendationHandler', function () { }) it('should set the completion type to block given a multi-line suggestion', function () { + const session = CodeWhispererSessionState.instance.getSession() session.setCompletionType(0, { content: 'test\n\n \t\r\nanother test' }) assert.strictEqual(session.getCompletionType(0), 'Block') @@ -233,6 +241,7 @@ describe('recommendationHandler', function () { }) it('should set the completion type to line given a single-line suggestion', function () { + const session = CodeWhispererSessionState.instance.getSession() session.setCompletionType(0, { content: 'test' }) assert.strictEqual(session.getCompletionType(0), 'Line') @@ -241,6 +250,7 @@ describe('recommendationHandler', function () { }) it('should set the completion type to line given a multi-line completion but only one-lien of non-blank sequence', function () { + const session = CodeWhispererSessionState.instance.getSession() session.setCompletionType(0, { content: 'test\n\t' }) assert.strictEqual(session.getCompletionType(0), 'Line') @@ -257,6 +267,7 @@ describe('recommendationHandler', function () { describe('on event change', async function () { beforeEach(function () { + const session = CodeWhispererSessionState.instance.getSession() const fakeReferences = [ { message: '', @@ -274,12 +285,14 @@ describe('recommendationHandler', function () { }) it('should remove inline reference onEditorChange', async function () { + const session = CodeWhispererSessionState.instance.getSession() session.sessionId = 'aSessionId' RecommendationHandler.instance.requestId = 'aRequestId' await RecommendationHandler.instance.onEditorChange() assert.strictEqual(ReferenceInlineProvider.instance.refs.length, 0) }) it('should remove inline reference onFocusChange', async function () { + const session = CodeWhispererSessionState.instance.getSession() session.sessionId = 'aSessionId' RecommendationHandler.instance.requestId = 'aRequestId' await RecommendationHandler.instance.onFocusChange() diff --git a/packages/amazonq/test/unit/codewhisperer/service/telemetry.test.ts b/packages/amazonq/test/unit/codewhisperer/service/telemetry.test.ts index 0f1429f130b..797de801428 100644 --- a/packages/amazonq/test/unit/codewhisperer/service/telemetry.test.ts +++ b/packages/amazonq/test/unit/codewhisperer/service/telemetry.test.ts @@ -22,7 +22,7 @@ import { invokeRecommendation, ConfigurationEntry, RecommendationHandler, - session, + CodeWhispererSessionState, vsCodeCursorUpdateDelay, AuthUtil, } from 'aws-core-vscode/codewhisperer' @@ -36,6 +36,7 @@ type CodeWhispererResponse = ListRecommendationsResponse & { let tempFolder: string describe.skip('CodeWhisperer telemetry', async function () { + const session = CodeWhispererSessionState.instance.getSession() let sandbox: sinon.SinonSandbox let client: DefaultCodeWhispererClient @@ -519,6 +520,7 @@ async function manualTrigger( // Note: RecommendationHandler.isSuggestionVisible seems not to work well, hence not using it async function waitUntilSuggestionSeen(index: number = 0) { + const session = CodeWhispererSessionState.instance.getSession() const state = await waitUntil( async () => { const r = session.getSuggestionState(index) diff --git a/packages/amazonq/test/unit/codewhisperer/util/telemetryHelper.test.ts b/packages/amazonq/test/unit/codewhisperer/util/telemetryHelper.test.ts index e042b1d43a2..99f2585a285 100644 --- a/packages/amazonq/test/unit/codewhisperer/util/telemetryHelper.test.ts +++ b/packages/amazonq/test/unit/codewhisperer/util/telemetryHelper.test.ts @@ -5,7 +5,7 @@ import assert from 'assert' import { assertTelemetryCurried, resetCodeWhispererGlobalVariables } from 'aws-core-vscode/test' -import { TelemetryHelper, Completion, session } from 'aws-core-vscode/codewhisperer' +import { TelemetryHelper, Completion, CodeWhispererSessionState } from 'aws-core-vscode/codewhisperer' import { CodewhispererCompletionType, CodewhispererSuggestionState, @@ -39,6 +39,7 @@ function aCompletion(): Completion { } describe('telemetryHelper', function () { + const session = CodeWhispererSessionState.instance.getSession() describe('clientComponentLatency', function () { let sut: TelemetryHelper @@ -48,6 +49,7 @@ describe('telemetryHelper', function () { afterEach(function () { sinon.restore() + session.reset() }) it('resetClientComponentLatencyTime should reset state variables', function () { diff --git a/packages/core/src/codewhisperer/client/codewhisperer.ts b/packages/core/src/codewhisperer/client/codewhisperer.ts index 7a869a68372..ceaf020fb02 100644 --- a/packages/core/src/codewhisperer/client/codewhisperer.ts +++ b/packages/core/src/codewhisperer/client/codewhisperer.ts @@ -17,7 +17,7 @@ import { isSsoConnection } from '../../auth/connection' import { pageableToCollection } from '../../shared/utilities/collectionUtils' import apiConfig = require('./service-2.json') import userApiConfig = require('./user-service-2.json') -import { session } from '../util/codeWhispererSession' +import { CodeWhispererSessionState } from '../util/codeWhispererSession' import { getLogger } from '../../shared/logger' import { indent } from '../../shared/utilities/textUtilities' import { keepAliveHeader } from './agent' @@ -133,6 +133,7 @@ export class DefaultCodeWhispererClient { } async createUserSdkClient(maxRetries?: number): Promise { + const session = CodeWhispererSessionState.instance.getSession() const isOptedOut = CodeWhispererSettings.instance.isOptoutEnabled() session.setFetchCredentialStart() const bearerToken = await AuthUtil.instance.getBearerToken() diff --git a/packages/core/src/codewhisperer/commands/invokeRecommendation.ts b/packages/core/src/codewhisperer/commands/invokeRecommendation.ts index 37fcb965774..e1f17e8a909 100644 --- a/packages/core/src/codewhisperer/commands/invokeRecommendation.ts +++ b/packages/core/src/codewhisperer/commands/invokeRecommendation.ts @@ -8,7 +8,7 @@ import { vsCodeState, ConfigurationEntry } from '../models/model' import { resetIntelliSenseState } from '../util/globalStateUtil' import { DefaultCodeWhispererClient } from '../client/codewhisperer' import { RecommendationHandler } from '../service/recommendationHandler' -import { session } from '../util/codeWhispererSession' +import { CodeWhispererSessionState } from '../util/codeWhispererSession' import { RecommendationService } from '../service/recommendationService' /** @@ -33,6 +33,7 @@ export async function invokeRecommendation( /** * When using intelliSense, if invocation position changed, reject previous active recommendations */ + const session = CodeWhispererSessionState.instance.getSession() if (vsCodeState.isIntelliSenseActive && editor.selection.active !== session.startPos) { resetIntelliSenseState( config.isManualTriggerEnabled, diff --git a/packages/core/src/codewhisperer/commands/onInlineAcceptance.ts b/packages/core/src/codewhisperer/commands/onInlineAcceptance.ts index da581d1aacc..a9a888ef8c4 100644 --- a/packages/core/src/codewhisperer/commands/onInlineAcceptance.ts +++ b/packages/core/src/codewhisperer/commands/onInlineAcceptance.ts @@ -26,7 +26,7 @@ import { import { ReferenceLogViewProvider } from '../service/referenceLogViewProvider' import { ReferenceHoverProvider } from '../service/referenceHoverProvider' import { ImportAdderProvider } from '../service/importAdderProvider' -import { session } from '../util/codeWhispererSession' +import { CodeWhispererSessionState } from '../util/codeWhispererSession' import path from 'path' import { RecommendationService } from '../service/recommendationService' import { Container } from '../service/serviceContainer' @@ -89,6 +89,7 @@ export async function onInlineAcceptance(acceptanceEntry: OnRecommendationAccept const end = acceptanceEntry.editor.selection.active vsCodeState.isCodeWhispererEditing = true + const session = CodeWhispererSessionState.instance.getSession() /** * Mitigation to right context handling mainly for auto closing bracket use case */ @@ -142,5 +143,17 @@ export async function onInlineAcceptance(acceptanceEntry: OnRecommendationAccept } RecommendationHandler.instance.reportUserDecisions(acceptanceEntry.acceptIndex) + await promoteNextSessionIfAvailable(acceptanceEntry) + } +} + +async function promoteNextSessionIfAvailable(acceptanceEntry: OnRecommendationAcceptanceEntry) { + if (acceptanceEntry.acceptIndex === 0 && acceptanceEntry.editor) { + const nextSession = CodeWhispererSessionState.instance.getNextSession() + nextSession.startPos = acceptanceEntry.editor.selection.active + CodeWhispererSessionState.instance.setSession(nextSession) + if (nextSession.recommendations.length) { + await RecommendationHandler.instance.tryShowRecommendation() + } } } diff --git a/packages/core/src/codewhisperer/index.ts b/packages/core/src/codewhisperer/index.ts index 3aea72fb4ca..69b7cbcf96f 100644 --- a/packages/core/src/codewhisperer/index.ts +++ b/packages/core/src/codewhisperer/index.ts @@ -59,7 +59,7 @@ export { onAcceptance } from './commands/onAcceptance' export { CodeWhispererTracker } from './tracker/codewhispererTracker' export { RecommendationHandler } from './service/recommendationHandler' export { CodeWhispererUserGroupSettings } from './util/userGroupUtil' -export { session } from './util/codeWhispererSession' +export { CodeWhispererSessionState, CodeWhispererSession } from './util/codeWhispererSession' export { onInlineAcceptance } from './commands/onInlineAcceptance' export { stopTransformByQ } from './commands/startTransformByQ' export { getCompletionItems, getCompletionItem, getLabel } from './service/completionProvider' diff --git a/packages/core/src/codewhisperer/service/completionProvider.ts b/packages/core/src/codewhisperer/service/completionProvider.ts index 226d04dec2b..df4f2f98466 100644 --- a/packages/core/src/codewhisperer/service/completionProvider.ts +++ b/packages/core/src/codewhisperer/service/completionProvider.ts @@ -9,12 +9,13 @@ import { runtimeLanguageContext } from '../util/runtimeLanguageContext' import { Recommendation } from '../client/codewhisperer' import { LicenseUtil } from '../util/licenseUtil' import { RecommendationHandler } from './recommendationHandler' -import { session } from '../util/codeWhispererSession' +import { CodeWhispererSessionState } from '../util/codeWhispererSession' import path from 'path' /** * completion provider for intelliSense popup */ export function getCompletionItems(document: vscode.TextDocument, position: vscode.Position) { + const session = CodeWhispererSessionState.instance.getSession() const completionItems: vscode.CompletionItem[] = [] for (const [index, recommendation] of session.recommendations.entries()) { completionItems.push(getCompletionItem(document, position, recommendation, index)) @@ -29,6 +30,7 @@ export function getCompletionItem( recommendationDetail: Recommendation, recommendationIndex: number ) { + const session = CodeWhispererSessionState.instance.getSession() const start = session.startPos const range = new vscode.Range(start, start) const recommendation = recommendationDetail.content diff --git a/packages/core/src/codewhisperer/service/inlineCompletionItemProvider.ts b/packages/core/src/codewhisperer/service/inlineCompletionItemProvider.ts index a6c424c321d..dedd2531901 100644 --- a/packages/core/src/codewhisperer/service/inlineCompletionItemProvider.ts +++ b/packages/core/src/codewhisperer/service/inlineCompletionItemProvider.ts @@ -5,7 +5,7 @@ import vscode, { Position } from 'vscode' import { getPrefixSuffixOverlap } from '../util/commonUtil' import { Recommendation } from '../client/codewhisperer' -import { session } from '../util/codeWhispererSession' +import { CodeWhispererSessionState } from '../util/codeWhispererSession' import { TelemetryHelper } from '../util/telemetryHelper' import { runtimeLanguageContext } from '../util/runtimeLanguageContext' import { ReferenceInlineProvider } from './referenceInlineProvider' @@ -21,6 +21,7 @@ export class CWInlineCompletionItemProvider implements vscode.InlineCompletionIt private requestId: string private startPos: Position private nextToken: string + private session = CodeWhispererSessionState.instance.getSession() private _onDidShow: vscode.EventEmitter = new vscode.EventEmitter() public readonly onDidShow: vscode.Event = this._onDidShow.event @@ -100,8 +101,8 @@ export class CWInlineCompletionItemProvider implements vscode.InlineCompletionIt const effectiveStart = document.positionAt(document.offsetAt(start) + prefix.length) const truncatedSuggestion = this.truncateOverlapWithRightContext(document, r.content, end) if (truncatedSuggestion.length === 0) { - if (session.getSuggestionState(index) !== 'Showed') { - session.setSuggestionState(index, 'Discard') + if (this.session.getSuggestionState(index) !== 'Showed') { + this.session.setSuggestionState(index, 'Discard') } return undefined } @@ -118,9 +119,9 @@ export class CWInlineCompletionItemProvider implements vscode.InlineCompletionIt index, truncatedSuggestion, this.requestId, - session.sessionId, - session.triggerType, - session.getCompletionType(index), + this.session.sessionId, + this.session.triggerType, + this.session.getCompletionType(index), runtimeLanguageContext.getLanguageContext(document.languageId, path.extname(document.fileName)) .language, r.references, @@ -155,22 +156,22 @@ export class CWInlineCompletionItemProvider implements vscode.InlineCompletionIt const end = position const iteratingIndexes = this.getIteratingIndexes() const prefix = document.getText(new vscode.Range(start, end)).replace(/\r\n/g, '\n') - const matchedCount = session.recommendations.filter( + const matchedCount = this.session.recommendations.filter( (r) => r.content.length > 0 && r.content.startsWith(prefix) && r.content !== prefix ).length for (const i of iteratingIndexes) { - const r = session.recommendations[i] + const r = this.recommendations[i] const item = this.getInlineCompletionItem(document, r, start, end, i, prefix) if (item === undefined) { continue } this.activeItemIndex = i - session.setSuggestionState(i, 'Showed') + this.session.setSuggestionState(i, 'Showed') ReferenceInlineProvider.instance.setInlineReference(this.startPos.line, r.content, r.references) ImportAdderProvider.instance.onShowRecommendation(document, this.startPos.line, r) this.nextMove = 0 TelemetryHelper.instance.setFirstSuggestionShowTime() - session.setPerceivedLatency() + this.session.setPerceivedLatency() UserWrittenCodeTracker.instance.onQStartsMakingEdits() this._onDidShow.fire() if (matchedCount >= 2 || this.nextToken !== '') { diff --git a/packages/core/src/codewhisperer/service/inlineCompletionService.ts b/packages/core/src/codewhisperer/service/inlineCompletionService.ts index 715fd93ad2d..9cd0dda781d 100644 --- a/packages/core/src/codewhisperer/service/inlineCompletionService.ts +++ b/packages/core/src/codewhisperer/service/inlineCompletionService.ts @@ -16,7 +16,7 @@ import { shared } from '../../shared/utilities/functionUtils' import { ClassifierTrigger } from './classifierTrigger' import { getSelectedCustomization } from '../util/customizationUtil' import { codicon, getIcon } from '../../shared/icons' -import { session } from '../util/codeWhispererSession' +import { CodeWhispererSessionState } from '../util/codeWhispererSession' import { noSuggestions } from '../models/constants' import { Commands } from '../../shared/vscode/commands2' import { listCodeWhispererCommandsId } from '../ui/statusBarMenu' @@ -119,6 +119,7 @@ export class InlineCompletionService { errorMessage: undefined, recommendationCount: 0, } + const session = CodeWhispererSessionState.instance.getSession() try { let page = 0 while (page < this.maxPage) { @@ -127,6 +128,7 @@ export class InlineCompletionService { editor, triggerType, config, + session, autoTriggerType, true, page diff --git a/packages/core/src/codewhisperer/service/recommendationHandler.ts b/packages/core/src/codewhisperer/service/recommendationHandler.ts index 1f5096ad1cc..92595df2fdb 100644 --- a/packages/core/src/codewhisperer/service/recommendationHandler.ts +++ b/packages/core/src/codewhisperer/service/recommendationHandler.ts @@ -5,10 +5,15 @@ import * as vscode from 'vscode' import { extensionVersion } from '../../shared/vscode/env' -import { RecommendationsList, DefaultCodeWhispererClient, CognitoCredentialsError } from '../client/codewhisperer' +import { + RecommendationsList, + DefaultCodeWhispererClient, + CognitoCredentialsError, + ListRecommendationsRequest, +} from '../client/codewhisperer' import * as EditorContext from '../util/editorContext' import * as CodeWhispererConstants from '../models/constants' -import { ConfigurationEntry, GetRecommendationsResponse, vsCodeState } from '../models/model' +import { CodeSuggestionsState, ConfigurationEntry, GetRecommendationsResponse, vsCodeState } from '../models/model' import { runtimeLanguageContext } from '../util/runtimeLanguageContext' import { AWSError } from 'aws-sdk' import { isAwsError } from '../../shared/errors' @@ -31,7 +36,7 @@ import { import { CodeWhispererCodeCoverageTracker } from '../tracker/codewhispererCodeCoverageTracker' import { invalidCustomizationMessage } from '../models/constants' import { getSelectedCustomization, switchToBaseCustomizationAndNotify } from '../util/customizationUtil' -import { session } from '../util/codeWhispererSession' +import { CodeWhispererSessionState, CodeWhispererSession } from '../util/codeWhispererSession' import { Commands } from '../../shared/vscode/commands2' import globals from '../../shared/extensionGlobals' import { noSuggestions, updateInlineLockKey } from '../models/constants' @@ -44,6 +49,8 @@ import { indent } from '../../shared/utilities/textUtilities' import path from 'path' import { isIamConnection } from '../../auth/connection' import { UserWrittenCodeTracker } from '../tracker/userWrittenCodeTracker' +import * as codewhispererClient from '../client/codewhisperer' +import { CodeWhispererSettings } from '../util/codewhispererSettings' /** * This class is for getRecommendation/listRecommendation API calls and its states @@ -103,6 +110,7 @@ export class RecommendationHandler { } isValidResponse(): boolean { + const session = CodeWhispererSessionState.instance.getSession() return session.recommendations.some((r) => r.content.trim() !== '') } @@ -154,14 +162,22 @@ export class RecommendationHandler { editor: vscode.TextEditor, triggerType: CodewhispererTriggerType, config: ConfigurationEntry, + session: CodeWhispererSession, autoTriggerType?: CodewhispererAutomatedTriggerType, pagination: boolean = true, page: number = 0, - generate: boolean = isIamConnection(AuthUtil.instance.conn) + generate: boolean = isIamConnection(AuthUtil.instance.conn), + isNextSession: boolean = false ): Promise { let invocationResult: 'Succeeded' | 'Failed' = 'Failed' let errorMessage: string | undefined = undefined let errorCode: string | undefined = undefined + let currentSession = session + if (isNextSession) { + getLogger().debug('pre-fetching next recommendation for model routing') + currentSession = new CodeWhispererSession() + CodeWhispererSessionState.instance.setNextSession(currentSession) + } if (!editor) { return Promise.resolve({ @@ -178,43 +194,60 @@ export class RecommendationHandler { let latency = 0 let nextToken = '' let shouldRecordServiceInvocation = true - session.language = runtimeLanguageContext.getLanguageContext( + currentSession.language = runtimeLanguageContext.getLanguageContext( editor.document.languageId, path.extname(editor.document.fileName) ).language - session.taskType = await this.getTaskTypeFromEditorFileName(editor.document.fileName) + currentSession.taskType = await this.getTaskTypeFromEditorFileName(editor.document.fileName) if (pagination && !generate) { if (page === 0) { - session.requestContext = await EditorContext.buildListRecommendationRequest( - editor as vscode.TextEditor, - this.nextToken, - config.isSuggestionsWithCodeReferencesEnabled - ) + if (isNextSession) { + const request = session.requestContext.request as ListRecommendationsRequest + currentSession.requestContext = { + request: { + ...request, + fileContext: { + ...request.fileContext, + leftFileContent: `${request.fileContext.leftFileContent}${session.recommendations[0].content}`, + }, + nextToken: undefined, + }, + supplementalMetadata: currentSession.requestContext.supplementalMetadata, + } + } else { + currentSession.requestContext = await EditorContext.buildListRecommendationRequest( + editor as vscode.TextEditor, + this.nextToken, + config.isSuggestionsWithCodeReferencesEnabled + ) + } } else { - session.requestContext = { + currentSession.requestContext = { request: { - ...session.requestContext.request, + ...currentSession.requestContext.request, // Putting nextToken assignment in the end so it overwrites the existing nextToken nextToken: this.nextToken, }, - supplementalMetadata: session.requestContext.supplementalMetadata, + supplementalMetadata: currentSession.requestContext.supplementalMetadata, } } + // } } else { - session.requestContext = await EditorContext.buildGenerateRecommendationRequest(editor as vscode.TextEditor) + currentSession.requestContext = await EditorContext.buildGenerateRecommendationRequest( + editor as vscode.TextEditor + ) } - const request = session.requestContext.request - // record preprocessing end time + const request = currentSession.requestContext.request TelemetryHelper.instance.setPreprocessEndTime() // set start pos for non pagination call or first pagination call if (!pagination || (pagination && page === 0)) { - session.startPos = editor.selection.active - session.startCursorOffset = editor.document.offsetAt(session.startPos) - session.leftContextOfCurrentLine = EditorContext.getLeftContext(editor, session.startPos.line) - session.triggerType = triggerType - session.autoTriggerType = autoTriggerType + currentSession.startPos = editor.selection.active + currentSession.startCursorOffset = editor.document.offsetAt(currentSession.startPos) + currentSession.leftContextOfCurrentLine = EditorContext.getLeftContext(editor, currentSession.startPos.line) + currentSession.triggerType = triggerType + currentSession.autoTriggerType = autoTriggerType /** * Validate request @@ -255,7 +288,7 @@ export class RecommendationHandler { sessionId = resp?.$response?.httpResponse?.headers['x-amzn-sessionid'] TelemetryHelper.instance.setFirstResponseRequestId(requestId) if (page === 0) { - session.setTimeToFirstRecommendation(performance.now()) + currentSession.setTimeToFirstRecommendation(performance.now()) } if (nextToken === '') { TelemetryHelper.instance.setAllPaginationEndTime() @@ -301,9 +334,9 @@ export class RecommendationHandler { vscode version: '${vscode.version}', extension version: '${extensionVersion}', filename: '${EditorContext.getFileName(editor)}', - left context of line: '${session.leftContextOfCurrentLine}', - line number: ${session.startPos.line}, - character location: ${session.startPos.character}, + left context of line: '${currentSession.leftContextOfCurrentLine}', + line number: ${currentSession.startPos.line}, + character location: ${currentSession.startPos.character}, latency: ${latency} ms. Recommendations:`, 4, @@ -311,11 +344,11 @@ export class RecommendationHandler { ).trimStart() for (const [index, item] of recommendations.entries()) { msg += `\n ${index.toString().padStart(2, '0')}: ${indent(item.content, 8, true).trim()}` - session.requestIdList.push(requestId) + currentSession.requestIdList.push(requestId) } getLogger().debug(msg) if (invocationResult === 'Succeeded') { - CodeWhispererCodeCoverageTracker.getTracker(session.language)?.incrementServiceInvocationCount() + CodeWhispererCodeCoverageTracker.getTracker(currentSession.language)?.incrementServiceInvocationCount() UserWrittenCodeTracker.instance.onQFeatureInvoked() } else { if ( @@ -331,6 +364,7 @@ export class RecommendationHandler { editor, triggerType, config, + currentSession, autoTriggerType, pagination, page, @@ -343,82 +377,86 @@ export class RecommendationHandler { TelemetryHelper.instance.recordServiceInvocationTelemetry( requestId, sessionId, - session.recommendations.length + recommendations.length - 1, + currentSession.recommendations.length + recommendations.length - 1, invocationResult, latency, - session.language, - session.taskType, + currentSession.language, + currentSession.taskType, reason, - session.requestContext.supplementalMetadata + currentSession.requestContext.supplementalMetadata ) } } - if (this.isCancellationRequested()) { + if (!isNextSession && this.isCancellationRequested()) { return Promise.resolve({ result: invocationResult, errorMessage: errorMessage, - recommendationCount: session.recommendations.length, + recommendationCount: currentSession.recommendations.length, }) } const typedPrefix = editor.document - .getText(new vscode.Range(session.startPos, editor.selection.active)) + .getText(new vscode.Range(currentSession.startPos, editor.selection.active)) .replace('\r\n', '\n') if (recommendations.length > 0) { TelemetryHelper.instance.setTypeAheadLength(typedPrefix.length) // mark suggestions that does not match typeahead when arrival as Discard // these suggestions can be marked as Showed if typeahead can be removed with new inline API for (const [i, r] of recommendations.entries()) { - const recommendationIndex = i + session.recommendations.length + const recommendationIndex = i + currentSession.recommendations.length if ( !r.content.startsWith(typedPrefix) && - session.getSuggestionState(recommendationIndex) === undefined + currentSession.getSuggestionState(recommendationIndex) === undefined ) { - session.setSuggestionState(recommendationIndex, 'Discard') + currentSession.setSuggestionState(recommendationIndex, 'Discard') } - session.setCompletionType(recommendationIndex, r) + currentSession.setCompletionType(recommendationIndex, r) } - session.recommendations = pagination ? session.recommendations.concat(recommendations) : recommendations - if (isInlineCompletionEnabled() && this.hasAtLeastOneValidSuggestion(typedPrefix)) { + currentSession.recommendations = pagination + ? currentSession.recommendations.concat(recommendations) + : recommendations + if (isInlineCompletionEnabled() && this.hasAtLeastOneValidSuggestion(typedPrefix, currentSession)) { this._onDidReceiveRecommendation.fire() } } this.requestId = requestId - session.sessionId = sessionId - this.nextToken = nextToken + currentSession.sessionId = sessionId + if (!isNextSession) { + this.nextToken = nextToken + } // send Empty userDecision event if user receives no recommendations in this session at all. if (invocationResult === 'Succeeded' && nextToken === '') { // case 1: empty list of suggestion [] - if (session.recommendations.length === 0) { - session.requestIdList.push(requestId) + if (currentSession.recommendations.length === 0) { + currentSession.requestIdList.push(requestId) // Received an empty list of recommendations TelemetryHelper.instance.recordUserDecisionTelemetryForEmptyList( - session.requestIdList, + currentSession.requestIdList, sessionId, page, runtimeLanguageContext.getLanguageContext( editor.document.languageId, path.extname(editor.document.fileName) ).language, - session.requestContext.supplementalMetadata + currentSession.requestContext.supplementalMetadata ) } // case 2: non empty list of suggestion but with (a) empty content or (b) non-matching typeahead - else if (!this.hasAtLeastOneValidSuggestion(typedPrefix)) { + else if (!this.hasAtLeastOneValidSuggestion(typedPrefix, currentSession)) { this.reportUserDecisions(-1) } } return Promise.resolve({ result: invocationResult, errorMessage: errorMessage, - recommendationCount: session.recommendations.length, + recommendationCount: currentSession.recommendations.length, }) } - hasAtLeastOneValidSuggestion(typedPrefix: string): boolean { + hasAtLeastOneValidSuggestion(typedPrefix: string, session: CodeWhispererSession): boolean { return session.recommendations.some((r) => r.content.trim() !== '' && r.content.startsWith(typedPrefix)) } @@ -444,6 +482,7 @@ export class RecommendationHandler { * Clear recommendation state */ clearRecommendations() { + const session = CodeWhispererSessionState.instance.getSession() session.requestIdList = [] session.recommendations = [] session.suggestionStates = new Map() @@ -473,6 +512,7 @@ export class RecommendationHandler { } reportDiscardedUserDecisions() { + const session = CodeWhispererSessionState.instance.getSession() for (const [i, _] of session.recommendations.entries()) { session.setSuggestionState(i, 'Discard') } @@ -483,6 +523,7 @@ export class RecommendationHandler { * Emits telemetry reflecting user decision for current recommendation. */ reportUserDecisions(acceptIndex: number) { + const session = CodeWhispererSessionState.instance.getSession() if (session.sessionId === '' || this.requestId === '') { return } @@ -512,6 +553,7 @@ export class RecommendationHandler { showPrompt: boolean = false, response: GetRecommendationsResponse ): boolean { + const session = CodeWhispererSessionState.instance.getSession() const reject = () => { this.reportUserDecisions(-1) } @@ -593,6 +635,11 @@ export class RecommendationHandler { } async showRecommendation(indexShift: number, noSuggestionVisible: boolean = false) { + const session = CodeWhispererSessionState.instance.getSession() + + if (!indexShift && session.recommendations.length) { + await this.fetchNextRecommendations() + } await lock.acquire(updateInlineLockKey, async () => { if (!vscode.window.state.focused) { this.reportDiscardedUserDecisions() @@ -623,7 +670,7 @@ export class RecommendationHandler { } if (noSuggestionVisible) { await vscode.commands.executeCommand(`editor.action.inlineSuggest.trigger`) - this.sendPerceivedLatencyTelemetry() + this.sendPerceivedLatencyTelemetry(session) } }) } @@ -654,7 +701,46 @@ export class RecommendationHandler { return this.inlineCompletionProvider?.getActiveItemIndex !== undefined } + async getConfigEntry(): Promise { + const codewhispererSettings = CodeWhispererSettings.instance + const isShowMethodsEnabled: boolean = + vscode.workspace.getConfiguration('editor').get('suggest.showMethods') || false + const isAutomatedTriggerEnabled: boolean = CodeSuggestionsState.instance.isSuggestionsEnabled() + const isManualTriggerEnabled: boolean = true + const isSuggestionsWithCodeReferencesEnabled = codewhispererSettings.isSuggestionsWithCodeReferencesEnabled() + + return { + isShowMethodsEnabled, + isManualTriggerEnabled, + isAutomatedTriggerEnabled, + isSuggestionsWithCodeReferencesEnabled, + } + } + + async fetchNextRecommendations() { + const session = CodeWhispererSessionState.instance.getSession() + const client = new codewhispererClient.DefaultCodeWhispererClient() + const editor = vscode.window.activeTextEditor + if (!editor) { + return + } + + await this.getRecommendations( + client, + editor, + session.triggerType, + await this.getConfigEntry(), + session, + session.autoTriggerType, + true, + 0, + false, + true + ) + } + async tryShowRecommendation() { + const session = CodeWhispererSessionState.instance.getSession() const editor = vscode.window.activeTextEditor if (editor === undefined) { return @@ -683,7 +769,7 @@ export class RecommendationHandler { } } - private sendPerceivedLatencyTelemetry() { + private sendPerceivedLatencyTelemetry(session: CodeWhispererSession) { if (vscode.window.activeTextEditor) { const languageContext = runtimeLanguageContext.getLanguageContext( vscode.window.activeTextEditor.document.languageId, diff --git a/packages/core/src/codewhisperer/service/referenceLogViewProvider.ts b/packages/core/src/codewhisperer/service/referenceLogViewProvider.ts index 9ec20b8cb44..02055f0aded 100644 --- a/packages/core/src/codewhisperer/service/referenceLogViewProvider.ts +++ b/packages/core/src/codewhisperer/service/referenceLogViewProvider.ts @@ -10,7 +10,7 @@ import * as CodeWhispererConstants from '../models/constants' import { CodeWhispererSettings } from '../util/codewhispererSettings' import globals from '../../shared/extensionGlobals' import { AuthUtil } from '../util/authUtil' -import { session } from '../util/codeWhispererSession' +import { CodeWhispererSessionState } from '../util/codeWhispererSession' export class ReferenceLogViewProvider implements vscode.WebviewViewProvider { public static readonly viewType = 'aws.codeWhisperer.referenceLog' @@ -68,6 +68,7 @@ export class ReferenceLogViewProvider implements vscode.WebviewViewProvider { reference.recommendationContentSpan.start, reference.recommendationContentSpan.end ) + const session = CodeWhispererSessionState.instance.getSession() const firstCharLineNumber = editor.document.positionAt(session.startCursorOffset + reference.recommendationContentSpan.start).line + 1 diff --git a/packages/core/src/codewhisperer/util/codeWhispererSession.ts b/packages/core/src/codewhisperer/util/codeWhispererSession.ts index 042cd947124..12222903568 100644 --- a/packages/core/src/codewhisperer/util/codeWhispererSession.ts +++ b/packages/core/src/codewhisperer/util/codeWhispererSession.ts @@ -14,41 +14,62 @@ import { GenerateRecommendationsRequest, ListRecommendationsRequest, Recommendat import { Position } from 'vscode' import { CodeWhispererSupplementalContext, vsCodeState } from '../models/model' -class CodeWhispererSession { - static #instance: CodeWhispererSession +export class CodeWhispererSessionState { + static #instance: CodeWhispererSessionState + session: CodeWhispererSession + nextSession: CodeWhispererSession + + constructor() { + this.session = new CodeWhispererSession() + this.nextSession = new CodeWhispererSession() + } + public static get instance() { + return (this.#instance ??= new CodeWhispererSessionState()) + } + + getSession() { + return this.session + } + + setSession(session: CodeWhispererSession) { + this.session = session + } + + getNextSession() { + return this.nextSession + } - // Per-session states - sessionId = '' + setNextSession(session: CodeWhispererSession) { + this.nextSession = session + } +} + +export class CodeWhispererSession { + sessionId: string = '' requestIdList: string[] = [] - startPos = new Position(0, 0) - startCursorOffset = 0 - leftContextOfCurrentLine = '' + startPos: Position = new Position(0, 0) + startCursorOffset: number = 0 + leftContextOfCurrentLine: string = '' requestContext: { request: ListRecommendationsRequest | GenerateRecommendationsRequest supplementalMetadata: CodeWhispererSupplementalContext | undefined - } = { request: {} as any, supplementalMetadata: {} as any } + } = { request: {} as any, supplementalMetadata: undefined } language: CodewhispererLanguage = 'python' - taskType: CodewhispererGettingStartedTask | undefined + taskType: CodewhispererGettingStartedTask | undefined = undefined triggerType: CodewhispererTriggerType = 'OnDemand' - autoTriggerType: CodewhispererAutomatedTriggerType | undefined - + autoTriggerType: CodewhispererAutomatedTriggerType | undefined = undefined // Various states of recommendations recommendations: Recommendation[] = [] - suggestionStates = new Map() - completionTypes = new Map() - + suggestionStates: Map = new Map() + completionTypes: Map = new Map() // Some other variables for client component latency - fetchCredentialStartTime = 0 - sdkApiCallStartTime = 0 - invokeSuggestionStartTime = 0 - preprocessEndTime = 0 - timeToFirstRecommendation = 0 - firstSuggestionShowTime = 0 - perceivedLatency = 0 - - public static get instance() { - return (this.#instance ??= new CodeWhispererSession()) - } + fetchCredentialStartTime: number = 0 + sdkApiCallStartTime: number = 0 + invokeSuggestionStartTime: number = 0 + preprocessEndTime: number = 0 + timeToFirstRecommendation: number = 0 + firstSuggestionShowTime: number = 0 + perceivedLatency: number = 0 setFetchCredentialStart() { if (this.fetchCredentialStartTime === 0 && this.invokeSuggestionStartTime !== 0) { @@ -89,7 +110,7 @@ class CodeWhispererSession { if (triggerType === 'OnDemand') { return this.timeToFirstRecommendation } else { - return session.firstSuggestionShowTime - vsCodeState.lastUserModificationTime + return this.firstSuggestionShowTime - vsCodeState.lastUserModificationTime } } @@ -118,6 +139,3 @@ class CodeWhispererSession { this.completionTypes.clear() } } - -// TODO: convert this to a function call -export const session = CodeWhispererSession.instance diff --git a/packages/core/src/codewhisperer/util/telemetryHelper.ts b/packages/core/src/codewhisperer/util/telemetryHelper.ts index 5276d869bb9..964145fd881 100644 --- a/packages/core/src/codewhisperer/util/telemetryHelper.ts +++ b/packages/core/src/codewhisperer/util/telemetryHelper.ts @@ -22,7 +22,7 @@ import { getSelectedCustomization } from './customizationUtil' import { AuthUtil } from './authUtil' import { isAwsError } from '../../shared/errors' import { getLogger } from '../../shared/logger' -import { session } from './codeWhispererSession' +import { CodeWhispererSessionState } from './codeWhispererSession' import { CodeWhispererSupplementalContext } from '../models/model' import { FeatureConfigProvider } from '../../shared/featureConfig' import { CodeScanRemediationsEventType } from '../client/codewhispereruserclient' @@ -124,6 +124,7 @@ export class TelemetryHelper { reason: string, supplementalContextMetadata?: CodeWhispererSupplementalContext | undefined ) { + const session = CodeWhispererSessionState.instance.getSession() const event = { codewhispererAutomatedTriggerType: session.autoTriggerType, codewhispererCursorOffset: session.startCursorOffset, @@ -157,6 +158,7 @@ export class TelemetryHelper { supplementalContextMetadata?: CodeWhispererSupplementalContext | undefined ) { const selectedCustomization = getSelectedCustomization() + const session = CodeWhispererSessionState.instance.getSession() telemetry.codewhisperer_userDecision.emit({ codewhispererCompletionType: 'Line', @@ -277,6 +279,7 @@ export class TelemetryHelper { if (_elem.content.length === 0) { recommendationSuggestionState?.set(i, 'Empty') } + const session = CodeWhispererSessionState.instance.getSession() const event: CodewhispererUserDecision = { // TODO: maintain a list of RecommendationContexts with both recommendation and requestId in it, instead of two separate list items. codewhispererCompletionType: this.getCompletionType(i, completionTypes), @@ -338,6 +341,7 @@ export class TelemetryHelper { if (!events.length) { return } + const session = CodeWhispererSessionState.instance.getSession() const aggregated: CodewhispererUserTriggerDecision = { codewhispererAutomatedTriggerType: session.autoTriggerType, codewhispererCompletionType: events[0].codewhispererCompletionType, @@ -387,6 +391,7 @@ export class TelemetryHelper { .map((e) => e.codewhispererSuggestionCount) .reduce((a, b) => a + b, 0) + const session = CodeWhispererSessionState.instance.getSession() const aggregated: CodewhispererUserTriggerDecision = { codewhispererAutomatedTriggerType: autoTriggerType, codewhispererCharactersAccepted: acceptedRecommendationContent.length, @@ -508,6 +513,7 @@ export class TelemetryHelper { } private resetUserTriggerDecisionTelemetry() { + const session = CodeWhispererSessionState.instance.getSession() this.sessionDecisions = [] this.triggerChar = '' this.typeAheadLength = 0 @@ -592,6 +598,7 @@ export class TelemetryHelper { } public resetClientComponentLatencyTime() { + const session = CodeWhispererSessionState.instance.getSession() session.invokeSuggestionStartTime = 0 session.preprocessEndTime = 0 session.sdkApiCallStartTime = 0 @@ -603,6 +610,7 @@ export class TelemetryHelper { } public setPreprocessEndTime() { + const session = CodeWhispererSessionState.instance.getSession() if (session.preprocessEndTime !== 0) { getLogger().warn(`inline completion preprocessEndTime has been set and not reset correctly`) } @@ -611,11 +619,13 @@ export class TelemetryHelper { /** This method is assumed to be invoked first at the start of execution **/ public setInvokeSuggestionStartTime() { + const session = CodeWhispererSessionState.instance.getSession() this.resetClientComponentLatencyTime() session.invokeSuggestionStartTime = performance.now() } public setSdkApiCallEndTime() { + const session = CodeWhispererSessionState.instance.getSession() if (this._sdkApiCallEndTime === 0 && session.sdkApiCallStartTime !== 0) { this._sdkApiCallEndTime = performance.now() } @@ -628,6 +638,7 @@ export class TelemetryHelper { } public setFirstSuggestionShowTime() { + const session = CodeWhispererSessionState.instance.getSession() if (session.firstSuggestionShowTime === 0 && this._sdkApiCallEndTime !== 0) { session.firstSuggestionShowTime = performance.now() } @@ -642,6 +653,7 @@ export class TelemetryHelper { // report client component latency after all pagination call finish // and at least one suggestion is shown to the user public tryRecordClientComponentLatency() { + const session = CodeWhispererSessionState.instance.getSession() if (session.firstSuggestionShowTime === 0 || this._allPaginationEndTime === 0) { return } diff --git a/packages/core/src/codewhisperer/views/lineAnnotationController.ts b/packages/core/src/codewhisperer/views/lineAnnotationController.ts index 12a52f9b4eb..39732c84064 100644 --- a/packages/core/src/codewhisperer/views/lineAnnotationController.ts +++ b/packages/core/src/codewhisperer/views/lineAnnotationController.ts @@ -16,7 +16,7 @@ import { Container } from '../service/serviceContainer' import { telemetry } from '../../shared/telemetry/telemetry' import { getLogger } from '../../shared/logger/logger' import { Commands } from '../../shared/vscode/commands2' -import { session } from '../util/codeWhispererSession' +import { CodeWhispererSessionState } from '../util/codeWhispererSession' import { RecommendationHandler } from '../service/recommendationHandler' import { runtimeLanguageContext } from '../util/runtimeLanguageContext' import { setContext } from '../../shared' @@ -75,6 +75,7 @@ export class AutotriggerState implements AnnotationState { static acceptedCount = 0 updateState(changeSource: AnnotationChangeSource, force: boolean): AnnotationState | undefined { + const session = CodeWhispererSessionState.instance.getSession() if (AutotriggerState.acceptedCount < RecommendationService.instance.acceptedSuggestionCount) { return new ManualtriggerState() } else if (session.recommendations.length > 0 && RecommendationHandler.instance.isSuggestionVisible()) { diff --git a/packages/core/src/test/codewhisperer/testUtil.ts b/packages/core/src/test/codewhisperer/testUtil.ts index f3b82fd3850..e7e1971a77d 100644 --- a/packages/core/src/test/codewhisperer/testUtil.ts +++ b/packages/core/src/test/codewhisperer/testUtil.ts @@ -16,7 +16,7 @@ import { MockDocument } from '../fake/fakeDocument' import { getLogger } from '../../shared/logger' import { CodeWhispererCodeCoverageTracker } from '../../codewhisperer/tracker/codewhispererCodeCoverageTracker' import globals from '../../shared/extensionGlobals' -import { session } from '../../codewhisperer/util/codeWhispererSession' +import { CodeWhispererSessionState } from '../../codewhisperer/util/codeWhispererSession' import { DefaultAWSClientBuilder, ServiceOptions } from '../../shared/awsClientBuilder' import { FakeAwsContext } from '../utilities/fakeAwsContext' import { HttpResponse, Service } from 'aws-sdk' @@ -33,6 +33,7 @@ export async function resetCodeWhispererGlobalVariables() { vsCodeState.isCodeWhispererEditing = false CodeWhispererCodeCoverageTracker.instances.clear() globals.telemetry.logger.clear() + const session = CodeWhispererSessionState.instance.getSession() session.reset() await globals.globalState.clear() await CodeSuggestionsState.instance.setSuggestionsEnabled(true) diff --git a/packages/core/src/testE2E/codewhisperer/referenceTracker.test.ts b/packages/core/src/testE2E/codewhisperer/referenceTracker.test.ts index 0038795ad89..6511c628eb9 100644 --- a/packages/core/src/testE2E/codewhisperer/referenceTracker.test.ts +++ b/packages/core/src/testE2E/codewhisperer/referenceTracker.test.ts @@ -10,7 +10,7 @@ import { setValidConnection, skipTestIfNoValidConn } from '../util/connection' import { RecommendationHandler } from '../../codewhisperer/service/recommendationHandler' import { createMockTextEditor, resetCodeWhispererGlobalVariables } from '../../test/codewhisperer/testUtil' import { invokeRecommendation } from '../../codewhisperer/commands/invokeRecommendation' -import { session } from '../../codewhisperer/util/codeWhispererSession' +import { CodeWhispererSessionState } from '../../codewhisperer/util/codeWhispererSession' /* New model deployment may impact references returned. @@ -47,6 +47,7 @@ describe('CodeWhisperer service invocation', async function () { isAutomatedTriggerEnabled: true, isSuggestionsWithCodeReferencesEnabled: false, } + const session = CodeWhispererSessionState.instance.getSession() before(async function () { validConnection = await setValidConnection() diff --git a/packages/core/src/testE2E/codewhisperer/serviceInvocations.test.ts b/packages/core/src/testE2E/codewhisperer/serviceInvocations.test.ts index d4265d13982..6770373489b 100644 --- a/packages/core/src/testE2E/codewhisperer/serviceInvocations.test.ts +++ b/packages/core/src/testE2E/codewhisperer/serviceInvocations.test.ts @@ -19,9 +19,10 @@ import { KeyStrokeHandler } from '../../codewhisperer/service/keyStrokeHandler' import { sleep } from '../../shared/utilities/timeoutUtils' import { invokeRecommendation } from '../../codewhisperer/commands/invokeRecommendation' import { getTestWorkspaceFolder } from '../../testInteg/integrationTestsUtilities' -import { session } from '../../codewhisperer/util/codeWhispererSession' +import { CodeWhispererSessionState } from '../../codewhisperer/util/codeWhispererSession' describe('CodeWhisperer service invocation', async function () { + const session = CodeWhispererSessionState.instance.getSession() let validConnection: boolean const client = new codewhispererClient.DefaultCodeWhispererClient() const config: ConfigurationEntry = {