feat(amazonq): auto-reject edit suggestions when cursor moves >25 lines away (#7868)

floralph · web-flow · commit 3135e98bd0ff · 2025-08-12T12:06:03.000-07:00
## Problem Amazon Q edit suggestions remain active even when users navigate far away from the suggestion location, leading to confusion and potential unintended edits when users accept suggestions they can no longer see. ## Solution Added automatic rejection of edit suggestions when the cursor moves more than 25 lines away from the suggestion location. Implemented cursor change listener with proper cleanup to prevent memory leaks and ensure suggestions are only accepted when contextually relevant. Fixed 2 failing recommendationService.test.ts tests --- - Treat all work as PUBLIC. Private `feature/x` branches will not be squash-merged at release time. - Your code changes must meet the guidelines in [CONTRIBUTING.md](https://github.com/aws/aws-toolkit-vscode/blob/master/CONTRIBUTING.md#guidelines). - License: I confirm that my contribution is made under the terms of the Apache 2.0 license.
diff --git a/packages/amazonq/src/app/inline/EditRendering/displayImage.ts b/packages/amazonq/src/app/inline/EditRendering/displayImage.ts
@@ -16,6 +16,8 @@ import { EditSuggestionState } from '../editSuggestionState'
 import type { AmazonQInlineCompletionItemProvider } from '../completion'
 import { vsCodeState } from 'aws-core-vscode/codewhisperer'
 
+const autoRejectEditCursorDistance = 25
+
 export class EditDecorationManager {
     private imageDecorationType: vscode.TextEditorDecorationType
     private removedCodeDecorationType: vscode.TextEditorDecorationType
@@ -346,6 +348,19 @@ export async function displaySvgDecoration(
             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) {
+            void vscode.commands.executeCommand('aws.amazonq.inline.rejectEdit')
+        }
+    })
     await decorationManager.displayEditSuggestion(
         editor,
         svgImage,
@@ -371,6 +386,7 @@ export async function displaySvgDecoration(
 
             await decorationManager.clearDecorations(editor)
             documentChangeListener.dispose()
+            cursorChangeListener.dispose()
             const params: LogInlineCompletionSessionResultsParams = {
                 sessionId: session.sessionId,
                 completionSessionResult: {
@@ -405,6 +421,7 @@ export async function displaySvgDecoration(
             getLogger().info('Edit suggestion rejected')
             await decorationManager.clearDecorations(editor)
             documentChangeListener.dispose()
+            cursorChangeListener.dispose()
             const params: LogInlineCompletionSessionResultsParams = {
                 sessionId: session.sessionId,
                 completionSessionResult: {
diff --git a/packages/amazonq/test/unit/amazonq/apps/inline/recommendationService.test.ts b/packages/amazonq/test/unit/amazonq/apps/inline/recommendationService.test.ts
@@ -129,6 +129,9 @@ describe('RecommendationService', () => {
 
     describe('getAllRecommendations', () => {
         it('should handle single request with no partial result token', async () => {
+            // Mock EditSuggestionState to return false (no edit suggestion active)
+            sandbox.stub(EditSuggestionState, 'isEditSuggestionActive').returns(false)
+
             const mockFirstResult = {
                 sessionId: 'test-session',
                 items: [mockInlineCompletionItemOne],
@@ -172,6 +175,9 @@ describe('RecommendationService', () => {
         })
 
         it('should handle multiple request with partial result token', async () => {
+            // Mock EditSuggestionState to return false (no edit suggestion active)
+            sandbox.stub(EditSuggestionState, 'isEditSuggestionActive').returns(false)
+
             const mockFirstResult = {
                 sessionId: 'test-session',
                 items: [mockInlineCompletionItemOne],
diff --git a/packages/amazonq/test/unit/app/inline/EditRendering/displayImage.test.ts b/packages/amazonq/test/unit/app/inline/EditRendering/displayImage.test.ts
@@ -6,7 +6,30 @@
 import * as vscode from 'vscode'
 import * as sinon from 'sinon'
 import assert from 'assert'
-import { EditDecorationManager } from '../../../../../src/app/inline/EditRendering/displayImage'
+import { EditDecorationManager, displaySvgDecoration } from '../../../../../src/app/inline/EditRendering/displayImage'
+import { EditSuggestionState } from '../../../../../src/app/inline/editSuggestionState'
+
+// Shared helper function to create common stubs
+function createCommonStubs(sandbox: sinon.SinonSandbox) {
+    const documentStub = {
+        getText: sandbox.stub().returns('Original code content'),
+        uri: vscode.Uri.file('/test/file.ts'),
+        lineAt: sandbox.stub().returns({
+            text: 'Line text content',
+            range: new vscode.Range(0, 0, 0, 18),
+            rangeIncludingLineBreak: new vscode.Range(0, 0, 0, 19),
+            firstNonWhitespaceCharacterIndex: 0,
+            isEmptyOrWhitespace: false,
+        }),
+    } as unknown as sinon.SinonStubbedInstance<vscode.TextDocument>
+
+    const editorStub = {
+        document: documentStub,
+        setDecorations: sandbox.stub(),
+    } as unknown as sinon.SinonStubbedInstance<vscode.TextEditor>
+
+    return { documentStub, editorStub }
+}
 
 describe('EditDecorationManager', function () {
     let sandbox: sinon.SinonSandbox
@@ -25,23 +48,13 @@ describe('EditDecorationManager', function () {
             dispose: sandbox.stub(),
         } as unknown as sinon.SinonStubbedInstance<vscode.TextEditorDecorationType>
 
-        documentStub = {
-            getText: sandbox.stub().returns('Original code content'),
-            lineCount: 5,
-            lineAt: sandbox.stub().returns({
-                text: 'Line text content',
-                range: new vscode.Range(0, 0, 0, 18),
-                rangeIncludingLineBreak: new vscode.Range(0, 0, 0, 19),
-                firstNonWhitespaceCharacterIndex: 0,
-                isEmptyOrWhitespace: false,
-            }),
-        } as unknown as sinon.SinonStubbedInstance<vscode.TextDocument>
-
-        editorStub = {
-            document: documentStub,
-            setDecorations: sandbox.stub(),
-            edit: sandbox.stub().resolves(true),
-        } as unknown as sinon.SinonStubbedInstance<vscode.TextEditor>
+        const commonStubs = createCommonStubs(sandbox)
+        documentStub = commonStubs.documentStub
+        editorStub = commonStubs.editorStub
+
+        // Add additional properties needed for this test suite - extend the stub objects
+        Object.assign(documentStub, { lineCount: 5 })
+        Object.assign(editorStub, { edit: sandbox.stub().resolves(true) })
 
         windowStub = sandbox.stub(vscode.window)
         windowStub.createTextEditorDecorationType.returns(decorationTypeStub as any)
@@ -174,3 +187,124 @@ describe('EditDecorationManager', function () {
         sinon.assert.calledWith(editorStub.setDecorations.secondCall, manager['removedCodeDecorationType'], [])
     })
 })
+
+describe('displaySvgDecoration cursor distance auto-reject', function () {
+    let sandbox: sinon.SinonSandbox
+    let editorStub: sinon.SinonStubbedInstance<vscode.TextEditor>
+    let windowStub: sinon.SinonStub
+    let commandsStub: sinon.SinonStub
+    let editSuggestionStateStub: sinon.SinonStub
+    let onDidChangeTextEditorSelectionStub: sinon.SinonStub
+    let selectionChangeListener: (e: vscode.TextEditorSelectionChangeEvent) => void
+
+    // Helper function to setup displaySvgDecoration
+    async function setupDisplaySvgDecoration(startLine: number) {
+        return await displaySvgDecoration(
+            editorStub as unknown as vscode.TextEditor,
+            vscode.Uri.parse('data:image/svg+xml;base64,test'),
+            startLine,
+            'new code',
+            [],
+            {} as any,
+            {} as any,
+            { itemId: 'test', insertText: 'patch' } as any
+        )
+    }
+
+    // Helper function to create selection change event
+    function createSelectionChangeEvent(line: number): vscode.TextEditorSelectionChangeEvent {
+        const position = new vscode.Position(line, 0)
+        const selection = new vscode.Selection(position, position)
+        return {
+            textEditor: editorStub,
+            selections: [selection],
+            kind: vscode.TextEditorSelectionChangeKind.Mouse,
+        } as vscode.TextEditorSelectionChangeEvent
+    }
+
+    beforeEach(function () {
+        sandbox = sinon.createSandbox()
+
+        const commonStubs = createCommonStubs(sandbox)
+        editorStub = commonStubs.editorStub
+
+        // Mock vscode.window.onDidChangeTextEditorSelection
+        onDidChangeTextEditorSelectionStub = sandbox.stub()
+        onDidChangeTextEditorSelectionStub.returns({ dispose: sandbox.stub() })
+        windowStub = sandbox.stub(vscode.window, 'onDidChangeTextEditorSelection')
+        windowStub.callsFake((callback) => {
+            selectionChangeListener = callback
+            return { dispose: sandbox.stub() }
+        })
+
+        // Mock vscode.commands.executeCommand
+        commandsStub = sandbox.stub(vscode.commands, 'executeCommand')
+
+        // Mock EditSuggestionState
+        editSuggestionStateStub = sandbox.stub(EditSuggestionState, 'isEditSuggestionActive')
+        editSuggestionStateStub.returns(true)
+
+        // Mock other required dependencies
+        sandbox.stub(vscode.workspace, 'onDidChangeTextDocument').returns({ dispose: sandbox.stub() })
+    })
+
+    afterEach(function () {
+        sandbox.restore()
+    })
+
+    it('should not reject when cursor moves less than 25 lines away', async function () {
+        const startLine = 50
+        await setupDisplaySvgDecoration(startLine)
+
+        selectionChangeListener(createSelectionChangeEvent(startLine + 24))
+
+        sinon.assert.notCalled(commandsStub)
+    })
+
+    it('should not reject when cursor moves exactly 25 lines away', async function () {
+        const startLine = 50
+        await setupDisplaySvgDecoration(startLine)
+
+        selectionChangeListener(createSelectionChangeEvent(startLine + 25))
+
+        sinon.assert.notCalled(commandsStub)
+    })
+
+    it('should reject when cursor moves more than 25 lines away', async function () {
+        const startLine = 50
+        await setupDisplaySvgDecoration(startLine)
+
+        selectionChangeListener(createSelectionChangeEvent(startLine + 26))
+
+        sinon.assert.calledOnceWithExactly(commandsStub, 'aws.amazonq.inline.rejectEdit')
+    })
+
+    it('should reject when cursor moves more than 25 lines before the edit', async function () {
+        const startLine = 50
+        await setupDisplaySvgDecoration(startLine)
+
+        selectionChangeListener(createSelectionChangeEvent(startLine - 26))
+
+        sinon.assert.calledOnceWithExactly(commandsStub, 'aws.amazonq.inline.rejectEdit')
+    })
+
+    it('should not reject when edit is near beginning of file and cursor cannot move far enough', async function () {
+        const startLine = 10
+        await setupDisplaySvgDecoration(startLine)
+
+        selectionChangeListener(createSelectionChangeEvent(0))
+
+        sinon.assert.notCalled(commandsStub)
+    })
+
+    it('should not reject when edit suggestion is not active', async function () {
+        editSuggestionStateStub.returns(false)
+
+        const startLine = 50
+        await setupDisplaySvgDecoration(startLine)
+
+        selectionChangeListener(createSelectionChangeEvent(startLine + 100))
+
+        sinon.assert.notCalled(commandsStub)
+    })
+})
