codewhisperer: fix bugs for proactive scans aws#4768

Problem
Fix remaining bugs in `feature/cw-proactive-scan` branch

Solution
- Add changelogs
- Update telemetry package and pass scope field to `codewhisperer_securityScan`
- Don't run file scans if auth is not connected
- Don't run file scans if monthly quota is exceeded

Author: ctlai95

.changes/next-release/Feature-29598070-7516-4255-8981-a3db0c4cb8de.json

@@ -0,0 +1,4 @@
+{
+	"type": "Feature",
+	"description": "CodeWhisperer: Security scans can now run on all files in the project"
+}

.changes/next-release/Feature-c45c0979-bb7a-43dd-b1e2-d046667ebbb1.json

@@ -0,0 +1,4 @@
+{
+	"type": "Feature",
+	"description": "CodeWhisperer: Security scans can now run automatically when file changes are made"
+}

package-lock.json

@@ -4245,7 +4245,7 @@
     },
     "devDependencies": {
         "@aws-sdk/types": "^3.13.1",
-        "@aws-toolkits/telemetry": "^1.0.201",
+        "@aws-toolkits/telemetry": "^1.0.203",
         "@aws/fully-qualified-names": "^2.1.4",
         "@cspotcode/source-map-support": "^0.8.1",
         "@sinonjs/fake-timers": "^10.0.2",

packages/core/package.json

@@ -96,11 +96,6 @@ export async function activate(context: ExtContext): Promise<void> {
      */
     const client = new codewhispererClient.DefaultCodeWhispererClient()
 
-    /**
-     * CodeWhisperer auto scans
-     */
-    setSubscriptionsForAutoScans()
-
     // Service initialization
     const container = Container.instance
     ReferenceInlineProvider.instance
@@ -286,11 +281,18 @@ export async function activate(context: ExtContext): Promise<void> {
         await CodeScansState.instance.setScansEnabled(false)
     }
 
+    /**
+     * CodeWhisperer auto scans
+     */
+    setSubscriptionsForAutoScans()
+
     function setSubscriptionsForAutoScans() {
         // Initial scan when the editor opens for the first time
         const editor = vscode.window.activeTextEditor
         if (
             CodeScansState.instance.isScansEnabled() &&
+            !CodeScansState.instance.isMonthlyQuotaExceeded() &&
+            auth.isConnected() &&
             !auth.isBuilderIdInUse() &&
             editor &&
             securityScanLanguageContext.isLanguageSupported(editor.document.languageId) &&
@@ -310,6 +312,8 @@ export async function activate(context: ExtContext): Promise<void> {
             vscode.window.onDidChangeActiveTextEditor(editor => {
                 if (
                     CodeScansState.instance.isScansEnabled() &&
+                    !CodeScansState.instance.isMonthlyQuotaExceeded() &&
+                    auth.isConnected() &&
                     !auth.isBuilderIdInUse() &&
                     editor &&
                     securityScanLanguageContext.isLanguageSupported(editor.document.languageId)
@@ -339,6 +343,8 @@ export async function activate(context: ExtContext): Promise<void> {
                 const editor = vscode.window.activeTextEditor
                 if (
                     CodeScansState.instance.isScansEnabled() &&
+                    !CodeScansState.instance.isMonthlyQuotaExceeded() &&
+                    auth.isConnected() &&
                     !auth.isBuilderIdInUse() &&
                     editor &&
                     event.document === editor.document &&
@@ -361,6 +367,8 @@ export async function activate(context: ExtContext): Promise<void> {
             const editor = vscode.window.activeTextEditor
             if (
                 isScansEnabled &&
+                !CodeScansState.instance.isMonthlyQuotaExceeded() &&
+                auth.isConnected() &&
                 !auth.isBuilderIdInUse() &&
                 editor &&
                 securityScanLanguageContext.isLanguageSupported(editor.document.languageId) &&

packages/core/src/codewhisperer/activation.ts

@@ -20,7 +20,7 @@ import {
     throwIfCancelled,
 } from '../service/securityScanHandler'
 import { runtimeLanguageContext } from '../util/runtimeLanguageContext'
-import { AggregatedCodeScanIssue, codeScanState, CodeScanTelemetryEntry } from '../models/model'
+import { AggregatedCodeScanIssue, CodeScansState, codeScanState, CodeScanTelemetryEntry } from '../models/model'
 import { cancel, ok } from '../../shared/localizedText'
 import { getFileExt } from '../util/commonUtil'
 import { getDirSize } from '../../shared/filesystemUtilities'
@@ -33,6 +33,7 @@ import { ZipMetadata, ZipUtil } from '../util/zipUtil'
 import { debounce } from 'lodash'
 import { once } from '../../shared/utilities/functionUtils'
 import { randomUUID } from '../../common/crypto'
+import { CodeAnalysisScope } from '../models/constants'
 
 const localize = nls.loadMessageBundle()
 export const stopScanButton = localize('aws.codewhisperer.stopscan', 'Stop Scan')
@@ -109,6 +110,7 @@ export async function startSecurityScan(
         codewhispererCodeScanTotalIssues: 0,
         codewhispererCodeScanIssuesWithFixes: 0,
         credentialStartUrl: AuthUtil.instance.startUrl,
+        codewhispererCodeScanScope: scope,
     }
     try {
         getLogger().verbose(`Starting security scan `)
@@ -219,14 +221,20 @@ export async function startSecurityScan(
             codeScanTelemetryEntry.result = 'Failed'
         }
 
-        if (isAwsError(error)) {
+        if (isAwsError(error) && error.code === 'ThrottlingException') {
             if (
-                error.code === 'ThrottlingException' &&
-                error.message.includes(CodeWhispererConstants.throttlingMessage)
+                scope === CodeAnalysisScope.PROJECT &&
+                error.message.includes(CodeWhispererConstants.projectScansThrottlingMessage)
             ) {
-                void vscode.window.showErrorMessage(CodeWhispererConstants.freeTierLimitReachedCodeScan)
+                void vscode.window.showErrorMessage(CodeWhispererConstants.projectScansLimitReached)
                 // TODO: Should we set a graphical state?
                 // We shouldn't set vsCodeState.isFreeTierLimitReached here because it will hide CW and Q chat options.
+            } else if (
+                scope === CodeAnalysisScope.FILE &&
+                error.message.includes(CodeWhispererConstants.fileScansThrottlingMessage)
+            ) {
+                void vscode.window.showErrorMessage(CodeWhispererConstants.fileScansLimitReached)
+                CodeScansState.instance.setMonthlyQuotaExceeded()
             }
         }
         codeScanTelemetryEntry.reason = (error as Error).message
@@ -298,7 +306,6 @@ function populateCodeScanLogStream(scannedFiles: Set<string>) {
         const uri = vscode.Uri.file(file)
         codeScanLogger.info(`File scanned: ${uri.fsPath}`)
     }
-    codeScanOutpuChan.show()
 }
 
 export async function confirmStopSecurityScan() {

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

@@ -290,10 +290,18 @@ export const freeTierLimitReached = 'You have reached the monthly fair use limit
 
 export const freeTierLimitReachedCodeScan = 'You have reached the monthly quota of code scans.'
 
+export const fileScansLimitReached = 'You have reached the monthly quota of automatic file scans.'
+
+export const projectScansLimitReached = 'You have reached the monthly quota of full project scans.'
+
 export const throttlingLearnMore = `Learn More`
 
 export const throttlingMessage = `Maximum recommendation count reached for this month`
 
+export const fileScansThrottlingMessage = `Maximum automatic file scan count reached for this month`
+
+export const projectScansThrottlingMessage = `Maximum full project scan count reached for this month`
+
 export const connectionChangeMessage = `Keep using CodeWhisperer with `
 
 // TODO: align this text with service side

packages/core/src/codewhisperer/models/constants.ts

@@ -6,6 +6,7 @@ import * as vscode from 'vscode'
 import { ToolkitError } from '../../shared/errors'
 import { getIcon } from '../../shared/icons'
 import {
+    CodewhispererCodeScanScope,
     CodewhispererCompletionType,
     CodewhispererLanguage,
     CodewhispererTriggerType,
@@ -119,6 +120,8 @@ export class CodeScansState {
     /** Set a callback for when state of code scans changes */
     onDidChangeState = this.#onDidChangeState.event
 
+    private exceedsMonthlyQuota = false
+
     static #instance: CodeScansState
     static get instance() {
         return (this.#instance ??= new this())
@@ -147,6 +150,14 @@ export class CodeScansState {
         const isEnabled = get(autoScansEnabledKey, this.#context)
         return isEnabled !== undefined ? isEnabled : this.#fallback
     }
+
+    setMonthlyQuotaExceeded() {
+        this.exceedsMonthlyQuota = true
+    }
+
+    isMonthlyQuotaExceeded() {
+        return this.exceedsMonthlyQuota
+    }
 }
 
 export interface AcceptedSuggestionEntry {
@@ -626,6 +637,7 @@ export interface CodeScanTelemetryEntry {
     codewhispererCodeScanTotalIssues: number
     codewhispererCodeScanIssuesWithFixes: number
     credentialStartUrl: string | undefined
+    codewhispererCodeScanScope: CodewhispererCodeScanScope
 }
 
 export interface RecommendationDescription {

packages/core/src/codewhisperer/models/model.ts

@@ -22,7 +22,7 @@ import {
     toggleCodeScans,
 } from '../commands/basicCommands'
 import { CodeWhispererCommandDeclarations } from '../commands/gettingStartedPageCommands'
-import { codeScanState } from '../models/model'
+import { CodeScansState, codeScanState } from '../models/model'
 import { getNewCustomizationsAvailable, getSelectedCustomization } from '../util/customizationUtil'
 import { cwQuickPickSource, cwTreeNodeSource } from '../commands/types'
 import { AuthUtil } from '../util/authUtil'
@@ -48,11 +48,12 @@ export function createAutoScans(pause: boolean): DataQuickPickItem<'autoScans'>
     const iconResume = getIcon('vscode-debug-alt')
     const labelPause = localize('AWS.codewhisperer.pauseCodeWhispererNode.label', 'Pause Automatic File Scanning')
     const iconPause = getIcon('vscode-debug-pause')
+    const monthlyQuotaExceeded = CodeScansState.instance.isMonthlyQuotaExceeded()
 
     return {
         data: 'autoScans',
         label: pause ? codicon`${iconPause} ${labelPause}` : codicon`${iconResume} ${labelResume}`,
-        description: pause ? 'RUNNING' : 'PAUSED',
+        description: monthlyQuotaExceeded ? 'Monthly quota exceeded' : pause ? 'RUNNING' : 'PAUSED',
         onClick: () => toggleCodeScans.execute(placeholder, cwQuickPickSource),
     } as DataQuickPickItem<'autoScans'>
 }