perf(amazonq): nep UI improvement

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 { BaseLanguageClient } 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 {
@@ -164,7 +163,10 @@ export class EditDecorationManager {
     /**
      * Clears all edit suggestion decorations
      */
-    public async clearDecorations(editor: vscode.TextEditor): Promise<void> {
+    public async clearDecorations(editor: vscode.TextEditor, disposables: vscode.Disposable[]): Promise<void> {
+        for (const d of disposables) {
+            d.dispose()
+        }
         editor.setDecorations(this.imageDecorationType, [])
         editor.setDecorations(this.removedCodeDecorationType, [])
         this.currentImageDecoration = undefined
@@ -311,6 +313,7 @@ export async function displaySvgDecoration(
     session: CodeWhispererSession,
     languageClient: BaseLanguageClient,
     item: InlineCompletionItemWithReferences,
+    listeners: vscode.Disposable[],
     inlineCompletionProvider?: AmazonQInlineCompletionItemProvider
 ) {
     function logSuggestionFailure(type: 'DISCARD' | 'REJECT', reason: string, suggestionContent: string) {
@@ -359,44 +362,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,
@@ -417,9 +383,8 @@ export async function displaySvgDecoration(
             const endPosition = getEndOfEditPosition(originalCode, newCode)
             editor.selection = new vscode.Selection(endPosition, endPosition)
 
-            await decorationManager.clearDecorations(editor)
-            documentChangeListener.dispose()
-            cursorChangeListener.dispose()
+            await decorationManager.clearDecorations(editor, listeners)
+
             const params: LogInlineCompletionSessionResultsParams = {
                 sessionId: session.sessionId,
                 completionSessionResult: {
@@ -443,9 +408,8 @@ export async function displaySvgDecoration(
             } else {
                 getLogger().info('Edit suggestion rejected')
             }
-            await decorationManager.clearDecorations(editor)
-            documentChangeListener.dispose()
-            cursorChangeListener.dispose()
+            await decorationManager.clearDecorations(editor, listeners)
+
             const suggestionState = isDiscard
                 ? {
                       seen: false,

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

@@ -4,56 +4,207 @@
  */
 
 import * as vscode from 'vscode'
-import { displaySvgDecoration } from './displayImage'
+import { displaySvgDecoration, decorationManager } from './displayImage'
 import { SvgGenerationService } from './svgGenerator'
-import { getLogger } from 'aws-core-vscode/shared'
+import { getContext, getLogger } from 'aws-core-vscode/shared'
 import { BaseLanguageClient } 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'
+import { debounce } from 'aws-core-vscode/utils'
 
-export async function showEdits(
-    item: InlineCompletionItemWithReferences,
-    editor: vscode.TextEditor | undefined,
-    session: CodeWhispererSession,
-    languageClient: BaseLanguageClient,
-    inlineCompletionProvider?: AmazonQInlineCompletionItemProvider
-) {
-    if (!editor) {
-        return
-    }
-    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
-        )
+const autoRejectEditCursorDistance = 25
+const maxPrefixRetryCharDiff = 5
+const rerenderDeboucneInMs = 500
+
+enum RejectReason {
+    DocumentChange = 'Invalid patch due to document change',
+    NotApplicableToOriginal = 'ApplyPatch fail for original code',
+    MaxRetry = `Already retry ${maxPrefixRetryCharDiff} times`,
+}
+
+export class EditsSuggestionSvg {
+    private readonly logger = getLogger('nextEditPrediction')
+    private documentChangedListener: vscode.Disposable | undefined
+    private cursorChangedListener: vscode.Disposable | undefined
+
+    private startLine = 0
+
+    private documentChangeTrace = {
+        contentChanged: '',
+        count: 0,
+    }
+
+    constructor(
+        private suggestion: InlineCompletionItemWithReferences,
+        private readonly editor: vscode.TextEditor,
+        private readonly languageClient: BaseLanguageClient,
+        private readonly session: CodeWhispererSession,
+        private readonly inlineCompletionProvider?: AmazonQInlineCompletionItemProvider
+    ) {}
+
+    async show(patchedSuggestion?: InlineCompletionItemWithReferences) {
+        if (!this.editor) {
+            this.logger.error(`attempting to render an edit suggestion while editor is undefined`)
+            return
+        }
+
+        const item = patchedSuggestion ? patchedSuggestion : this.suggestion
 
-        // 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')
+        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
+
+            if (newCode.length === 0) {
+                this.logger.warn('not able to apply provided edit suggestion, skip rendering')
+                return
+            }
+
+            if (svgImage) {
+                const documentChangedListener = (this.documentChangedListener ??=
+                    vscode.workspace.onDidChangeTextDocument(async (e) => {
+                        await this.onDocChange(e)
+                    }))
+
+                const cursorChangedListener = (this.cursorChangedListener ??=
+                    vscode.window.onDidChangeTextEditorSelection((e) => {
+                        this.onCursorChange(e)
+                    }))
+
+                // display the SVG image
+                await displaySvgDecoration(
+                    this.editor,
+                    svgImage,
+                    startLine,
+                    newCode,
+                    originalCodeHighlightRange,
+                    this.session,
+                    this.languageClient,
+                    item,
+                    [documentChangedListener, 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) {
+            this.autoReject(`cursor position move too far away off ${autoRejectEditCursorDistance} lines`)
+        }
+    }
 
-        if (svgImage) {
-            // display the SVG image
-            await displaySvgDecoration(
-                editor,
-                svgImage,
-                startLine,
-                newCode,
-                originalCodeHighlightRange,
-                session,
-                languageClient,
-                item,
-                inlineCompletionProvider
+    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
+        }
+
+        // TODO: handle multi-contentChanges scenario
+        const diff = e.contentChanges[0] ? e.contentChanges[0].text : ''
+        this.logger.info(`docChange sessionId=${this.session.sessionId}, contentChange=${diff}`)
+
+        // Track document changes because we might need to hide/reject suggestions while users are typing for better UX
+        this.documentChangeTrace.contentChanged += e.contentChanges[0].text
+        this.documentChangeTrace.count++
+        /**
+         * 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.documentChangeTrace.contentChanged.length > maxPrefixRetryCharDiff ||
+                    this.documentChangeTrace.count > maxPrefixRetryCharDiff
+                ) {
+                    // Reject the suggestion if users've typed over 5 characters while the suggestion is shown
+                    this.autoReject(RejectReason.MaxRetry)
+                } else if (applyPatch(this.editor.document.getText(), updatedPatch.insertText as string) === false) {
+                    this.autoReject(RejectReason.DocumentChange)
+                } else {
+                    // Close the previoius popup and rerender it
+                    this.logger.debug(`calling rerender with suggestion\n ${updatedPatch.insertText as string}`)
+                    await this.debouncedRerender(updatedPatch)
+                }
+            } else {
+                this.autoReject(RejectReason.NotApplicableToOriginal)
+            }
+        } catch (e) {
+            this.logger.error(`encountered error while processing edit suggestion when users type ${e}`)
+            // TODO: Maybe we should auto reject/hide suggestions in this scenario
+        }
+    }
+
+    async dispose() {
+        this.documentChangedListener?.dispose()
+        this.cursorChangedListener?.dispose()
+        await decorationManager.clearDecorations(this.editor, [])
+    }
+
+    debouncedRerender = debounce(
+        async (suggestion: InlineCompletionItemWithReferences) => await this.rerender(suggestion),
+        rerenderDeboucneInMs,
+        true
+    )
+
+    private async rerender(suggestion: InlineCompletionItemWithReferences) {
+        await decorationManager.clearDecorations(this.editor, [])
+        await this.show(suggestion)
+    }
+
+    private autoReject(reason: string) {
+        function logSuggestionFailure(type: 'REJECT', reason: string, suggestionContent: string) {
+            getLogger('nextEditPrediction').debug(
+                `Auto ${type} edit suggestion with reason=${reason}, suggetion: ${suggestionContent}`
             )
-        } else {
-            getLogger('nextEditPrediction').error('SVG image generation returned an empty result.')
         }
-    } catch (error) {
-        getLogger('nextEditPrediction').error(`Error generating SVG image: ${error}`)
+
+        logSuggestionFailure('REJECT', reason, this.suggestion.insertText as string)
+        void vscode.commands.executeCommand('aws.amazonq.inline.rejectEdit')
+    }
+
+    private patchSuggestion(appliedToOriginal: string): InlineCompletionItemWithReferences {
+        const updatedPatch = createPatch(
+            this.editor.document.fileName,
+            this.editor.document.getText(),
+            appliedToOriginal
+        )
+        this.logger.info(`Update edit suggestion\n ${updatedPatch}`)
+        return { ...this.suggestion, insertText: updatedPatch }
     }
 }

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

@@ -42,9 +42,9 @@ import {
 import { LineTracker } from './stateTracker/lineTracker'
 import { InlineTutorialAnnotation } from './tutorials/inlineTutorialAnnotation'
 import { TelemetryHelper } from './telemetryHelper'
-import { Experiments, getLogger, sleep } from 'aws-core-vscode/shared'
+import { Experiments, getContext, 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'
 
@@ -215,6 +215,7 @@ export class InlineCompletionManager implements Disposable {
 export class AmazonQInlineCompletionItemProvider implements InlineCompletionItemProvider {
     private logger = getLogger()
     private pendingRequest: Promise<InlineCompletionItem[]> | undefined
+    private lastEdit: EditsSuggestionSvg | undefined
 
     constructor(
         private readonly languageClient: BaseLanguageClient,
@@ -350,6 +351,11 @@ export class AmazonQInlineCompletionItemProvider implements InlineCompletionItem
             return []
         }
 
+        // Make edit suggestion blocking
+        if (getContext('aws.amazonq.editSuggestionActive') === true) {
+            return []
+        }
+
         // there is a bug in VS Code, when hitting Enter, the context.triggerKind is Invoke (0)
         // when hitting other keystrokes, the context.triggerKind is Automatic (1)
         // we only mark option + C as manual trigger
@@ -531,7 +537,12 @@ ${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)
+                        if (this.lastEdit) {
+                            await this.lastEdit.dispose()
+                        }
+                        const e = new EditsSuggestionSvg(item, editor, this.languageClient, session, this)
+                        await e.show()
+                        this.lastEdit = e
                         logstr += `- duration between trigger to edits suggestion is displayed: ${Date.now() - t0}ms`
                     }
                     return []