feat(amazonq): apply fix without save aws#4855

ctlai95 · web-flow · commit 72d8c3485746 · 2024-04-27T10:28:12.000-07:00
Problem
Scans can now run without saving the file to disk, but "Apply Fix" still saves the file to disk.

Solution
- Use WorkspaceEdit to apply fix
- If auto-scans is enabled, immediately remove the issue from hover and diagnostics after fixing it
diff --git a/packages/core/src/codewhisperer/commands/basicCommands.ts b/packages/core/src/codewhisperer/commands/basicCommands.ts
@@ -27,7 +27,6 @@ import {
 } from '../util/customizationUtil'
 import { applyPatch } from 'diff'
 import { closeSecurityIssueWebview, showSecurityIssueWebview } from '../views/securityIssue/securityIssueWebview'
-import { fsCommon } from '../../srcShared/fs'
 import { Mutable } from '../../shared/utilities/tsUtils'
 import { CodeWhispererSource } from './types'
 import { FeatureConfigProvider } from '../service/featureConfigProvider'
@@ -36,6 +35,9 @@ import { Auth, AwsConnection } from '../../auth'
 import { once } from '../../shared/utilities/functionUtils'
 import { isTextEditor } from '../../shared/utilities/editorUtilities'
 import { focusAmazonQPanel } from '../../codewhispererChat/commands/registerCommands'
