perf(amazonq): nep UI improvement

aws-toolkit-vscode.code-workspace

@@ -12,6 +12,9 @@
         {
             "path": "packages/amazonq",
         },
+        {
+            "path": "../language-servers",
+        },
     ],
     "settings": {
         "typescript.tsdk": "node_modules/typescript/lib",

packages/amazonq/.vscode/launch.json

@@ -13,9 +13,9 @@
             "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_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",

packages/amazonq/src/app/inline/EditRendering/displayImage.ts

@@ -3,7 +3,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { getContext, getLogger, setContext } from 'aws-core-vscode/shared'
+import { getLogger, setContext } from 'aws-core-vscode/shared'
 import * as vscode from 'vscode'
 import { applyPatch, diffLines } from 'diff'
 import { LanguageClient } from 'vscode-languageclient'
@@ -16,7 +16,6 @@ import { EditSuggestionState } from '../editSuggestionState'
 import type { AmazonQInlineCompletionItemProvider } from '../completion'
 import { vsCodeState } from 'aws-core-vscode/codewhisperer'
 
-const autoRejectEditCursorDistance = 25
 const autoDiscardEditCursorDistance = 10
 
 export class EditDecorationManager {
@@ -311,6 +310,7 @@ export async function displaySvgDecoration(
     session: CodeWhispererSession,
     languageClient: LanguageClient,
     item: InlineCompletionItemWithReferences,
+    listeners: vscode.Disposable[],
     inlineCompletionProvider?: AmazonQInlineCompletionItemProvider
 ) {
     function logSuggestionFailure(type: 'DISCARD' | 'REJECT', reason: string, suggestionContent: string) {
@@ -359,44 +359,7 @@ export async function displaySvgDecoration(
         logSuggestionFailure('DISCARD', 'Invalid patch', item.insertText as string)
         return
     }
-    const documentChangeListener = vscode.workspace.onDidChangeTextDocument((e) => {
-        if (e.contentChanges.length <= 0) {
-            return
-        }
-        if (e.document !== editor.document) {
-            return
-        }
-        if (vsCodeState.isCodeWhispererEditing) {
-            return
-        }
-        if (getContext('aws.amazonq.editSuggestionActive') === false) {
-            return
-        }
 
-        const isPatchValid = applyPatch(e.document.getText(), item.insertText as string)
-        if (!isPatchValid) {
-            logSuggestionFailure('REJECT', 'Invalid patch due to document change', item.insertText as string)
-            void vscode.commands.executeCommand('aws.amazonq.inline.rejectEdit')
-        }
-    })
-    const cursorChangeListener = vscode.window.onDidChangeTextEditorSelection((e) => {
-        if (!EditSuggestionState.isEditSuggestionActive()) {
-            return
-        }
-        if (e.textEditor !== editor) {
-            return
-        }
-        const currentPosition = e.selections[0].active
-        const distance = Math.abs(currentPosition.line - startLine)
-        if (distance > autoRejectEditCursorDistance) {
-            logSuggestionFailure(
-                'REJECT',
-                `cursor position move too far away off ${autoRejectEditCursorDistance} lines`,
-                item.insertText as string
-            )
-            void vscode.commands.executeCommand('aws.amazonq.inline.rejectEdit')
-        }
-    })
     await decorationManager.displayEditSuggestion(
         editor,
         svgImage,
@@ -418,8 +381,12 @@ export async function displaySvgDecoration(
             editor.selection = new vscode.Selection(endPosition, endPosition)
 
             await decorationManager.clearDecorations(editor)
-            documentChangeListener.dispose()
-            cursorChangeListener.dispose()
+
+            // Dispose registered listeners on popup close
+            for (const listener of listeners) {
+                listener.dispose()
+            }
+
             const params: LogInlineCompletionSessionResultsParams = {
                 sessionId: session.sessionId,
                 completionSessionResult: {
@@ -444,8 +411,12 @@ export async function displaySvgDecoration(
                 getLogger().info('Edit suggestion rejected')
             }
             await decorationManager.clearDecorations(editor)
-            documentChangeListener.dispose()
-            cursorChangeListener.dispose()
+
+            // Dispose registered listeners on popup close
+            for (const listener of listeners) {
+                listener.dispose()
+            }
+
             const suggestionState = isDiscard
                 ? {
                       seen: false,

packages/amazonq/src/app/inline/EditRendering/imageRenderer.ts

@@ -6,54 +6,177 @@
 import * as vscode from 'vscode'
 import { displaySvgDecoration } from './displayImage'
 import { SvgGenerationService } from './svgGenerator'
-import { getLogger } from 'aws-core-vscode/shared'
+import { getContext, getLogger } from 'aws-core-vscode/shared'
 import { LanguageClient } from 'vscode-languageclient'
 import { InlineCompletionItemWithReferences } from '@aws/language-server-runtimes/protocol'
 import { CodeWhispererSession } from '../sessionManager'
 import type { AmazonQInlineCompletionItemProvider } from '../completion'
+import { vsCodeState } from 'aws-core-vscode/codewhisperer'
+import { applyPatch, createPatch } from 'diff'
+import { EditSuggestionState } from '../editSuggestionState'
 
-export async function showEdits(
-    item: InlineCompletionItemWithReferences,
-    editor: vscode.TextEditor | undefined,
-    session: CodeWhispererSession,
-    languageClient: LanguageClient,
-    inlineCompletionProvider?: AmazonQInlineCompletionItemProvider
-) {
-    if (!editor) {
-        return
+function logSuggestionFailure(type: 'DISCARD' | 'REJECT', reason: string, suggestionContent: string) {
+    getLogger('nextEditPrediction').debug(
+        `Auto ${type} edit suggestion with reason=${reason}, suggetion: ${suggestionContent}`
+    )
+}
+
+const autoRejectEditCursorDistance = 25
+const maxPrefixRetryCount = 5
+
+enum RejectReason {
+    DocumentChange = 'Invalid patch due to document change',
+    NotApplicableToOriginal = 'ApplyPatch fail for original code',
+    MaxRetry = 'Already retry 10 times',
+}
+
+export class EditsSuggestionSvg {
+    private readonly logger = getLogger('nextEditPrediction')
+    private readonly documentChangedListener: vscode.Disposable
+    private readonly cursorChangedListener: vscode.Disposable
+    private readonly updatedSuggestions: InlineCompletionItemWithReferences[] = []
+    private startLine = 0
+
+    constructor(
+        private suggestion: InlineCompletionItemWithReferences,
+        private readonly editor: vscode.TextEditor,
+        private readonly languageClient: LanguageClient,
+        private readonly session: CodeWhispererSession,
+        private readonly inlineCompletionProvider?: AmazonQInlineCompletionItemProvider // why nullable?
+    ) {
+        this.documentChangedListener = vscode.workspace.onDidChangeTextDocument(async (e) => {
+            await this.onDocChange(e)
+        })
+
+        this.cursorChangedListener = vscode.window.onDidChangeTextEditorSelection((e) => {
+            this.onCursorChange(e)
+        })
     }
-    try {
-        const svgGenerationService = new SvgGenerationService()
-        // Generate your SVG image with the file contents
-        const currentFile = editor.document.uri.fsPath
-        const { svgImage, startLine, newCode, originalCodeHighlightRange } = await svgGenerationService.generateDiffSvg(
-            currentFile,
-            item.insertText as string
-        )
 
-        // TODO: To investigate why it fails and patch [generateDiffSvg]
-        if (newCode.length === 0) {
-            getLogger('nextEditPrediction').warn('not able to apply provided edit suggestion, skip rendering')
+    async show() {
+        if (!this.editor) {
             return
         }
 
-        if (svgImage) {
-            // display the SVG image
-            await displaySvgDecoration(
-                editor,
-                svgImage,
-                startLine,
-                newCode,
-                originalCodeHighlightRange,
-                session,
-                languageClient,
-                item,
-                inlineCompletionProvider
+        const item =
+            this.updatedSuggestions.length > 0
+                ? this.updatedSuggestions[this.updatedSuggestions.length - 1]
+                : this.suggestion
+
+        try {
+            const svgGenerationService = new SvgGenerationService()
+            // Generate your SVG image with the file contents
+            const currentFile = this.editor.document.uri.fsPath
+            const { svgImage, startLine, newCode, originalCodeHighlightRange } =
+                await svgGenerationService.generateDiffSvg(currentFile, this.suggestion.insertText as string)
+
+            // For cursorChangeListener to access
+            this.startLine = startLine
+
+            // TODO: To investigate why it fails and patch [generateDiffSvg]
+            if (newCode.length === 0) {
+                this.logger.warn('not able to apply provided edit suggestion, skip rendering')
+                return
+            }
+
+            if (svgImage) {
+                // display the SVG image
+                await displaySvgDecoration(
+                    this.editor,
+                    svgImage,
+                    startLine,
+                    newCode,
+                    originalCodeHighlightRange,
+                    this.session,
+                    this.languageClient,
+                    item,
+                    [this.documentChangedListener, this.cursorChangedListener],
+                    this.inlineCompletionProvider
+                )
+            } else {
+                this.logger.error('SVG image generation returned an empty result.')
+            }
+        } catch (error) {
+            this.logger.error(`Error generating SVG image: ${error}`)
+        }
+    }
+
+    private onCursorChange(e: vscode.TextEditorSelectionChangeEvent) {
+        if (!EditSuggestionState.isEditSuggestionActive()) {
+            return
+        }
+        if (e.textEditor !== this.editor) {
+            return
+        }
+        const currentPosition = e.selections[0].active
+        const distance = Math.abs(currentPosition.line - this.startLine)
+        if (distance > autoRejectEditCursorDistance) {
+            logSuggestionFailure(
+                'REJECT',
+                `cursor position move too far away off ${autoRejectEditCursorDistance} lines`,
+                this.suggestion.insertText as string
             )
-        } else {
-            getLogger('nextEditPrediction').error('SVG image generation returned an empty result.')
+            void vscode.commands.executeCommand('aws.amazonq.inline.rejectEdit')
         }
-    } catch (error) {
-        getLogger('nextEditPrediction').error(`Error generating SVG image: ${error}`)
+    }
+
+    private async onDocChange(e: vscode.TextDocumentChangeEvent) {
+        if (e.contentChanges.length <= 0) {
+            return
+        }
+        if (e.document !== this.editor.document) {
+            return
+        }
+        if (vsCodeState.isCodeWhispererEditing) {
+            return
+        }
+        if (getContext('aws.amazonq.editSuggestionActive') === false) {
+            return
+        }
+
+        /**
+         * 1. Take the diff returned by the model and apply it to the code we originally sent to the model
+         * 2. Do a diff between the above code and what's currently in the editor
+         * 3. Show this second diff to the user as the edit suggestion
+         */
+        // Users' file content when the request fires (best guess because the actual process happens in language server)
+        const originalCode = this.session.fileContent
+        const appliedToOriginal = applyPatch(originalCode, this.suggestion.insertText as string)
+        try {
+            if (appliedToOriginal) {
+                const updatedPatch = this.patchSuggestion(appliedToOriginal)
+
+                if (this.updatedSuggestions.length > maxPrefixRetryCount) {
+                    this.autoReject(RejectReason.MaxRetry)
+                } else if (applyPatch(this.editor.document.getText(), updatedPatch) === false) {
+                    this.autoReject(RejectReason.DocumentChange)
+                }
+
+                await this.show()
+            } else {
+                this.autoReject(RejectReason.NotApplicableToOriginal)
+            }
+        } catch (e) {
+            // TODO: format
+            this.logger.error(`${e}`)
+        }
+    }
+
+    private autoReject(reason: string) {
+        logSuggestionFailure('REJECT', reason, this.suggestion.insertText as string)
+        void vscode.commands.executeCommand('aws.amazonq.inline.rejectEdit')
+    }
+
+    private patchSuggestion(appliedToOriginal: string): string {
+        const updatedPatch = createPatch(
+            this.editor.document.fileName,
+            this.editor.document.getText(),
+            appliedToOriginal
+        )
+
+        this.logger.info(`Update edit suggestion\n ${updatedPatch}`)
+        const updated: InlineCompletionItemWithReferences = { ...this.suggestion, insertText: updatedPatch }
+        this.updatedSuggestions.push(updated)
+        return updatedPatch
     }
 }

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

@@ -43,7 +43,7 @@ import { InlineTutorialAnnotation } from './tutorials/inlineTutorialAnnotation'
 import { TelemetryHelper } from './telemetryHelper'
 import { Experiments, getLogger, sleep } from 'aws-core-vscode/shared'
 import { messageUtils } from 'aws-core-vscode/utils'
-import { showEdits } from './EditRendering/imageRenderer'
+import { EditsSuggestionSvg } from './EditRendering/imageRenderer'
 import { ICursorUpdateRecorder } from './cursorUpdateManager'
 import { DocumentEventListener } from './documentEventListener'
 
@@ -529,7 +529,7 @@ ${itemLog}
                 if (item.isInlineEdit) {
                     // Check if Next Edit Prediction feature flag is enabled
                     if (Experiments.instance.get('amazonqLSPNEP', true)) {
-                        await showEdits(item, editor, session, this.languageClient, this)
+                        await new EditsSuggestionSvg(item, editor, this.languageClient, session, this).show()
                         logstr += `- duration between trigger to edits suggestion is displayed: ${Date.now() - t0}ms`
                     }
                     return []

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

@@ -141,12 +141,12 @@ export class RecommendationService {
              * Edits leverage partialResultToken to achieve EditStreak such that clients can pull all continuous suggestions generated by the model within 1 EOS block.
              */
             if (!isTriggerByDeletion && !request.partialResultToken && !EditSuggestionState.isEditSuggestionActive()) {
-                const completionPromise: Promise<InlineCompletionListWithReferences> = languageClient.sendRequest(
-                    inlineCompletionWithReferencesRequestType.method,
-                    request,
-                    token
-                )
-                ps.push(completionPromise)
+                // const completionPromise: Promise<InlineCompletionListWithReferences> = languageClient.sendRequest(
+                //     inlineCompletionWithReferencesRequestType.method,
+                //     request,
+                //     token
+                // )
+                // ps.push(completionPromise)
             }
 
             /**
@@ -241,6 +241,7 @@ export class RecommendationService {
                 result.items,
                 requestStartTime,
                 position,
+                document,
                 firstCompletionDisplayLatency
             )