fix: scope inline chat accept/reject lenses to the task document

packages/amazonq/.changes/next-release/Bug Fix-b6847722-4622-43ce-be85-ac4828a99a13.json

@@ -0,0 +1,4 @@
+{
+	"type": "Bug Fix",
+	"description": "Inline chat: Accept/Reject code lenses no longer appear in other files when switching editors before accepting or rejecting a suggestion"
+}

packages/amazonq/src/inlineChat/codeLenses/codeLenseProvider.ts

@@ -9,6 +9,7 @@ import { InlineTask, TaskState } from '../controller/inlineTask'
 
 export class CodelensProvider implements vscode.CodeLensProvider {
     private codeLenses: vscode.CodeLens[] = []
+    private taskDocumentUri: vscode.Uri | undefined
     private _onDidChangeCodeLenses: vscode.EventEmitter<void> = new vscode.EventEmitter<void>()
     public readonly onDidChangeCodeLenses: vscode.Event<void> = this._onDidChangeCodeLenses.event
 
@@ -17,16 +18,24 @@ export class CodelensProvider implements vscode.CodeLensProvider {
         this.provideCodeLenses = this.provideCodeLenses.bind(this)
     }
 
-    public provideCodeLenses(_document: vscode.TextDocument, _token: vscode.CancellationToken): vscode.CodeLens[] {
+    public provideCodeLenses(document: vscode.TextDocument, _token: vscode.CancellationToken): vscode.CodeLens[] {
+        // Only surface the inline-chat lenses in the document that owns the active task.
+        // The provider is registered for all documents ('*'), so without this guard the
+        // accept/reject lenses would also appear in other files the user switches to.
+        if (this.taskDocumentUri === undefined || document.uri.toString() !== this.taskDocumentUri.toString()) {
+            return []
+        }
         return this.codeLenses
     }
 
     public updateLenses(task: InlineTask): void {
         if (task.state === TaskState.Complete) {
             this.codeLenses = []
+            this.taskDocumentUri = undefined
             this._onDidChangeCodeLenses.fire()
             return
         }
+        this.taskDocumentUri = task.document.uri
         switch (task.state) {
             case TaskState.InProgress: {
                 this.codeLenses = []

packages/amazonq/test/unit/inlineChat/codeLenseProvider.test.ts

@@ -0,0 +1,57 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import assert from 'assert'
+import * as vscode from 'vscode'
+import { FakeExtensionContext } from 'aws-core-vscode/test'
+import { CodelensProvider } from '../../../src/inlineChat/codeLenses/codeLenseProvider'
+import { InlineTask, TaskState } from '../../../src/inlineChat/controller/inlineTask'
+
+describe('inline chat CodelensProvider', function () {
+    let provider: CodelensProvider
+    const token = new vscode.CancellationTokenSource().token
+    const uriA = vscode.Uri.parse('file:///tmp/a.ts')
+    const uriB = vscode.Uri.parse('file:///tmp/b.ts')
+
+    // Lightweight stand-in for an InlineTask exposing only the fields updateLenses reads.
+    function makeTask(uri: vscode.Uri, state: TaskState): InlineTask {
+        return {
+            state,
+            selectedRange: new vscode.Range(0, 0, 0, 0),
+            document: { uri } as vscode.TextDocument,
+        } as unknown as InlineTask
+    }
+
+    function docFor(uri: vscode.Uri): vscode.TextDocument {
+        return { uri } as vscode.TextDocument
+    }
+
+    beforeEach(async function () {
+        provider = new CodelensProvider(await FakeExtensionContext.create())
+    })
+
+    it('shows accept/reject lenses only in the task document, not other files', function () {
+        provider.updateLenses(makeTask(uriA, TaskState.WaitingForDecision))
+
+        // The document that owns the task shows the two decision lenses.
+        assert.strictEqual(provider.provideCodeLenses(docFor(uriA), token).length, 2)
+        // Switching to another file must not surface the lenses (regression guard).
+        assert.strictEqual(provider.provideCodeLenses(docFor(uriB), token).length, 0)
+    })
+
+    it('shows the in-progress lens only in the task document', function () {
+        provider.updateLenses(makeTask(uriA, TaskState.InProgress))
+
+        assert.strictEqual(provider.provideCodeLenses(docFor(uriA), token).length, 1)
+        assert.strictEqual(provider.provideCodeLenses(docFor(uriB), token).length, 0)
+    })
+
+    it('clears the lenses once the task completes', function () {
+        provider.updateLenses(makeTask(uriA, TaskState.WaitingForDecision))
+        provider.updateLenses(makeTask(uriA, TaskState.Complete))
+
+        assert.strictEqual(provider.provideCodeLenses(docFor(uriA), token).length, 0)
+    })
+})