fix(amazonq): apply fix removes other issues (#6236)

## Problem

Apply fix removes other issues in the same file. This is happening
because apply fix command will always replace the entire file contents
which triggers a document change event that intersects with all issues
in the file.


## Solution

Look through the diff hunks and apply them in a range.


---

- 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.

Author: ctlai95

packages/amazonq/.changes/next-release/Bug Fix-1fd5992c-cc2e-49ac-8afc-360d7f0a3966.json

@@ -0,0 +1,4 @@
+{
+	"type": "Bug Fix",
+	"description": "/review: Apply fix removes other issues in the same file."
+}

packages/core/src/codewhisperer/activation.ts

@@ -670,27 +670,28 @@ export async function activate(context: ExtContext): Promise<void> {
     function setSubscriptionsForCodeIssues() {
         context.extensionContext.subscriptions.push(
             vscode.workspace.onDidChangeTextDocument(async (e) => {
-                // verify the document is something with a finding
-                for (const issue of SecurityIssueProvider.instance.issues) {
-                    if (issue.filePath === e.document.uri.fsPath) {
-                        disposeSecurityDiagnostic(e)
-
-                        SecurityIssueProvider.instance.handleDocumentChange(e)
-                        SecurityIssueTreeViewProvider.instance.refresh()
-                        await syncSecurityIssueWebview(context)
-
-                        toggleIssuesVisibility((issue, filePath) =>
-                            filePath !== e.document.uri.fsPath
-                                ? issue.visible
-                                : !detectCommentAboveLine(
-                                      e.document,
-                                      issue.startLine,
-                                      CodeWhispererConstants.amazonqIgnoreNextLine
-                                  )
-                        )
-                        break
-                    }
+                if (e.document.uri.scheme !== 'file') {
+                    return
                 }
+                const diagnostics = securityScanRender.securityDiagnosticCollection?.get(e.document.uri)
+                if (!diagnostics || diagnostics.length === 0) {
+                    return
+                }
+                disposeSecurityDiagnostic(e)
+
+                SecurityIssueProvider.instance.handleDocumentChange(e)
+                SecurityIssueTreeViewProvider.instance.refresh()
+                await syncSecurityIssueWebview(context)
+
+                toggleIssuesVisibility((issue, filePath) =>
+                    filePath !== e.document.uri.fsPath
+                        ? issue.visible
+                        : !detectCommentAboveLine(
+                              e.document,
+                              issue.startLine,
+                              CodeWhispererConstants.amazonqIgnoreNextLine
+                          )
+                )
             })
         )
     }

packages/core/src/codewhisperer/commands/basicCommands.ts

@@ -66,6 +66,7 @@ import { cancel, confirm } from '../../shared'
 import { startCodeFixGeneration } from './startCodeFixGeneration'
 import { DefaultAmazonQAppInitContext } from '../../amazonq/apps/initContext'
 import path from 'path'
+import { parsePatch } from 'diff'
 
 const MessageTimeOut = 5_000
 
@@ -459,11 +460,21 @@ export const applySecurityFix = Commands.declare(
             }
 
             const edit = new vscode.WorkspaceEdit()
-            edit.replace(
-                document.uri,
-                new vscode.Range(document.lineAt(0).range.start, document.lineAt(document.lineCount - 1).range.end),
-                updatedContent
-            )
+            const diffs = parsePatch(suggestedFix.code)
+            for (const diff of diffs) {
+                for (const hunk of [...diff.hunks].reverse()) {
+                    const startLine = document.lineAt(hunk.oldStart - 1)
+                    const endLine = document.lineAt(hunk.oldStart - 1 + hunk.oldLines - 1)
+                    const range = new vscode.Range(startLine.range.start, endLine.range.end)
+
+                    const newText = updatedContent
+                        .split('\n')
+                        .slice(hunk.newStart - 1, hunk.newStart - 1 + hunk.newLines)
+                        .join('\n')
+
+                    edit.replace(document.uri, range, newText)
+                }
+            }
             const isApplied = await vscode.workspace.applyEdit(edit)
             if (isApplied) {
                 void document.save().then((didSave) => {

packages/core/src/codewhisperer/service/securityIssueProvider.ts

@@ -4,7 +4,7 @@
  */
 
 import * as vscode from 'vscode'
-import { AggregatedCodeScanIssue, CodeScanIssue, SuggestedFix } from '../models/model'
+import { AggregatedCodeScanIssue, CodeScanIssue, CodeScansState, SuggestedFix } from '../models/model'
 export class SecurityIssueProvider {
     static #instance: SecurityIssueProvider
     public static get instance() {
@@ -25,13 +25,15 @@ export class SecurityIssueProvider {
         if (!event.contentChanges || event.contentChanges.length === 0) {
             return
         }
-        const { changedRange, lineOffset } = event.contentChanges.reduce(
+        const { changedRange, changedText, lineOffset } = event.contentChanges.reduce(
             (acc, change) => ({
                 changedRange: acc.changedRange.union(change.range),
+                changedText: acc.changedText + change.text,
                 lineOffset: acc.lineOffset + this._getLineOffset(change.range, change.text),
             }),
             {
                 changedRange: event.contentChanges[0].range,
+                changedText: '',
                 lineOffset: 0,
             }
         )
@@ -43,18 +45,20 @@ export class SecurityIssueProvider {
             return {
                 ...group,
                 issues: group.issues
-                    .filter(
-                        (issue) =>
-                            // Filter out any modified issues
-                            !changedRange.intersection(
-                                new vscode.Range(
-                                    issue.startLine,
-                                    event.document.lineAt(issue.startLine)?.range.start.character ?? 0,
-                                    issue.endLine,
-                                    event.document.lineAt(issue.endLine)?.range.end.character ?? 0
-                                )
-                            )
-                    )
+                    .filter((issue) => {
+                        const range = new vscode.Range(
+                            issue.startLine,
+                            event.document.lineAt(issue.startLine)?.range.start.character ?? 0,
+                            issue.endLine,
+                            event.document.lineAt(issue.endLine - 1)?.range.end.character ?? 0
+                        )
+                        const intersection = changedRange.intersection(range)
+                        return !(
+                            intersection &&
+                            (/\S/.test(changedText) || changedText === '') &&
+                            !CodeScansState.instance.isScansEnabled()
+                        )
+                    })
                     .map((issue) => {
                         if (issue.startLine < changedRange.end.line) {
                             return issue

packages/core/src/test/codewhisperer/commands/basicCommands.test.ts

@@ -672,6 +672,81 @@ describe('CodeWhisperer-basicCommands', function () {
                 reasonDesc: 'Failed to apply edit to the workspace.',
             })
         })
+
+        it('should apply the edit at the correct range', async function () {
+            const fileName = 'sample.py'
+            const textDocumentMock = createMockDocument(
+                `from flask import app
+
+
+@app.route('/')
+def execute_input_noncompliant():
+    from flask import request
+    module_version = request.args.get("module_version")
+    # Noncompliant: executes unsanitized inputs.
+    exec("import urllib%s as urllib" % module_version)
+# {/fact}
+
+
+# {fact rule=code-injection@v1.0 defects=0}
+from flask import app
+
+
+@app.route('/')
+def execute_input_compliant():
+    from flask import request
+    module_version = request.args.get("module_version")
+    # Compliant: executes sanitized inputs.
+    exec("import urllib%d as urllib" % int(module_version))
+# {/fact}`,
+                fileName
+            )
+            openTextDocumentMock.resolves(textDocumentMock)
+            sandbox.stub(vscode.workspace, 'openTextDocument').value(openTextDocumentMock)
+
+            sandbox.stub(vscode.WorkspaceEdit.prototype, 'replace').value(replaceMock)
+            applyEditMock.resolves(true)
+            sandbox.stub(vscode.workspace, 'applyEdit').value(applyEditMock)
+            sandbox.stub(diagnosticsProvider, 'removeDiagnostic').value(removeDiagnosticMock)
+            sandbox.stub(SecurityIssueProvider.instance, 'removeIssue').value(removeIssueMock)
+            sandbox.stub(vscode.window, 'showTextDocument').value(showTextDocumentMock)
+
+            targetCommand = testCommand(applySecurityFix)
+            codeScanIssue.suggestedFixes = [
+                {
+                    code: `@@ -6,4 +6,5 @@
+     from flask import request
+     module_version = request.args.get("module_version")
+     # Noncompliant: executes unsanitized inputs.
+-    exec("import urllib%d as urllib" % int(module_version))
++    __import__("urllib" + module_version)
++#import importlib`,
+                    description: 'dummy',
+                },
+            ]
+            await targetCommand.execute(codeScanIssue, fileName, 'webview')
+            assert.ok(
+                replaceMock.calledOnceWith(
+                    textDocumentMock.uri,
+                    new vscode.Range(5, 0, 8, 54),
+                    `    from flask import request
+    module_version = request.args.get("module_version")
+    # Noncompliant: executes unsanitized inputs.
+    __import__("urllib" + module_version)
+#import importlib`
+                )
+            )
+            assert.ok(applyEditMock.calledOnce)
+            assert.ok(removeDiagnosticMock.calledOnceWith(textDocumentMock.uri, codeScanIssue))
+            assert.ok(removeIssueMock.calledOnce)
+
+            assertTelemetry('codewhisperer_codeScanIssueApplyFix', {
+                detectorId: codeScanIssue.detectorId,
+                findingId: codeScanIssue.findingId,
+                component: 'webview',
+                result: 'Succeeded',
+            })
+        })
     })
 
     // describe('generateFix', function () {