poc model routing

Author: zixlin7

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,
@@ -18,6 +18,7 @@ import { resetCodeWhispererGlobalVariables, createMockTextEditor } from 'aws-cor
 import { assertTelemetryCurried } from 'aws-core-vscode/test'
 
 describe('onAcceptance', function () {
+    const session = CodeWhispererSessionState.instance.getSession()
     describe('onAcceptance', function () {
         beforeEach(async function () {
             await resetCodeWhispererGlobalVariables()

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

@@ -8,11 +8,17 @@ 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 () {
+    const session = CodeWhispererSessionState.instance.getSession()
     describe('onInlineAcceptance', function () {
         beforeEach(async function () {
             await resetCodeWhispererGlobalVariables()

packages/amazonq/test/unit/codewhisperer/service/completionProvider.test.ts

@@ -12,11 +12,12 @@ import {
     getLabel,
     Recommendation,
     RecommendationHandler,
-    session,
+    CodeWhispererSessionState,
 } from 'aws-core-vscode/codewhisperer'
 import { createMockDocument, resetCodeWhispererGlobalVariables } from 'aws-core-vscode/test'
 
 describe('completionProviderService', function () {
+    const session = CodeWhispererSessionState.instance.getSession()
     beforeEach(async function () {
         await resetCodeWhispererGlobalVariables()
     })

packages/amazonq/test/unit/codewhisperer/service/inlineCompletionService.test.ts

@@ -14,14 +14,15 @@ import {
     CodeSuggestionsState,
     ConfigurationEntry,
     CWInlineCompletionItemProvider,
-    session,
+    CodeWhispererSessionState,
     AuthUtil,
     listCodeWhispererCommandsId,
     DefaultCodeWhispererClient,
 } from 'aws-core-vscode/codewhisperer'
 import { createMockTextEditor, resetCodeWhispererGlobalVariables, createMockDocument } from 'aws-core-vscode/test'
 
 describe('inlineCompletionService', function () {
+    const session = CodeWhispererSessionState.instance.getSession()
     beforeEach(async function () {
         await resetCodeWhispererGlobalVariables()
     })

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,
@@ -26,6 +26,7 @@ import {
 // import * as supplementalContextUtil from 'aws-core-vscode/codewhisperer'
 
 describe('recommendationHandler', function () {
+    const session = CodeWhispererSessionState.instance.getSession()
     const config: ConfigurationEntry = {
         isShowMethodsEnabled: true,
         isManualTriggerEnabled: true,

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
 

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,12 @@ export async function onInlineAcceptance(acceptanceEntry: OnRecommendationAccept
         }
 
         RecommendationHandler.instance.reportUserDecisions(acceptanceEntry.acceptIndex)
+        if (acceptanceEntry.acceptIndex == 0) {
+            // TODO: gate behind A/B group if needed
+            const nextSession = CodeWhispererSessionState.instance.getNextSession()
+            nextSession.startPos = acceptanceEntry.editor.selection.active
+            CodeWhispererSessionState.instance.setSession(nextSession)
+            await RecommendationHandler.instance.showNextRecommendations()
+        }
     }
 }