From da7c7168703fd95b1bf52197acb8789aa57aecf1 Mon Sep 17 00:00:00 2001
From: Will Lo <xshaohua@amazon.com>
Date: Thu, 3 Jul 2025 02:43:43 -0700
Subject: [PATCH] fix(amazonq): nep typeahead

---
 aws-toolkit-vscode.code-workspace             |  3 ++
 packages/amazonq/.vscode/launch.json          |  6 +--
 packages/amazonq/src/app/inline/completion.ts | 43 ++++++++++++-------
 3 files changed, 33 insertions(+), 19 deletions(-)

diff --git a/aws-toolkit-vscode.code-workspace b/aws-toolkit-vscode.code-workspace
index f03aafae2fe..479f9e8fd66 100644
--- a/aws-toolkit-vscode.code-workspace
+++ b/aws-toolkit-vscode.code-workspace
@@ -12,6 +12,9 @@
         {
             "path": "packages/amazonq",
         },
+        {
+            "path": "../language-servers",
+        },
     ],
     "settings": {
         "typescript.tsdk": "node_modules/typescript/lib",
diff --git a/packages/amazonq/.vscode/launch.json b/packages/amazonq/.vscode/launch.json
index b00c5071ce5..cdeabe152a9 100644
--- a/packages/amazonq/.vscode/launch.json
+++ b/packages/amazonq/.vscode/launch.json
@@ -13,10 +13,10 @@
             "args": ["--extensionDevelopmentPath=${workspaceFolder}"],
             "env": {
                 "SSMDOCUMENT_LANGUAGESERVER_PORT": "6010",
-                "WEBPACK_DEVELOPER_SERVER": "http://localhost:8080"
+                "WEBPACK_DEVELOPER_SERVER": "http://localhost:8080",
                 // Below allows for overrides used during development
-                // "__AMAZONQLSP_PATH": "${workspaceFolder}/../../../language-servers/app/aws-lsp-codewhisperer-runtimes/out/agent-standalone.js",
-                // "__AMAZONQLSP_UI": "${workspaceFolder}/../../../language-servers/chat-client/build/amazonq-ui.js"
+                "__AMAZONQLSP_PATH": "${workspaceFolder}/../../../language-servers/app/aws-lsp-codewhisperer-runtimes/out/agent-standalone.js",
+                "__AMAZONQLSP_UI": "${workspaceFolder}/../../../language-servers/chat-client/build/amazonq-ui.js"
             },
             "envFile": "${workspaceFolder}/.local.env",
             "outFiles": ["${workspaceFolder}/dist/**/*.js", "${workspaceFolder}/../core/dist/**/*.js"],
diff --git a/packages/amazonq/src/app/inline/completion.ts b/packages/amazonq/src/app/inline/completion.ts
index 069a6bc5128..efd5e1cbec6 100644
--- a/packages/amazonq/src/app/inline/completion.ts
+++ b/packages/amazonq/src/app/inline/completion.ts
@@ -239,30 +239,41 @@ export class AmazonQInlineCompletionItemProvider implements InlineCompletionItem
             const editor = window.activeTextEditor
             if (prevSession && prevSessionId && prevItemId && prevStartPosition) {
                 const prefix = document.getText(new Range(prevStartPosition, position))
-                const prevItemMatchingPrefix = []
+                const prevItemMatchingPrefix: InlineCompletionItemWithReferences[] = []
                 for (const item of this.sessionManager.getActiveRecommendation()) {
                     const text = typeof item.insertText === 'string' ? item.insertText : item.insertText.value
-                    if (text.startsWith(prefix) && position.isAfterOrEqual(prevStartPosition)) {
-                        item.command = {
-                            command: 'aws.amazonq.acceptInline',
-                            title: 'On acceptance',
-                            arguments: [
-                                prevSessionId,
-                                item,
-                                editor,
-                                prevSession?.requestStartTime,
-                                position.line,
-                                prevSession?.firstCompletionDisplayLatency,
-                            ],
+                    if (prefix.length > 0 && text.startsWith(prefix) && position.isAfterOrEqual(prevStartPosition)) {
+                        if (item.isInlineEdit) {
+                            prevItemMatchingPrefix.push(item)
+                        } else {
+                            item.command = {
+                                command: 'aws.amazonq.acceptInline',
+                                title: 'On acceptance',
+                                arguments: [
+                                    prevSessionId,
+                                    item,
+                                    editor,
+                                    prevSession?.requestStartTime,
+                                    position.line,
+                                    prevSession?.firstCompletionDisplayLatency,
+                                ],
+                            }
+                            item.range = new Range(prevStartPosition, position)
+                            prevItemMatchingPrefix.push(item)
                         }
-                        item.range = new Range(prevStartPosition, position)
-                        prevItemMatchingPrefix.push(item as InlineCompletionItem)
                     }
                 }
                 // re-use previous suggestions as long as new typed prefix matches
                 if (prevItemMatchingPrefix.length > 0) {
+                    if (prevItemMatchingPrefix[0].isInlineEdit) {
+                        // Check if Next Edit Prediction feature flag is enabled
+                        if (Experiments.instance.isExperimentEnabled('amazonqLSPNEP')) {
+                            void showEdits(prevItemMatchingPrefix[0], editor, prevSession, this.languageClient).then()
+                        }
+                        return []
+                    }
                     getLogger().debug(`Re-using suggestions that match user typed characters`)
-                    return prevItemMatchingPrefix
+                    return prevItemMatchingPrefix as InlineCompletionItem[]
                 }
                 getLogger().debug(`Auto rejecting suggestions from previous session`)
                 // if no such suggestions, report the previous suggestion as Reject