feat(amazonq): pre-fetch next recommendations for inline completions

packages/amazonq/test/unit/codewhisperer/commands/onAcceptance.test.ts

@@ -9,7 +9,7 @@ import * as sinon from 'sinon'
 import {
     onAcceptance,
     AcceptedSuggestionEntry,
-    session,
+    CodeWhispererSessionState,
     CodeWhispererTracker,
     RecommendationHandler,
     AuthUtil,
@@ -20,11 +20,13 @@ import { assertTelemetryCurried } from 'aws-core-vscode/test'
 describe('onAcceptance', function () {
     describe('onAcceptance', function () {
         beforeEach(async function () {
+            const session = CodeWhispererSessionState.instance.getSession()
             await resetCodeWhispererGlobalVariables()
             session.reset()
         })
 
         afterEach(function () {
+            const session = CodeWhispererSessionState.instance.getSession()
             sinon.restore()
             session.reset()
         })
@@ -70,6 +72,7 @@ describe('onAcceptance', function () {
         })
 
         it('Should report telemetry that records this user decision event', async function () {
+            const session = CodeWhispererSessionState.instance.getSession()
             const testStartUrl = 'testStartUrl'
             sinon.stub(AuthUtil.instance, 'startUrl').value(testStartUrl)
             const mockEditor = createMockTextEditor()

packages/amazonq/test/unit/codewhisperer/commands/onInlineAcceptance.test.ts

@@ -8,18 +8,25 @@ 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,
+} from 'aws-core-vscode/codewhisperer'
 import { globals } from 'aws-core-vscode/shared'
 import { extensionVersion } from 'aws-core-vscode/shared'
 
 describe('onInlineAcceptance', function () {
     describe('onInlineAcceptance', function () {
         beforeEach(async function () {
+            const session = CodeWhispererSessionState.instance.getSession()
             await resetCodeWhispererGlobalVariables()
             session.reset()
         })
 
         afterEach(function () {
+            const session = CodeWhispererSessionState.instance.getSession()
             sinon.restore()
             session.reset()
         })
@@ -44,6 +51,7 @@ describe('onInlineAcceptance', function () {
         })
 
         it('Should report telemetry that records this user decision event', async function () {
+            const session = CodeWhispererSessionState.instance.getSession()
             await globals.globalState.update('CODEWHISPERER_USER_GROUP', {
                 version: extensionVersion,
             })

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,
@@ -51,6 +51,7 @@ describe('telemetryHelper', function () {
         })
 
         it('resetClientComponentLatencyTime should reset state variables', function () {
+            const session = CodeWhispererSessionState.instance.getSession()
             session.invokeSuggestionStartTime = 100
             session.preprocessEndTime = 200
             session.sdkApiCallStartTime = 300
@@ -289,6 +290,7 @@ describe('telemetryHelper', function () {
         })
 
         it('Should call telemetry record for each recommendations with proper arguments', async function () {
+            const session = CodeWhispererSessionState.instance.getSession()
             const telemetryHelper = new TelemetryHelper()
             const response = [{ content: "print('Hello')" }]
             const requestIdList = ['test_x', 'test_x', 'test_y']

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,

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

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 } from './util/codeWhispererSession'
 export { onInlineAcceptance } from './commands/onInlineAcceptance'
 export { stopTransformByQ } from './commands/startTransformByQ'
 export { getCompletionItems, getCompletionItem, getLabel } from './service/completionProvider'