Merge master into feature/q-transform-staging

aws-toolkit-automation · web-flow · commit e22a0a94bb86 · 2024-05-02T15:28:32.000-07:00
diff --git a/packages/core/src/codewhisperer/commands/startSecurityScan.ts b/packages/core/src/codewhisperer/commands/startSecurityScan.ts
@@ -7,7 +7,7 @@ import * as vscode from 'vscode'
 import * as nls from 'vscode-nls'
 import { ArtifactMap, DefaultCodeWhispererClient } from '../client/codewhisperer'
 import { isCloud9 } from '../../shared/extensionUtilities'
-import { initSecurityScanRender, securityScanRender } from '../service/diagnosticsProvider'
+import { initSecurityScanRender } from '../service/diagnosticsProvider'
 import { SecurityPanelViewProvider } from '../views/securityPanelViewProvider'
 import { getLogger } from '../../shared/logger'
 import { makeLogger } from '../../shared/logger/activation'
@@ -21,7 +21,13 @@ import {
     getLoggerForScope,
 } from '../service/securityScanHandler'
 import { runtimeLanguageContext } from '../util/runtimeLanguageContext'
-import { AggregatedCodeScanIssue, CodeScansState, codeScanState, CodeScanTelemetryEntry } from '../models/model'
+import {
+    AggregatedCodeScanIssue,
+    CodeScansState,
+    codeScanState,
+    CodeScanStoppedError,
+    CodeScanTelemetryEntry,
+} from '../models/model'
 import { cancel, ok } from '../../shared/localizedText'
 import { getFileExt } from '../util/commonUtil'
 import { getDirSize } from '../../shared/filesystemUtilities'
