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

package.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",

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

@@ -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/lsp/client.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()
         }),