feat(amazonq): start AB test for inline completion rollback (aws#8006)

## Problem


## Solution


---

- 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: atontb <[email protected]>
Co-authored-by: aws-toolkit-automation <[email protected]>
Co-authored-by: andrewyuq <[email protected]>
Co-authored-by: Will Lo <[email protected]>

Author: 5 people

package-lock.json

@@ -74,7 +74,7 @@
         "webpack-merge": "^5.10.0"
     },
     "dependencies": {
-        "@aws/language-server-runtimes": "^0.2.125",
+        "@aws/language-server-runtimes": "^0.2.128",
         "@types/node": "^22.7.5",
         "jaro-winkler": "^0.2.8",
         "vscode-nls": "^5.2.0",

package.json

@@ -5,29 +5,73 @@
 
 import vscode from 'vscode'
 import {
+    acceptSuggestion,
     AuthUtil,
+    CodeSuggestionsState,
+    CodeWhispererCodeCoverageTracker,
     CodeWhispererConstants,
+    CodeWhispererSettings,
+    ConfigurationEntry,
+    DefaultCodeWhispererClient,
+    invokeRecommendation,
     isInlineCompletionEnabled,
+    KeyStrokeHandler,
+    RecommendationHandler,
     runtimeLanguageContext,
+    TelemetryHelper,
     UserWrittenCodeTracker,
     vsCodeState,
 } from 'aws-core-vscode/codewhisperer'
-import { globals, sleep } from 'aws-core-vscode/shared'
+import { Commands, getLogger, globals, sleep } from 'aws-core-vscode/shared'
+import { LanguageClient } from 'vscode-languageclient'
 
-export async function activate() {
-    if (isInlineCompletionEnabled()) {
-        // Debugging purpose: only initialize NextEditPredictionPanel when development
-        // NextEditPredictionPanel.getInstance()
+export async function activate(languageClient: LanguageClient) {
+    const codewhispererSettings = CodeWhispererSettings.instance
+    const client = new DefaultCodeWhispererClient()
 
+    if (isInlineCompletionEnabled()) {
         await setSubscriptionsforInlineCompletion()
         await AuthUtil.instance.setVscodeContextProps()
+        RecommendationHandler.instance.setLanguageClient(languageClient)
+    }
+
+    function getAutoTriggerStatus(): boolean {
+        return CodeSuggestionsState.instance.isSuggestionsEnabled()
+    }
+
+    async function getConfigEntry(): Promise<ConfigurationEntry> {
+        const isShowMethodsEnabled: boolean =
+            vscode.workspace.getConfiguration('editor').get('suggest.showMethods') || false
+        const isAutomatedTriggerEnabled: boolean = getAutoTriggerStatus()
+        const isManualTriggerEnabled: boolean = true
+        const isSuggestionsWithCodeReferencesEnabled = codewhispererSettings.isSuggestionsWithCodeReferencesEnabled()
+
+        // TODO:remove isManualTriggerEnabled
+        return {
+            isShowMethodsEnabled,
+            isManualTriggerEnabled,
+            isAutomatedTriggerEnabled,
+            isSuggestionsWithCodeReferencesEnabled,
+        }
     }
 
     async function setSubscriptionsforInlineCompletion() {
+        RecommendationHandler.instance.subscribeSuggestionCommands()
+
         /**
          * Automated trigger
          */
         globals.context.subscriptions.push(
+            acceptSuggestion.register(globals.context),
+            vscode.window.onDidChangeActiveTextEditor(async (editor) => {
+                await RecommendationHandler.instance.onEditorChange()
+            }),
+            vscode.window.onDidChangeWindowState(async (e) => {
+                await RecommendationHandler.instance.onFocusChange()
+            }),
+            vscode.window.onDidChangeTextEditorSelection(async (e) => {
+                await RecommendationHandler.instance.onCursorChange(e)
+            }),
             vscode.workspace.onDidChangeTextDocument(async (e) => {
                 const editor = vscode.window.activeTextEditor
                 if (!editor) {
@@ -40,6 +84,7 @@ export async function activate() {
                     return
                 }
 
+                CodeWhispererCodeCoverageTracker.getTracker(e.document.languageId)?.countTotalTokens(e)
                 UserWrittenCodeTracker.instance.onTextDocumentChange(e)
                 /**
                  * Handle this keystroke event only when
@@ -51,6 +96,11 @@ export async function activate() {
                     return
                 }
 
+                if (vsCodeState.lastUserModificationTime) {
+                    TelemetryHelper.instance.setTimeSinceLastModification(
+                        performance.now() - vsCodeState.lastUserModificationTime
+                    )
+                }
                 vsCodeState.lastUserModificationTime = performance.now()
                 /**
                  * Important:  Doing this sleep(10) is to make sure
@@ -59,6 +109,19 @@ export async function activate() {
                  * Then this event can be processed by our code.
                  */
                 await sleep(CodeWhispererConstants.vsCodeCursorUpdateDelay)
+                if (!RecommendationHandler.instance.isSuggestionVisible()) {
+                    await KeyStrokeHandler.instance.processKeyStroke(e, editor, client, await getConfigEntry())
+                }
+            }),
+            // manual trigger
+            Commands.register({ id: 'aws.amazonq.invokeInlineCompletion', autoconnect: true }, async () => {
+                invokeRecommendation(
+                    vscode.window.activeTextEditor as vscode.TextEditor,
+                    client,
+                    await getConfigEntry()
+                ).catch((e: Error) => {
+                    getLogger().error('invokeRecommendation failed: %s', (e as Error).message)
+                })
             })
         )
     }

packages/amazonq/src/app/inline/activation.ts

@@ -21,6 +21,7 @@ import {
 import {
     AuthUtil,
     CodeWhispererSettings,
+    FeatureConfigProvider,
     getSelectedCustomization,
     TelemetryHelper,
     vsCodeState,
@@ -45,6 +46,7 @@ import {
 } from 'aws-core-vscode/shared'
 import { processUtils } from 'aws-core-vscode/shared'
 import { activate } from './chat/activation'
+import { activate as activateInline } from '../app/inline/activation'
 import { AmazonQResourcePaths } from './lspInstaller'
 import { ConfigSection, isValidConfigSection, pushConfigUpdate, toAmazonQLSPLogLevel } from './config'
 import { activate as activateInlineChat } from '../inlineChat/activation'
@@ -338,8 +340,42 @@ async function onLanguageServerReady(
     // tutorial for inline chat
     const inlineChatTutorialAnnotation = new InlineChatTutorialAnnotation(inlineTutorialAnnotation)
 
-    const inlineManager = new InlineCompletionManager(client, sessionManager, lineTracker, inlineTutorialAnnotation)
-    inlineManager.registerInlineCompletion()
+    const enableInlineRollback = FeatureConfigProvider.instance.getPreFlareRollbackGroup() === 'treatment'
+    if (enableInlineRollback) {
+        // use VSC inline
+        getLogger().info('Entering preflare logic')
+        await activateInline(client)
+    } else {
+        // use language server for inline completion
+        getLogger().info('Entering postflare logic')
+        const inlineManager = new InlineCompletionManager(client, sessionManager, lineTracker, inlineTutorialAnnotation)
+        inlineManager.registerInlineCompletion()
+        toDispose.push(
+            inlineManager,
+            Commands.register('aws.amazonq.showPrev', async () => {
+                await sessionManager.maybeRefreshSessionUx()
+                await vscode.commands.executeCommand('editor.action.inlineSuggest.showPrevious')
+                sessionManager.onPrevSuggestion()
+            }),
+            Commands.register('aws.amazonq.showNext', async () => {
+                await sessionManager.maybeRefreshSessionUx()
+                await vscode.commands.executeCommand('editor.action.inlineSuggest.showNext')
+                sessionManager.onNextSuggestion()
+            }),
+            // this is a workaround since handleDidShowCompletionItem is not public API
+            Commands.register('aws.amazonq.checkInlineSuggestionVisibility', async () => {
+                sessionManager.checkInlineSuggestionVisibility()
+            }),
+            Commands.register({ id: 'aws.amazonq.invokeInlineCompletion', autoconnect: true }, async () => {
+                vsCodeState.lastManualTriggerTime = performance.now()
+                await vscode.commands.executeCommand('editor.action.inlineSuggest.trigger')
+            }),
+            vscode.workspace.onDidCloseTextDocument(async () => {
+                await vscode.commands.executeCommand('aws.amazonq.rejectCodeSuggestion')
+            })
+        )
+    }
+
     activateInlineChat(extensionContext, client, encryptionKey, inlineChatTutorialAnnotation)
 
     if (Experiments.instance.get('amazonqChatLSP', true)) {
@@ -354,25 +390,6 @@ async function onLanguageServerReady(
     await initializeLanguageServerConfiguration(client, 'startup')
 
     toDispose.push(
-        inlineManager,
-        Commands.register('aws.amazonq.showPrev', async () => {
-            await sessionManager.maybeRefreshSessionUx()
-            await vscode.commands.executeCommand('editor.action.inlineSuggest.showPrevious')
-            sessionManager.onPrevSuggestion()
-        }),
-        Commands.register('aws.amazonq.showNext', async () => {
-            await sessionManager.maybeRefreshSessionUx()
-            await vscode.commands.executeCommand('editor.action.inlineSuggest.showNext')
-            sessionManager.onNextSuggestion()
-        }),
-        // this is a workaround since handleDidShowCompletionItem is not public API
-        Commands.register('aws.amazonq.checkInlineSuggestionVisibility', async () => {
-            sessionManager.checkInlineSuggestionVisibility()
-        }),
-        Commands.register({ id: 'aws.amazonq.invokeInlineCompletion', autoconnect: true }, async () => {
-            vsCodeState.lastManualTriggerTime = performance.now()
-            await vscode.commands.executeCommand('editor.action.inlineSuggest.trigger')
-        }),
         Commands.register('aws.amazonq.refreshAnnotation', async (forceProceed: boolean) => {
             telemetry.record({
                 traceId: TelemetryHelper.instance.traceId,
@@ -398,9 +415,6 @@ async function onLanguageServerReady(
                 getLogger().debug(`codewhisperer: user dismiss tutorial.`)
             }
         }),
-        vscode.workspace.onDidCloseTextDocument(async () => {
-            await vscode.commands.executeCommand('aws.amazonq.rejectCodeSuggestion')
-        }),
         AuthUtil.instance.auth.onDidChangeActiveConnection(async () => {
             await auth.refreshConnection()
         }),

packages/amazonq/src/lsp/client.ts

@@ -0,0 +1,43 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import assert from 'assert'
+import * as sinon from 'sinon'
+import { resetCodeWhispererGlobalVariables, createMockTextEditor } from 'aws-core-vscode/test'
+import {
+    ConfigurationEntry,
+    invokeRecommendation,
+    InlineCompletionService,
+    isInlineCompletionEnabled,
+    DefaultCodeWhispererClient,
+} from 'aws-core-vscode/codewhisperer'
+
+describe('invokeRecommendation', function () {
+    describe('invokeRecommendation', function () {
+        let getRecommendationStub: sinon.SinonStub
+        let mockClient: DefaultCodeWhispererClient
+
+        beforeEach(async function () {
+            await resetCodeWhispererGlobalVariables()
+            getRecommendationStub = sinon.stub(InlineCompletionService.instance, 'getPaginatedRecommendation')
+        })
+
+        afterEach(function () {
+            sinon.restore()
+        })
+
+        it('Should call getPaginatedRecommendation with OnDemand as trigger type when inline completion is enabled', async function () {
+            const mockEditor = createMockTextEditor()
+            const config: ConfigurationEntry = {
+                isShowMethodsEnabled: true,
+                isManualTriggerEnabled: true,
+                isAutomatedTriggerEnabled: true,
+                isSuggestionsWithCodeReferencesEnabled: true,
+            }
+            await invokeRecommendation(mockEditor, mockClient, config)
+            assert.strictEqual(getRecommendationStub.called, isInlineCompletionEnabled())
+        })
+    })
+})

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

@@ -0,0 +1,64 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import assert from 'assert'
+import * as vscode from 'vscode'
+import * as sinon from 'sinon'
+import { onAcceptance, AcceptedSuggestionEntry, session, CodeWhispererTracker } from 'aws-core-vscode/codewhisperer'
+import { resetCodeWhispererGlobalVariables, createMockTextEditor } from 'aws-core-vscode/test'
+
+describe('onAcceptance', function () {
+    describe('onAcceptance', function () {
+        beforeEach(async function () {
+            await resetCodeWhispererGlobalVariables()
+            session.reset()
+        })
+
+        afterEach(function () {
+            sinon.restore()
+            session.reset()
+        })
+
+        it('Should enqueue an event object to tracker', async function () {
+            const mockEditor = createMockTextEditor()
+            const trackerSpy = sinon.spy(CodeWhispererTracker.prototype, 'enqueue')
+            const fakeReferences = [
+                {
+                    message: '',
+                    licenseName: 'MIT',
+                    repository: 'http://github.com/fake',
+                    recommendationContentSpan: {
+                        start: 0,
+                        end: 10,
+                    },
+                },
+            ]
+            await onAcceptance({
+                editor: mockEditor,
+                range: new vscode.Range(new vscode.Position(1, 0), new vscode.Position(1, 26)),
+                effectiveRange: new vscode.Range(new vscode.Position(1, 0), new vscode.Position(1, 26)),
+                acceptIndex: 0,
+                recommendation: "print('Hello World!')",
+                requestId: '',
+                sessionId: '',
+                triggerType: 'OnDemand',
+                completionType: 'Line',
+                language: 'python',
+                references: fakeReferences,
+            })
+            const actualArg = trackerSpy.getCall(0).args[0] as AcceptedSuggestionEntry
+            assert.ok(trackerSpy.calledOnce)
+            assert.strictEqual(actualArg.originalString, 'def two_sum(nums, target):')
+            assert.strictEqual(actualArg.requestId, '')
+            assert.strictEqual(actualArg.sessionId, '')
+            assert.strictEqual(actualArg.triggerType, 'OnDemand')
+            assert.strictEqual(actualArg.completionType, 'Line')
+            assert.strictEqual(actualArg.language, 'python')
+            assert.deepStrictEqual(actualArg.startPosition, new vscode.Position(1, 0))
+            assert.deepStrictEqual(actualArg.endPosition, new vscode.Position(1, 26))
+            assert.strictEqual(actualArg.index, 0)
+        })
+    })
+})

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

@@ -0,0 +1,43 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import assert from 'assert'
+import * as vscode from 'vscode'
+import * as sinon from 'sinon'
+import { resetCodeWhispererGlobalVariables, createMockTextEditor } from 'aws-core-vscode/test'
+import { onInlineAcceptance, RecommendationHandler, session } from 'aws-core-vscode/codewhisperer'
+
+describe('onInlineAcceptance', function () {
+    describe('onInlineAcceptance', function () {
+        beforeEach(async function () {
+            await resetCodeWhispererGlobalVariables()
+            session.reset()
+        })
+
+        afterEach(function () {
+            sinon.restore()
+            session.reset()
+        })
+
+        it('Should dispose inline completion provider', async function () {
+            const mockEditor = createMockTextEditor()
+            const spy = sinon.spy(RecommendationHandler.instance, 'disposeInlineCompletion')
+            await onInlineAcceptance({
+                editor: mockEditor,
+                range: new vscode.Range(new vscode.Position(1, 0), new vscode.Position(1, 21)),
+                effectiveRange: new vscode.Range(new vscode.Position(1, 0), new vscode.Position(1, 21)),
+                acceptIndex: 0,
+                recommendation: "print('Hello World!')",
+                requestId: '',
+                sessionId: '',
+                triggerType: 'OnDemand',
+                completionType: 'Line',
+                language: 'python',
+                references: undefined,
+            })
+            assert.ok(spy.calledWith())
+        })
+    })
+})