From d9d0b546e060f26435bfdb474af33535c0ae9700 Mon Sep 17 00:00:00 2001
From: Boyu Wang <bywang@amazon.com>
Date: Fri, 3 Oct 2025 15:21:07 -0700
Subject: [PATCH] fix(amazonq): block completion before active edit is
 accepted/rejected

---
 .../app/inline/EditRendering/displayImage.ts  | 21 +++++++++++++++----
 .../src/app/inline/recommendationService.ts   |  2 +-
 .../apps/inline/recommendationService.test.ts |  8 ++++---
 3 files changed, 23 insertions(+), 8 deletions(-)

diff --git a/packages/amazonq/src/app/inline/EditRendering/displayImage.ts b/packages/amazonq/src/app/inline/EditRendering/displayImage.ts
index d719191035e..0af4d4801c0 100644
--- a/packages/amazonq/src/app/inline/EditRendering/displayImage.ts
+++ b/packages/amazonq/src/app/inline/EditRendering/displayImage.ts
@@ -136,10 +136,11 @@ export class EditDecorationManager {
         newCode: string,
         originalCodeHighlightRanges: Array<{ line: number; start: number; end: number }>
     ): Promise<void> {
-        await this.clearDecorations(editor)
-
-        await setContext('aws.amazonq.editSuggestionActive' as any, true)
-        EditSuggestionState.setEditSuggestionActive(true)
+        // Clear old decorations but don't reset state (state is already set in displaySvgDecoration)
+        editor.setDecorations(this.imageDecorationType, [])
+        editor.setDecorations(this.removedCodeDecorationType, [])
+        this.currentImageDecoration = undefined
+        this.currentRemovedCodeDecorations = []
 
         this.acceptHandler = onAccept
         this.rejectHandler = onReject
@@ -313,8 +314,16 @@ export async function displaySvgDecoration(
 ) {
     const originalCode = editor.document.getText()
 
+    // Set edit state immediately to prevent race condition with completion requests
+    await setContext('aws.amazonq.editSuggestionActive' as any, true)
+    EditSuggestionState.setEditSuggestionActive(true)
+
     // Check if a completion suggestion is currently active - if so, discard edit suggestion
     if (inlineCompletionProvider && (await inlineCompletionProvider.isCompletionActive())) {
+        // Clean up state since we're not showing the edit
+        await setContext('aws.amazonq.editSuggestionActive' as any, false)
+        EditSuggestionState.setEditSuggestionActive(false)
+
         // Emit DISCARD telemetry for edit suggestion that can't be shown due to active completion
         const params = createDiscardTelemetryParams(session, item)
         languageClient.sendNotification('aws/logInlineCompletionSessionResults', params)
@@ -326,6 +335,10 @@ export async function displaySvgDecoration(
 
     const isPatchValid = applyPatch(editor.document.getText(), item.insertText as string)
     if (!isPatchValid) {
+        // Clean up state since we're not showing the edit
+        await setContext('aws.amazonq.editSuggestionActive' as any, false)
+        EditSuggestionState.setEditSuggestionActive(false)
+
         const params = createDiscardTelemetryParams(session, item)
         // TODO: this session is closed on flare side hence discarded is not emitted in flare
         languageClient.sendNotification('aws/logInlineCompletionSessionResults', params)
diff --git a/packages/amazonq/src/app/inline/recommendationService.ts b/packages/amazonq/src/app/inline/recommendationService.ts
index e9075ff426e..ad28910611b 100644
--- a/packages/amazonq/src/app/inline/recommendationService.ts
+++ b/packages/amazonq/src/app/inline/recommendationService.ts
@@ -138,7 +138,7 @@ export class RecommendationService {
              * Completions use PartialResultToken with single 1 call of [getAllRecommendations].
              * 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) {
+            if (!isTriggerByDeletion && !request.partialResultToken && !EditSuggestionState.isEditSuggestionActive()) {
                 const completionPromise: Promise<InlineCompletionListWithReferences> = languageClient.sendRequest(
                     inlineCompletionWithReferencesRequestType.method,
                     request,
diff --git a/packages/amazonq/test/unit/amazonq/apps/inline/recommendationService.test.ts b/packages/amazonq/test/unit/amazonq/apps/inline/recommendationService.test.ts
index 0b983f333b4..a051ef94abb 100644
--- a/packages/amazonq/test/unit/amazonq/apps/inline/recommendationService.test.ts
+++ b/packages/amazonq/test/unit/amazonq/apps/inline/recommendationService.test.ts
@@ -333,7 +333,7 @@ describe('RecommendationService', () => {
             }
         })
 
-        it('should make completion request when edit suggestion is active', async () => {
+        it('should not make completion request when edit suggestion is active', async () => {
             // Mock EditSuggestionState to return true (edit suggestion is active)
             sandbox.stub(EditSuggestionState, 'isEditSuggestionActive').returns(true)
 
@@ -360,13 +360,15 @@ describe('RecommendationService', () => {
             const completionCalls = cs.filter((c) => c.firstArg === completionApi)
             const editCalls = cs.filter((c) => c.firstArg === editApi)
 
-            assert.strictEqual(cs.length, 2) // Only edit call
-            assert.strictEqual(completionCalls.length, 1) // No completion calls
+            assert.strictEqual(cs.length, 1) // Only edit call
+            assert.strictEqual(completionCalls.length, 0) // No completion calls
             assert.strictEqual(editCalls.length, 1) // One edit call
         })
 
         it('should make completion request when edit suggestion is not active', async () => {
             // Mock EditSuggestionState to return false (no edit suggestion active)
+            sandbox.stub(EditSuggestionState, 'isEditSuggestionActive').returns(false)
+
             const mockResult = {
                 sessionId: 'test-session',
                 items: [mockInlineCompletionItemOne],