feat(amazonq): trigger inline suggestion via acceptance for EDITS

Merge pull request #7592 from atonaamz/nep-flare

Author: atonaamz

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

@@ -16,21 +16,13 @@ export type LineDiff =
  * @param unifiedDiff The unified diff content
  * @returns The modified code after applying the diff
  */
-export function applyUnifiedDiff(
-    docText: string,
-    unifiedDiff: string
-): { appliedCode: string; addedCharacterCount: number; deletedCharacterCount: number } {
+export function applyUnifiedDiff(docText: string, unifiedDiff: string): string {
     try {
-        const { addedCharacterCount, deletedCharacterCount } = getAddedAndDeletedCharCount(unifiedDiff)
         // First try the standard diff package
         try {
             const result = applyPatch(docText, unifiedDiff)
             if (result !== false) {
-                return {
-                    appliedCode: result,
-                    addedCharacterCount: addedCharacterCount,
-                    deletedCharacterCount: deletedCharacterCount,
-                }
+                return result
             }
         } catch (error) {}
 
@@ -94,49 +86,8 @@ export function applyUnifiedDiff(
             // Replace the text
             result = result.replace(textToReplace, newText)
         }
-        return {
-            appliedCode: result,
-            addedCharacterCount: addedCharacterCount,
-            deletedCharacterCount: deletedCharacterCount,
-        }
+        return result
     } catch (error) {
-        return {
-            appliedCode: docText, // Return original text if all methods fail
-            addedCharacterCount: 0,
-            deletedCharacterCount: 0,
-        }
-    }
-}
-
-export function getAddedAndDeletedCharCount(diff: string): {
-    addedCharacterCount: number
-    deletedCharacterCount: number
-} {
-    let addedCharacterCount = 0
-    let deletedCharacterCount = 0
-    let i = 0
-    const lines = diff.split('\n')
-    while (i < lines.length) {
-        const line = lines[i]
-        if (line.startsWith('+') && !line.startsWith('+++')) {
-            addedCharacterCount += line.length - 1
-        } else if (line.startsWith('-') && !line.startsWith('---')) {
-            const removedLine = line.substring(1)
-            deletedCharacterCount += removedLine.length
-
-            // Check if this is a modified line rather than a pure deletion
-            const nextLine = lines[i + 1]
-            if (nextLine && nextLine.startsWith('+') && !nextLine.startsWith('+++') && nextLine.includes(removedLine)) {
-                // This is a modified line, not a pure deletion
-                // We've already counted the deletion, so we'll just increment i to skip the next line
-                // since we'll process the addition on the next iteration
-                i += 1
-            }
-        }
-        i += 1
-    }
-    return {
-        addedCharacterCount,
-        deletedCharacterCount,
+        return docText // Return original text if all methods fail
     }
 }

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

@@ -12,6 +12,7 @@ import { LogInlineCompletionSessionResultsParams } from '@aws/language-server-ru
 import { InlineCompletionItemWithReferences } from '@aws/language-server-runtimes/protocol'
 import path from 'path'
 import { imageVerticalOffset } from './svgGenerator'
+import { AmazonQInlineCompletionItemProvider } from '../completion'
 
 export class EditDecorationManager {
     private imageDecorationType: vscode.TextEditorDecorationType
@@ -280,16 +281,15 @@ export async function displaySvgDecoration(
     session: CodeWhispererSession,
     languageClient: LanguageClient,
     item: InlineCompletionItemWithReferences,
-    addedCharacterCount: number,
-    deletedCharacterCount: number
+    inlineCompletionProvider?: AmazonQInlineCompletionItemProvider
 ) {
     const originalCode = editor.document.getText()
 
     decorationManager.displayEditSuggestion(
         editor,
         svgImage,
         startLine,
-        () => {
+        async () => {
             // Handle accept
             getLogger().info('Edit suggestion accepted')
 
@@ -315,10 +315,20 @@ export async function displaySvgDecoration(
                 },
                 totalSessionDisplayTime: Date.now() - session.requestStartTime,
                 firstCompletionDisplayLatency: session.firstCompletionDisplayLatency,
-                addedCharacterCount: addedCharacterCount,
-                deletedCharacterCount: deletedCharacterCount,
             }
             languageClient.sendNotification('aws/logInlineCompletionSessionResults', params)
+            if (inlineCompletionProvider) {
+                await inlineCompletionProvider.provideInlineCompletionItems(
+                    editor.document,
+                    endPosition,
+                    {
+                        triggerKind: vscode.InlineCompletionTriggerKind.Automatic,
+                        selectedCompletionInfo: undefined,
+                    },
+                    new vscode.CancellationTokenSource().token,
+                    { emitTelemetry: false, showUi: false }
+                )
+            }
         },
         () => {
             // Handle reject
@@ -333,8 +343,6 @@ export async function displaySvgDecoration(
                         discarded: false,
                     },
                 },
-                // addedCharacterCount: addedCharacterCount,
-                // deletedCharacterCount: deletedCharacterCount,
             }
             languageClient.sendNotification('aws/logInlineCompletionSessionResults', params)
         },

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

@@ -10,12 +10,14 @@ import { getLogger } from 'aws-core-vscode/shared'
 import { LanguageClient } from 'vscode-languageclient'
 import { InlineCompletionItemWithReferences } from '@aws/language-server-runtimes/protocol'
 import { CodeWhispererSession } from '../sessionManager'
+import { AmazonQInlineCompletionItemProvider } from '../completion'
 
 export async function showEdits(
     item: InlineCompletionItemWithReferences,
     editor: vscode.TextEditor | undefined,
     session: CodeWhispererSession,
-    languageClient: LanguageClient
+    languageClient: LanguageClient,
+    inlineCompletionProvider?: AmazonQInlineCompletionItemProvider
 ) {
     if (!editor) {
         return
@@ -24,14 +26,8 @@ export async function showEdits(
         const svgGenerationService = new SvgGenerationService()
         // Generate your SVG image with the file contents
         const currentFile = editor.document.uri.fsPath
-        const {
-            svgImage,
-            startLine,
-            newCode,
-            origionalCodeHighlightRange,
-            addedCharacterCount,
-            deletedCharacterCount,
-        } = await svgGenerationService.generateDiffSvg(currentFile, item.insertText as string)
+        const { svgImage, startLine, newCode, origionalCodeHighlightRange } =
+            await svgGenerationService.generateDiffSvg(currentFile, item.insertText as string)
 
         if (svgImage) {
             // display the SVG image
@@ -44,8 +40,7 @@ export async function showEdits(
                 session,
                 languageClient,
                 item,
-                addedCharacterCount,
-                deletedCharacterCount
+                inlineCompletionProvider
             )
         } else {
             getLogger('nextEditPrediction').error('SVG image generation returned an empty result.')

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

@@ -7,7 +7,6 @@ import { diffChars } from 'diff'
 import * as vscode from 'vscode'
 import { ToolkitError, getLogger } from 'aws-core-vscode/shared'
 import { diffUtilities } from 'aws-core-vscode/shared'
-import { applyUnifiedDiff } from './diffUtils'
 type Range = { line: number; start: number; end: number }
 
 const logger = getLogger('nextEditPrediction')
@@ -29,16 +28,13 @@ export class SvgGenerationService {
         startLine: number
         newCode: string
         origionalCodeHighlightRange: Range[]
-        addedCharacterCount: number
-        deletedCharacterCount: number
     }> {
         const textDoc = await vscode.workspace.openTextDocument(filePath)
         const originalCode = textDoc.getText()
         if (originalCode === '') {
             logger.error(`udiff format error`)
             throw new ToolkitError('udiff format error')
         }
-        const { addedCharacterCount, deletedCharacterCount } = applyUnifiedDiff(originalCode, udiff)
         const newCode = await diffUtilities.getPatchedCode(filePath, udiff)
         const modifiedLines = diffUtilities.getModifiedLinesFromUnifiedDiff(udiff)
         // TODO remove
@@ -91,8 +87,6 @@ export class SvgGenerationService {
             startLine: editStartLine,
             newCode: newCode,
             origionalCodeHighlightRange: highlightRanges.removedRanges,
-            addedCharacterCount,
-            deletedCharacterCount,
         }
     }
 

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

@@ -354,7 +354,7 @@ ${itemLog}
                 if (item.isInlineEdit) {
                     // Check if Next Edit Prediction feature flag is enabled
                     if (Experiments.instance.isExperimentEnabled('amazonqLSPNEP')) {
-                        void showEdits(item, editor, session, this.languageClient).then(() => {
+                        void showEdits(item, editor, session, this.languageClient, this).then(() => {
                             const t3 = performance.now()
                             logstr = logstr + `- duration since trigger to NEP suggestion is displayed: ${t3 - t0}ms`
                             this.logger.info(logstr)

packages/amazonq/test/unit/app/inline/EditRendering/diffUtils.test.ts

@@ -4,7 +4,7 @@
  */
 
 import * as assert from 'assert'
-import { applyUnifiedDiff, getAddedAndDeletedCharCount } from '../../../../../src/app/inline/EditRendering/diffUtils'
+import { applyUnifiedDiff } from '../../../../../src/app/inline/EditRendering/diffUtils'
 
 describe('diffUtils', function () {
     describe('applyUnifiedDiff', function () {
@@ -27,35 +27,13 @@ describe('diffUtils', function () {
             const expectedResult = 'function add(a, b) {\n  // Add two numbers\n  return a + b; // Return the sum\n}'
 
             // Apply the diff
-            const { appliedCode } = applyUnifiedDiff(originalCode, unifiedDiff)
+            const appliedCode = applyUnifiedDiff(originalCode, unifiedDiff)
 
             // Verify the result
             assert.strictEqual(appliedCode, expectedResult)
         })
     })
 
-    describe('getAddedAndDeletedCharCount', function () {
-        it('should correctly calculate added and deleted character counts', function () {
-            // Unified diff with additions and deletions
-            const unifiedDiff =
-                '--- a/file.js\n' +
-                '+++ b/file.js\n' +
-                '@@ -1,3 +1,4 @@\n' +
-                ' function add(a, b) {\n' +
-                '+  // Add two numbers\n' +
-                '-  return a + b;\n' +
-                '+  return a + b; // Return the sum\n' +
-                ' }'
-
-            // Calculate character counts
-            const { addedCharacterCount, deletedCharacterCount } = getAddedAndDeletedCharCount(unifiedDiff)
-
-            // Verify the counts with the actual values from the implementation
-            assert.strictEqual(addedCharacterCount, 20)
-            assert.strictEqual(deletedCharacterCount, 15)
-        })
-    })
-
     describe('applyUnifiedDiff with complex changes', function () {
         it('should handle multiple hunks in a diff', function () {
             // Original code with multiple functions
@@ -96,7 +74,7 @@ describe('diffUtils', function () {
                 '}'
 
             // Apply the diff
-            const { appliedCode } = applyUnifiedDiff(originalCode, unifiedDiff)
+            const appliedCode = applyUnifiedDiff(originalCode, unifiedDiff)
 
             // Verify the result
             assert.strictEqual(appliedCode, expectedResult)

packages/amazonq/test/unit/app/inline/EditRendering/imageRenderer.test.ts

@@ -31,8 +31,6 @@ describe('showEdits', function () {
             startLine: 5,
             newCode: 'console.log("Hello World");',
             origionalCodeHighlightRange: [{ line: 5, start: 0, end: 10 }],
-            addedCharacterCount: 25,
-            deletedCharacterCount: 0,
             ...overrides,
         }
     }
@@ -170,9 +168,7 @@ describe('showEdits', function () {
             mockSvgResult.origionalCodeHighlightRange,
             sessionStub,
             languageClientStub,
-            itemStub,
-            mockSvgResult.addedCharacterCount,
-            mockSvgResult.deletedCharacterCount
+            itemStub
         )
 
         // Verify no errors were logged