@@ -96,6 +102,9 @@ export async function startSecurityScan(
      * Step 0: Initial Code Scan telemetry
      */
     const codeScanStartTime = performance.now()
+    if (scope === CodeAnalysisScope.FILE) {
+        CodeScansState.instance.setLatestScanTime(codeScanStartTime)
+    }
     let serviceInvocationStartTime = 0
     const codeScanTelemetryEntry: CodeScanTelemetryEntry = {
         codewhispererLanguage: runtimeLanguageContext.getLanguageContext(
@@ -120,7 +129,7 @@ export async function startSecurityScan(
         /**
          * Step 1: Generate zip
          */
-        throwIfCancelled(scope)
+        throwIfCancelled(scope, codeScanStartTime)
         const zipUtil = new ZipUtil()
         const zipMetadata = await zipUtil.generateZip(editor.document.uri, scope)
         const projectPath = zipUtil.getProjectPath(editor.document.uri)
@@ -139,7 +148,7 @@ export async function startSecurityScan(
         /**
          * Step 2: Get presigned Url, upload and clean up
          */
-        throwIfCancelled(scope)
+        throwIfCancelled(scope, codeScanStartTime)
         let artifactMap: ArtifactMap = {}
         const uploadStartTime = performance.now()
         const scanName = randomUUID()
@@ -153,7 +162,7 @@ export async function startSecurityScan(
         /**
          * Step 3:  Create scan job
          */
-        throwIfCancelled(scope)
+        throwIfCancelled(scope, codeScanStartTime)
         serviceInvocationStartTime = performance.now()
         const scanJob = await createScanJob(
             client,
@@ -171,16 +180,16 @@ export async function startSecurityScan(
         /**
          * Step 4:  Polling mechanism on scan job status
          */
-        throwIfCancelled(scope)
-        const jobStatus = await pollScanJobStatus(client, scanJob.jobId, scope)
+        throwIfCancelled(scope, codeScanStartTime)
+        const jobStatus = await pollScanJobStatus(client, scanJob.jobId, scope, codeScanStartTime)
         if (jobStatus === 'Failed') {
             throw new Error('Security scan job failed.')
         }
 
         /**
          * Step 5: Process and render scan results
          */
-        throwIfCancelled(scope)
+        throwIfCancelled(scope, codeScanStartTime)
         logger.verbose(`Security scan job succeeded and start processing result.`)
         const securityRecommendationCollection = await listScanResults(
             client,
@@ -198,27 +207,22 @@ export async function startSecurityScan(
         )
         codeScanTelemetryEntry.codewhispererCodeScanTotalIssues = total
         codeScanTelemetryEntry.codewhispererCodeScanIssuesWithFixes = withFixes
-        throwIfCancelled(scope)
+        throwIfCancelled(scope, codeScanStartTime)
         logger.verbose(`Security scan totally found ${total} issues. ${withFixes} of them have fixes.`)
-        if (codeScanStartTime > securityScanRender.lastUpdated) {
-            showSecurityScanResults(
-                securityPanelViewProvider,
-                securityRecommendationCollection,
-                editor,
-                context,
-                scope,
-                zipMetadata,
-                total,
-                codeScanStartTime
-            )
-        } else {
-            logger.verbose('Received issues from older scan, discarding the results')
-        }
+        showSecurityScanResults(
+            securityPanelViewProvider,
+            securityRecommendationCollection,
+            editor,
+            context,
+            scope,
+            zipMetadata,
+            total
+        )
 
         logger.verbose(`Security scan completed.`)
     } catch (error) {
         getLogger().error('Security scan failed.', error)
-        if (codeScanState.isCancelling()) {
+        if (error instanceof CodeScanStoppedError) {
             codeScanTelemetryEntry.result = 'Cancelled'
         } else {
             errorPromptHelper(error as Error, scope)
@@ -257,14 +261,13 @@ export function showSecurityScanResults(
     context: vscode.ExtensionContext,
     scope: CodeWhispererConstants.CodeAnalysisScope,
     zipMetadata: ZipMetadata,
-    totalIssues: number,
-    codeScanStartTime: number
+    totalIssues: number
 ) {
     if (isCloud9()) {
         securityPanelViewProvider.addLines(securityRecommendationCollection, editor)
         void vscode.commands.executeCommand('workbench.view.extension.aws-codewhisperer-security-panel')
     } else {
-        initSecurityScanRender(securityRecommendationCollection, context, editor, scope, codeScanStartTime)
+        initSecurityScanRender(securityRecommendationCollection, context, editor, scope)
         if (scope === CodeWhispererConstants.CodeAnalysisScope.PROJECT) {
             void vscode.commands.executeCommand('workbench.action.problems.focus')
         }
diff --git a/packages/core/src/codewhisperer/models/constants.ts b/packages/core/src/codewhisperer/models/constants.ts
@@ -232,13 +232,17 @@ export const codeFileScanJobTimeoutSeconds = 60 //1 minute
 
 export const projectSizeCalculateTimeoutSeconds = 10
 
-export const codeScanJobPollingIntervalSeconds = 1
+export const codeScanJobPollingIntervalSeconds = 5
+
+export const fileScanPollingDelaySeconds = 10
+
+export const projectScanPollingDelaySeconds = 30
 
 export const artifactTypeSource = 'SourceCode'
 
 export const codeScanFindingsSchema = 'codescan/findings/1.0'
 
-export const autoScanDebounceDelaySeconds = 2
+export const autoScanDebounceDelaySeconds = 5
 
 export const codewhispererDiagnosticSourceLabel = 'Amazon Q '
 
diff --git a/packages/core/src/codewhisperer/models/model.ts b/packages/core/src/codewhisperer/models/model.ts
@@ -121,6 +121,7 @@ export class CodeScansState {
     onDidChangeState = this.#onDidChangeState.event
 
     private exceedsMonthlyQuota = false
+    private latestScanTime: number | undefined = undefined
 
     static #instance: CodeScansState
     static get instance() {
@@ -158,6 +159,14 @@ export class CodeScansState {
     isMonthlyQuotaExceeded() {
         return this.exceedsMonthlyQuota
     }
+
+    setLatestScanTime(time: number) {
+        this.latestScanTime = time
+    }
+
+    getLatestScanTime() {
+        return this.latestScanTime
+    }
 }
 
 export interface AcceptedSuggestionEntry {
diff --git a/packages/core/src/codewhisperer/service/diagnosticsProvider.ts b/packages/core/src/codewhisperer/service/diagnosticsProvider.ts
@@ -12,21 +12,18 @@ import { CodeAnalysisScope, codewhispererDiagnosticSourceLabel } from '../models
 interface SecurityScanRender {
     securityDiagnosticCollection: vscode.DiagnosticCollection | undefined
     initialized: boolean
-    lastUpdated: number
 }
 
 export const securityScanRender: SecurityScanRender = {
     securityDiagnosticCollection: undefined,
     initialized: false,
-    lastUpdated: 0,
 }
 
 export function initSecurityScanRender(
     securityRecommendationList: AggregatedCodeScanIssue[],
     context: vscode.ExtensionContext,
     editor: vscode.TextEditor,
-    scope: CodeAnalysisScope,
-    codeScanStartTime: number
+    scope: CodeAnalysisScope
 ) {
     securityScanRender.initialized = false
     if (scope === CodeAnalysisScope.FILE) {
@@ -39,7 +36,6 @@ export function initSecurityScanRender(
         updateSecurityIssueHoverAndCodeActions(securityRecommendation)
     })
     securityScanRender.initialized = true
-    securityScanRender.lastUpdated = codeScanStartTime
 }
 
 function updateSecurityIssueHoverAndCodeActions(securityRecommendation: AggregatedCodeScanIssue) {
diff --git a/packages/core/src/codewhisperer/service/securityScanHandler.ts b/packages/core/src/codewhisperer/service/securityScanHandler.ts
@@ -99,14 +99,18 @@ function mapToAggregatedList(codeScanIssueMap: Map<string, RawCodeScanIssue[]>,
 export async function pollScanJobStatus(
     client: DefaultCodeWhispererClient,
     jobId: string,
-    scope: CodeWhispererConstants.CodeAnalysisScope
+    scope: CodeWhispererConstants.CodeAnalysisScope,
+    codeScanStartTime: number
 ) {
+    // We don't expect to get results immediately, so sleep for some time initially to not make unnecessary calls
+    await sleep(getPollingDelayForScope(scope))
+
     const logger = getLoggerForScope(scope)
     logger.verbose(`Polling scan job status...`)
     let status: string = 'Pending'
     let timer: number = 0
     while (true) {
-        throwIfCancelled(scope)
+        throwIfCancelled(scope, codeScanStartTime)
         const req: codewhispererClient.GetCodeScanRequest = {
             jobId: jobId,
         }
@@ -118,7 +122,7 @@ export async function pollScanJobStatus(
             logger.verbose(`Complete Polling scan job status.`)
             break
         }
-        throwIfCancelled(scope)
+        throwIfCancelled(scope, codeScanStartTime)
         await sleep(CodeWhispererConstants.codeScanJobPollingIntervalSeconds * 1000)
         timer += CodeWhispererConstants.codeScanJobPollingIntervalSeconds
         const timeoutSeconds =
@@ -208,18 +212,23 @@ function getMd5(fileName: string) {
     return hasher.digest('base64')
 }
 
-export function throwIfCancelled(scope: CodeWhispererConstants.CodeAnalysisScope) {
+export function throwIfCancelled(scope: CodeWhispererConstants.CodeAnalysisScope, codeScanStartTime: number) {
     switch (scope) {
         case CodeWhispererConstants.CodeAnalysisScope.PROJECT:
             if (codeScanState.isCancelling()) {
                 throw new CodeScanStoppedError()
             }
             break
-        case CodeWhispererConstants.CodeAnalysisScope.FILE:
-            if (!CodeScansState.instance.isScansEnabled()) {
+        case CodeWhispererConstants.CodeAnalysisScope.FILE: {
+            const latestCodeScanStartTime = CodeScansState.instance.getLatestScanTime()
+            if (
+                !CodeScansState.instance.isScansEnabled() ||
+                (latestCodeScanStartTime && latestCodeScanStartTime > codeScanStartTime)
+            ) {
                 throw new CodeScanStoppedError()
             }
             break
+        }
         default:
             getLogger().warn(`Unknown code analysis scope: ${scope}`)
             break
@@ -263,3 +272,9 @@ export async function uploadArtifactToS3(
 export function getLoggerForScope(scope: CodeWhispererConstants.CodeAnalysisScope) {
     return scope === CodeWhispererConstants.CodeAnalysisScope.FILE ? getNullLogger() : getLogger()
 }
+
+function getPollingDelayForScope(scope: CodeWhispererConstants.CodeAnalysisScope) {
+    return scope === CodeWhispererConstants.CodeAnalysisScope.FILE
+        ? CodeWhispererConstants.fileScanPollingDelaySeconds
+        : CodeWhispererConstants.projectScanPollingDelaySeconds
+}
diff --git a/packages/core/src/test/codewhisperer/commands/startSecurityScan.test.ts b/packages/core/src/test/codewhisperer/commands/startSecurityScan.test.ts
@@ -296,21 +296,64 @@ describe('startSecurityScan', function () {
         } as CodewhispererSecurityScan)
     })
 
-    it('Should not show security scan results if a later scan already finished', async function () {
+    it('Should cancel a scan if a newer one has started', async function () {
         getFetchStubWithResponse({ status: 200, statusText: 'testing stub' })
-        const commandSpy = sinon.spy(vscode.commands, 'executeCommand')
-        const securityScanRenderSpy = sinon.spy(diagnosticsProvider, 'initSecurityScanRender')
-        diagnosticsProvider.securityScanRender.lastUpdated = Date.now() + 60
+        await model.CodeScansState.instance.setScansEnabled(true)
+
+        const scanPromise = startSecurityScan.startSecurityScan(
+            mockSecurityPanelViewProvider,
+            editor,
+            createClient(),
+            extensionContext,
+            CodeAnalysisScope.FILE
+        )
+        await startSecurityScan.startSecurityScan(
+            mockSecurityPanelViewProvider,
+            editor,
+            createClient(),
+            extensionContext,
+            CodeAnalysisScope.FILE
+        )
+        await scanPromise
+        assertTelemetry('codewhisperer_securityScan', [
+            {
+                result: 'Cancelled',
+                reason: 'Security scan stopped by user.',
+            },
+            {
+                result: 'Succeeded',
+            },
+        ])
+    })
 
+    it('Should not cancel a project scan if a file scan has started', async function () {
+        getFetchStubWithResponse({ status: 200, statusText: 'testing stub' })
         await model.CodeScansState.instance.setScansEnabled(true)
+
+        const scanPromise = startSecurityScan.startSecurityScan(
+            mockSecurityPanelViewProvider,
+            editor,
+            createClient(),
+            extensionContext,
+            CodeAnalysisScope.PROJECT
+        )
         await startSecurityScan.startSecurityScan(
             mockSecurityPanelViewProvider,
             editor,
             createClient(),
             extensionContext,
             CodeAnalysisScope.FILE
         )
-        assert.ok(commandSpy.neverCalledWith('workbench.action.problems.focus'))
-        assert.ok(securityScanRenderSpy.notCalled)
+        await scanPromise
+        assertTelemetry('codewhisperer_securityScan', [
+            {
+                result: 'Succeeded',
+                codewhispererCodeScanScope: 'FILE',
+            },
+            {
+                result: 'Succeeded',
+                codewhispererCodeScanScope: 'PROJECT',
+            },
+        ])
     })
 })
diff --git a/packages/core/src/testE2E/codewhisperer/securityScan.test.ts b/packages/core/src/testE2E/codewhisperer/securityScan.test.ts
@@ -88,6 +88,7 @@ describe('CodeWhisperer security scan', async function () {
     returns artifactMap, projectPath and codeScanName
     */
     async function securityJobSetup(editor: vscode.TextEditor) {
+        const codeScanStartTime = performance.now()
         const zipUtil = new ZipUtil()
         const uri = editor.document.uri
 
@@ -106,6 +107,7 @@ describe('CodeWhisperer security scan', async function () {
             artifactMap: artifactMap,
             projectPath: projectPath,
             codeScanName: codeScanName,
+            codeScanStartTime: codeScanStartTime,
         }
     }
 
@@ -130,7 +132,12 @@ describe('CodeWhisperer security scan', async function () {
             scope,
             securityJobSetupResult.codeScanName
         )
-        const jobStatus = await pollScanJobStatus(client, scanJob.jobId, scope)
+        const jobStatus = await pollScanJobStatus(
+            client,
+            scanJob.jobId,
+            scope,
+            securityJobSetupResult.codeScanStartTime
+        )
         const securityRecommendationCollection = await listScanResults(
             client,
             scanJob.jobId,
@@ -165,7 +172,12 @@ describe('CodeWhisperer security scan', async function () {
         )
 
         //get job status and result
-        const jobStatus = await pollScanJobStatus(client, scanJob.jobId, scope)
+        const jobStatus = await pollScanJobStatus(
+            client,
+            scanJob.jobId,
+            scope,
+            securityJobSetupResult.codeScanStartTime
+        )
         const securityRecommendationCollection = await listScanResults(
             client,
             scanJob.jobId,
