123

Author: Will-ShaoHua

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

@@ -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,7 +310,7 @@ export async function displaySvgDecoration(
     session: CodeWhispererSession,
     languageClient: LanguageClient,
     item: InlineCompletionItemWithReferences,
-    documentChangeListener: vscode.Disposable,
+    listeners: vscode.Disposable[],
     inlineCompletionProvider?: AmazonQInlineCompletionItemProvider
 ) {
     function logSuggestionFailure(type: 'DISCARD' | 'REJECT', reason: string, suggestionContent: string) {
@@ -361,24 +360,6 @@ export async function displaySvgDecoration(
         return
     }
 
-    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,
@@ -400,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: {
@@ -426,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

@@ -13,20 +13,29 @@ 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'
 
 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
-    // TODO:
-    // private readonly cursorChangedListener: vscode.Disposable
-
+    private readonly cursorChangedListener: vscode.Disposable
     private readonly updatedSuggestions: InlineCompletionItemWithReferences[] = []
+    private startLine = 0
 
     constructor(
         private suggestion: InlineCompletionItemWithReferences,
@@ -38,6 +47,10 @@ export class EditsSuggestionSvg {
         this.documentChangedListener = vscode.workspace.onDidChangeTextDocument(async (e) => {
             await this.onDocChange(e)
         })
+
+        this.cursorChangedListener = vscode.window.onDidChangeTextEditorSelection((e) => {
+            this.onCursorChange(e)
+        })
     }
 
     async show() {
@@ -57,6 +70,9 @@ export class EditsSuggestionSvg {
             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')
@@ -74,7 +90,7 @@ export class EditsSuggestionSvg {
                     this.session,
                     this.languageClient,
                     item,
-                    this.documentChangedListener,
+                    [this.documentChangedListener, this.cursorChangedListener],
                     this.inlineCompletionProvider
                 )
             } else {
@@ -85,6 +101,25 @@ export class EditsSuggestionSvg {
         }
     }
 
+    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
+            )
+            void vscode.commands.executeCommand('aws.amazonq.inline.rejectEdit')
+        }
+    }
+
     private async onDocChange(e: vscode.TextDocumentChangeEvent) {
         if (e.contentChanges.length <= 0) {
             return
@@ -104,40 +139,44 @@ export class EditsSuggestionSvg {
          * 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 = createPatch(
-                    this.editor.document.fileName,
-                    this.editor.document.getText(),
-                    appliedToOriginal
-                )
-                // show updatedPatch
-                this.updateSuggestion(updatedPatch)
+                const updatedPatch = this.patchSuggestion(appliedToOriginal)
 
-                if (this.updatedSuggestions.length > 5) {
-                    this.autoReject()
+                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()
+                this.autoReject(RejectReason.NotApplicableToOriginal)
             }
         } catch (e) {
-            this.logger.info('123')
+            // TODO: format
+            this.logger.error(`${e}`)
         }
     }
 
-    private autoReject() {
-        logSuggestionFailure('REJECT', 'Invalid patch due to document change', this.suggestion.insertText as string)
+    private autoReject(reason: string) {
+        logSuggestionFailure('REJECT', reason, this.suggestion.insertText as string)
         void vscode.commands.executeCommand('aws.amazonq.inline.rejectEdit')
     }
 
-    private updateSuggestion(content: string) {
-        this.logger.info(`update edit suggestion\n ${content}`)
-        const updated: InlineCompletionItemWithReferences = { ...this.suggestion, insertText: content }
+    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
     }
 }