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 <[email protected]>
aws · Feb 4, 2025 · 50c9121 · 50c9121
1 parent 9afc587
commit 50c9121
Show file tree

Hide file tree

Showing 23 changed files with 292 additions and 116 deletions.
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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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<CodeWhispererUserClient> {
+        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
@@ -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
@@ -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
@@ -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'