+import { removeDiagnostic } from '../service/diagnosticsProvider'
+import { SecurityIssueHoverProvider } from '../service/securityIssueHoverProvider'
+import { SecurityIssueCodeActionProvider } from '../service/securityIssueCodeActionProvider'
 
 export const toggleCodeSuggestions = Commands.declare(
     { id: 'aws.amazonq.toggleCodeSuggestion', compositeKey: { 1: 'source' } },
@@ -343,14 +345,23 @@ export const applySecurityFix = Commands.declare(
                 throw Error('Failed to get updated content from applying diff patch')
             }
 
-            // saving the document text if not save
-            const isSaved = await document.save()
-            if (!isSaved) {
-                throw Error('Failed to save editor text changes into the file.')
+            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 isApplied = await vscode.workspace.applyEdit(edit)
+            if (!isApplied) {
+                throw Error('Failed to apply edit to the workspace.')
+            }
+
+            if (CodeScansState.instance.isScansEnabled()) {
+                removeDiagnostic(document.uri, issue)
+                SecurityIssueHoverProvider.instance.removeIssue(document.uri, issue)
+                SecurityIssueCodeActionProvider.instance.removeIssue(document.uri, issue)
             }
 
-            // writing the patch applied version of document into the file
-            await fsCommon.writeFile(filePath, updatedContent)
             await closeSecurityIssueWebview(issue.findingId)
         } catch (err) {
             getLogger().error(`Apply fix command failed. ${err}`)
diff --git a/packages/core/src/codewhisperer/service/diagnosticsProvider.ts b/packages/core/src/codewhisperer/service/diagnosticsProvider.ts
@@ -128,3 +128,19 @@ function getLineOffset(range: vscode.Range, text: string) {
     const changedLines = text.split('\n').length
     return changedLines - originLines
 }
+
+export function removeDiagnostic(uri: vscode.Uri, issue: CodeScanIssue) {
+    const currentSecurityDiagnostics = securityScanRender.securityDiagnosticCollection?.get(uri)
+    if (currentSecurityDiagnostics) {
+        const newSecurityDiagnostics = currentSecurityDiagnostics.filter(diagnostic => {
+            return (
+                typeof diagnostic.code !== 'string' &&
+                typeof diagnostic.code !== 'number' &&
+                diagnostic.code?.value !== issue.detectorId &&
+                diagnostic.message !== issue.title &&
+                diagnostic.range !== new vscode.Range(issue.startLine, 0, issue.endLine, 0)
+            )
+        })
+        securityScanRender.securityDiagnosticCollection?.set(uri, newSecurityDiagnostics)
+    }
+}
diff --git a/packages/core/src/codewhisperer/service/securityIssueProvider.ts b/packages/core/src/codewhisperer/service/securityIssueProvider.ts
@@ -4,7 +4,7 @@
  */
 
 import * as vscode from 'vscode'
-import { AggregatedCodeScanIssue, CodeScansState } from '../models/model'
+import { AggregatedCodeScanIssue, CodeScanIssue, CodeScansState } from '../models/model'
 export abstract class SecurityIssueProvider {
     private _issues: AggregatedCodeScanIssue[] = []
     public get issues() {
@@ -64,4 +64,16 @@ export abstract class SecurityIssueProvider {
         const changedLines = text.split('\n').length
         return changedLines - originLines
     }
+
+    public removeIssue(uri: vscode.Uri, issue: CodeScanIssue) {
+        this._issues = this._issues.map(group => {
+            if (group.filePath !== uri.fsPath) {
+                return group
+            }
+            return {
+                ...group,
+                issues: group.issues.filter(i => i.findingId !== issue.findingId),
+            }
+        })
+    }
 }
diff --git a/packages/core/src/test/codewhisperer/commands/basicCommands.test.ts b/packages/core/src/test/codewhisperer/commands/basicCommands.test.ts
@@ -28,8 +28,6 @@ import { AuthUtil } from '../../../codewhisperer/util/authUtil'
 import { getTestWindow } from '../../shared/vscode/window'
 import { ExtContext } from '../../../shared/extensions'
 import { get, set } from '../../../codewhisperer/util/commonUtil'
-import { MockDocument } from '../../fake/fakeDocument'
-import { FileSystemCommon } from '../../../srcShared/fs'
 import { getLogger } from '../../../shared/logger/logger'
 import {
     createAutoScans,
@@ -55,6 +53,9 @@ import { cwQuickPickSource } from '../../../codewhisperer/commands/types'
 import { isTextEditor } from '../../../shared/utilities/editorUtilities'
 import { refreshStatusBar } from '../../../codewhisperer/service/inlineCompletionService'
 import { focusAmazonQPanel } from '../../../codewhispererChat/commands/registerCommands'
+import * as diagnosticsProvider from '../../../codewhisperer/service/diagnosticsProvider'
+import { SecurityIssueHoverProvider } from '../../../codewhisperer/service/securityIssueHoverProvider'
+import { SecurityIssueCodeActionProvider } from '../../../codewhisperer/service/securityIssueCodeActionProvider'
 
 describe('CodeWhisperer-basicCommands', function () {
     let targetCommand: Command<any> & vscode.Disposable
@@ -392,15 +393,19 @@ describe('CodeWhisperer-basicCommands', function () {
 
     describe('applySecurityFix', function () {
         let sandbox: sinon.SinonSandbox
-        let saveStub: sinon.SinonStub
         let openTextDocumentMock: sinon.SinonStub
-        let writeFileMock: sinon.SinonStub
+        let replaceMock: sinon.SinonStub
+        let applyEditMock: sinon.SinonStub
+        let removeDiagnosticMock: sinon.SinonStub
+        let removeIssueMock: sinon.SinonStub
 
         beforeEach(function () {
             sandbox = sinon.createSandbox()
-            saveStub = sinon.stub()
             openTextDocumentMock = sinon.stub()
-            writeFileMock = sinon.stub()
+            replaceMock = sinon.stub()
+            applyEditMock = sinon.stub()
+            removeDiagnosticMock = sinon.stub()
+            removeIssueMock = sinon.stub()
         })
 
         afterEach(function () {
@@ -409,14 +414,17 @@ describe('CodeWhisperer-basicCommands', function () {
 
         it('should call applySecurityFix command successfully', async function () {
             const fileName = 'sample.py'
-            saveStub.resolves(true)
-            const textDocumentMock = new MockDocument('first line\n second line\n fourth line', fileName, saveStub)
+            const textDocumentMock = createMockDocument('first line\n second line\n fourth line', fileName)
 
             openTextDocumentMock.resolves(textDocumentMock)
             sandbox.stub(vscode.workspace, 'openTextDocument').value(openTextDocumentMock)
 
-            writeFileMock.resolves(true)
-            sinon.stub(FileSystemCommon.prototype, 'writeFile').value(writeFileMock)
+            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(SecurityIssueHoverProvider.instance, 'removeIssue').value(removeIssueMock)
+            sandbox.stub(SecurityIssueCodeActionProvider.instance, 'removeIssue').value(removeIssueMock)
 
             targetCommand = testCommand(applySecurityFix)
             const codeScanIssue = createCodeScanIssue({
@@ -428,8 +436,16 @@ describe('CodeWhisperer-basicCommands', function () {
                 ],
             })
             await targetCommand.execute(codeScanIssue, fileName, 'hover')
-            assert.ok(saveStub.calledOnce)
-            assert.ok(writeFileMock.calledOnceWith(fileName, 'first line\n third line\n fourth line'))
+            assert.ok(
+                replaceMock.calledOnceWith(
+                    textDocumentMock.uri,
+                    new vscode.Range(0, 0, 2, 12),
+                    'first line\n third line\n fourth line'
+                )
+            )
+            assert.ok(applyEditMock.calledOnce)
+            assert.ok(removeDiagnosticMock.calledOnceWith(textDocumentMock.uri, codeScanIssue))
+            assert.ok(removeIssueMock.calledTwice)
 
             assertTelemetry('codewhisperer_codeScanIssueApplyFix', {
                 detectorId: codeScanIssue.detectorId,
@@ -439,83 +455,91 @@ describe('CodeWhisperer-basicCommands', function () {
             })
         })
 
-        it('handles patch failure', async function () {
-            const textDocumentMock = createMockDocument()
+        it('should call applySecurityFix command successfully but not remove issues if auto-scans is disabled', async function () {
+            const fileName = 'sample.py'
+            const textDocumentMock = createMockDocument('first line\n second line\n fourth line', 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(SecurityIssueHoverProvider.instance, 'removeIssue').value(removeIssueMock)
+            sandbox.stub(SecurityIssueCodeActionProvider.instance, 'removeIssue').value(removeIssueMock)
+            await CodeScansState.instance.setScansEnabled(false)
+
             targetCommand = testCommand(applySecurityFix)
             const codeScanIssue = createCodeScanIssue({
                 suggestedFixes: [
                     {
-                        code: '@@ -1,1 -1,1 @@\n-mock\n+line5',
-                        description: 'dummy',
+                        description: 'fix',
+                        code: '@@ -1,3 +1,3 @@\n first line\n- second line\n+ third line\n  fourth line',
                     },
                 ],
             })
-            await targetCommand.execute(codeScanIssue, 'test.py', 'webview')
+            await targetCommand.execute(codeScanIssue, fileName, 'hover')
+            assert.ok(
+                replaceMock.calledOnceWith(
+                    textDocumentMock.uri,
+                    new vscode.Range(0, 0, 2, 12),
+                    'first line\n third line\n fourth line'
+                )
+            )
+            assert.ok(applyEditMock.calledOnce)
+            assert.ok(removeDiagnosticMock.notCalled)
+            assert.ok(removeIssueMock.notCalled)
 
-            assert.strictEqual(getTestWindow().shownMessages[0].message, 'Failed to apply suggested code fix.')
             assertTelemetry('codewhisperer_codeScanIssueApplyFix', {
                 detectorId: codeScanIssue.detectorId,
                 findingId: codeScanIssue.findingId,
-                component: 'webview',
-                result: 'Failed',
-                reason: 'Error: Failed to get updated content from applying diff patch',
+                component: 'hover',
+                result: 'Succeeded',
             })
         })
 
-        it('handles document save failure', async function () {
-            const fileName = 'sample.py'
-            saveStub.resolves(false)
-            const textDocumentMock = new MockDocument('first line\n second line\n fourth line', fileName, saveStub)
+        it('handles patch failure', async function () {
+            const textDocumentMock = createMockDocument()
 
             openTextDocumentMock.resolves(textDocumentMock)
 
             sandbox.stub(vscode.workspace, 'openTextDocument').value(openTextDocumentMock)
-            const loggerStub = sinon.stub(getLogger(), 'error')
 
             targetCommand = testCommand(applySecurityFix)
             const codeScanIssue = createCodeScanIssue({
                 suggestedFixes: [
                     {
-                        description: 'fix',
-                        code: '@@ -1,3 +1,3 @@\n first line\n- second line\n+ third line\n  fourth line',
+                        code: '@@ -1,1 -1,1 @@\n-mock\n+line5',
+                        description: 'dummy',
                     },
                 ],
             })
-            await targetCommand.execute(codeScanIssue, fileName, 'quickfix')
+            await targetCommand.execute(codeScanIssue, 'test.py', 'webview')
 
-            assert.ok(saveStub.calledOnce)
-            assert.ok(loggerStub.calledOnce)
-            const actual = loggerStub.getCall(0).args[0]
-            assert.strictEqual(
-                actual,
-                'Apply fix command failed. Error: Failed to save editor text changes into the file.'
-            )
+            assert.strictEqual(getTestWindow().shownMessages[0].message, 'Failed to apply suggested code fix.')
             assertTelemetry('codewhisperer_codeScanIssueApplyFix', {
                 detectorId: codeScanIssue.detectorId,
                 findingId: codeScanIssue.findingId,
-                component: 'quickfix',
+                component: 'webview',
                 result: 'Failed',
-                reason: 'Error: Failed to save editor text changes into the file.',
+                reason: 'Error: Failed to get updated content from applying diff patch',
             })
         })
 
-        it('handles document write failure', async function () {
+        it('handles apply edit failure', async function () {
             const fileName = 'sample.py'
-            saveStub.resolves(true)
-            const textDocumentMock = new MockDocument('first line\n second line\n fourth line', fileName, saveStub)
+            const textDocumentMock = createMockDocument('first line\n second line\n fourth line', fileName)
 
             openTextDocumentMock.resolves(textDocumentMock)
-            writeFileMock.rejects('Error: Writing to file failed.')
 
             sandbox.stub(vscode.workspace, 'openTextDocument').value(openTextDocumentMock)
-            sinon.stub(FileSystemCommon.prototype, 'writeFile').value(writeFileMock)
             const loggerStub = sinon.stub(getLogger(), 'error')
 
+            sinon.stub(vscode.WorkspaceEdit.prototype, 'replace').value(replaceMock)
+            applyEditMock.resolves(false)
+            sinon.stub(vscode.workspace, 'applyEdit').value(applyEditMock)
+
             targetCommand = testCommand(applySecurityFix)
             const codeScanIssue = createCodeScanIssue({
                 suggestedFixes: [
@@ -525,18 +549,18 @@ describe('CodeWhisperer-basicCommands', function () {
                     },
                 ],
             })
-            await targetCommand.execute(codeScanIssue, fileName, 'hover')
+            await targetCommand.execute(codeScanIssue, fileName, 'quickfix')
 
-            assert.ok(saveStub.calledOnce)
+            assert.ok(replaceMock.calledOnce)
             assert.ok(loggerStub.calledOnce)
             const actual = loggerStub.getCall(0).args[0]
-            assert.strictEqual(actual, 'Apply fix command failed. Error: Writing to file failed.')
+            assert.strictEqual(actual, 'Apply fix command failed. Error: Failed to apply edit to the workspace.')
             assertTelemetry('codewhisperer_codeScanIssueApplyFix', {
                 detectorId: codeScanIssue.detectorId,
                 findingId: codeScanIssue.findingId,
-                component: 'hover',
+                component: 'quickfix',
                 result: 'Failed',
-                reason: 'Error: Writing to file failed.',
+                reason: 'Error: Failed to apply edit to the workspace.',
             })
         })
     })
