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]>

Author: zixlin7

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."
+}

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 () {

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 () {

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()

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: '',

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()

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)

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 () {

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()

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,