Skip to content

Commit ce0972b

Browse files
HweinstockHweinstock
committed
refactor: move scoping logic up the call chain
1 parent 41e1055 commit ce0972b

File tree

6 files changed

+111
-103
lines changed

6 files changed

+111
-103
lines changed

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@ import { telemetry } from '../../shared/telemetry/telemetry'
3535
import { ToolkitError, getTelemetryReasonDesc, isAwsError } from '../../shared/errors'
3636
import { AuthUtil } from '../util/authUtil'
3737
import path from 'path'
38-
import { ZipMetadata, ZipUtil } from '../util/zipUtil'
38+
import { getPrefixFromUseCase, ZipMetadata, ZipUtil } from '../util/zipUtil'
3939
import { debounce } from 'lodash'
4040
import { once } from '../../shared/utilities/functionUtils'
4141
import { randomUUID } from '../../shared/crypto'
@@ -106,7 +106,7 @@ export async function startSecurityScan(
106106
context: vscode.ExtensionContext,
107107
scope: CodeWhispererConstants.CodeAnalysisScope,
108108
initiatedByChat: boolean,
109-
zipUtil: ZipUtil = new ZipUtil(),
109+
zipUtil: ZipUtil = new ZipUtil(getPrefixFromUseCase(CodeWhispererConstants.FeatureUseCase.CODE_SCAN)),
110110
scanUuid?: string
111111
) {
112112
const logger = getLoggerForScope(scope)

‎packages/core/src/codewhisperer/commands/startTestGeneration.ts‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
*/
55

66
import { getLogger } from '../../shared/logger/logger'
7-
import { ZipUtil } from '../util/zipUtil'
7+
import { getPrefixFromUseCase, ZipUtil } from '../util/zipUtil'
88
import { ArtifactMap } from '../client/codewhisperer'
99
import { testGenerationLogsDir } from '../../shared/filesystemUtilities'
1010
import {
@@ -22,6 +22,7 @@ import { BuildStatus } from '../../amazonqTest/chat/session/session'
2222
import { fs } from '../../shared/fs/fs'
2323
import { Range } from '../client/codewhispereruserclient'
2424
import { getWorkspaceForFile } from '../../shared/utilities/workspaceUtils'
25+
import { FeatureUseCase } from '../indexNode'
2526

2627
// eslint-disable-next-line unicorn/no-null
2728
let spawnResult: ChildProcess | null = null
@@ -47,7 +48,7 @@ export async function startTestGenerationProcess(
4748
* Step 1: Zip the project
4849
*/
4950

50-
const zipUtil = new ZipUtil()
51+
const zipUtil = new ZipUtil(getPrefixFromUseCase(FeatureUseCase.TEST_GENERATION))
5152
if (initialExecution) {
5253
const projectPath = getWorkspaceForFile(filePath) ?? ''
5354
const relativeTargetPath = path.relative(projectPath, filePath)

‎packages/core/src/codewhisperer/util/zipUtil.ts‎

Lines changed: 60 additions & 53 deletions
Original file line numberDiff line numberDiff line change
@@ -56,28 +56,28 @@ export class ZipUtil {
5656
protected _pickedSourceFiles: Set<string> = new Set<string>()
5757
protected _pickedBuildFiles: Set<string> = new Set<string>()
5858
protected _totalSize: number = 0
59+
protected _zipDir: string = ''
5960
protected _totalBuildSize: number = 0
6061
protected _tmpDir: string = tempDirPath
61-
protected _zipDir: string = ''
6262
protected _totalLines: number = 0
6363
protected _fetchedDirs: Set<string> = new Set<string>()
6464
protected _language: CodewhispererLanguage | undefined
6565
protected _timestamp: string = Date.now().toString()
66-
protected _payloadByteLimits = {
66+
protected static _payloadByteLimits = {
6767
file: CodeWhispererConstants.fileScanPayloadSizeLimitBytes,
6868
project: CodeWhispererConstants.projectScanPayloadSizeLimitBytes,
6969
}
70-
constructor() {}
70+
constructor(protected _zipDirPrefix: string) {}
7171

72-
public aboveByteLimit(size: number, limitType: ZipType): boolean {
72+
public static aboveByteLimit(size: number, limitType: ZipType): boolean {
7373
return size > this._payloadByteLimits[limitType]
7474
}
7575

76-
public willReachProjectByteLimit(current: number, adding: number): boolean {
77-
return this.aboveByteLimit(current + adding, 'project')
76+
public static willReachProjectByteLimit(current: number, adding: number): boolean {
77+
return ZipUtil.aboveByteLimit(current + adding, 'project')
7878
}
7979

80-
protected async zipFile(uri: vscode.Uri | undefined, scope: CodeWhispererConstants.CodeAnalysisScope) {
80+
protected async zipFile(uri: vscode.Uri | undefined, includeGitDiffHeader?: boolean) {
8181
if (!uri) {
8282
throw new NoActiveFileError()
8383
}
@@ -95,7 +95,7 @@ export class ZipUtil {
9595
const zipEntryPath = this.getZipEntryPath(projectName, relativePath)
9696
zip.writeString(content, zipEntryPath)
9797

98-
if (scope === CodeWhispererConstants.CodeAnalysisScope.FILE_ON_DEMAND) {
98+
if (includeGitDiffHeader) {
9999
const gitDiffContent = `+++ b/${normalize(zipEntryPath)}` // Sending file path in payload for LLM code review
100100
zip.writeString(gitDiffContent, ZipConstants.codeDiffFilePath)
101101
}
@@ -107,7 +107,7 @@ export class ZipUtil {
107107
this._totalSize += (await fs.stat(uri.fsPath)).size
108108
this._totalLines += content.split(ZipConstants.newlineRegex).length
109109

110-
if (this.aboveByteLimit(this._totalSize, 'file')) {
110+
if (ZipUtil.aboveByteLimit(this._totalSize, 'file')) {
111111
throw new FileSizeExceededError()
112112
}
113113
const zipFilePath = this.getZipDirPath(FeatureUseCase.CODE_SCAN) + CodeWhispererConstants.codeScanZipExt
@@ -186,11 +186,18 @@ export class ZipUtil {
186186
projectPaths = getWorkspacePaths()
187187
}
188188
if (useCase === FeatureUseCase.CODE_SCAN) {
189-
await this.processCombinedGitDiff(zip, projectPaths, '', CodeWhispererConstants.CodeAnalysisScope.PROJECT)
189+
await this.processCombinedGitDiff(zip, projectPaths, '')
190190
}
191191
const languageCount = new Map<CodewhispererLanguage, number>()
192192

193-
await this.processSourceFiles(zip, languageCount, projectPaths, useCase)
193+
await this.processSourceFiles(
194+
zip,
195+
languageCount,
196+
projectPaths,
197+
this.getWorkspaceFolders(useCase),
198+
this.getExcludePatterns(useCase),
199+
useCase === FeatureUseCase.TEST_GENERATION
200+
)
194201
if (metadataDir) {
195202
await this.processMetadataDir(zip, metadataDir)
196203
}
@@ -207,51 +214,51 @@ export class ZipUtil {
207214
return zipFilePath
208215
}
209216

210-
protected async processCombinedGitDiff(
211-
zip: ZipStream,
212-
projectPaths: string[],
213-
filePath?: string,
214-
scope?: CodeWhispererConstants.CodeAnalysisScope
215-
) {
217+
protected async processCombinedGitDiff(zip: ZipStream, projectPaths: string[], filePath?: string) {
216218
const gitDiffContent = await getGitDiffContentForProjects(projectPaths, filePath)
217219
if (gitDiffContent) {
218220
zip.writeString(gitDiffContent, ZipConstants.codeDiffFilePath)
219221
}
220222
}
221223

224+
// TODO: remove
225+
private getWorkspaceFolders(useCase: FeatureUseCase) {
226+
return (
227+
useCase === FeatureUseCase.TEST_GENERATION
228+
? [...(vscode.workspace.workspaceFolders ?? [])].sort(
229+
(a, b) => b.uri.fsPath.length - a.uri.fsPath.length
230+
)
231+
: vscode.workspace.workspaceFolders
232+
) as CurrentWsFolders
233+
}
234+
235+
private getExcludePatterns(useCase: FeatureUseCase) {
236+
return useCase === FeatureUseCase.TEST_GENERATION
237+
? [...CodeWhispererConstants.testGenExcludePatterns, ...defaultExcludePatterns]
238+
: defaultExcludePatterns
239+
}
240+
222241
protected async processSourceFiles(
223242
zip: ZipStream,
224243
languageCount: Map<CodewhispererLanguage, number>,
225244
projectPaths: string[] | undefined,
226-
useCase: FeatureUseCase
245+
workspaceFolders: CurrentWsFolders,
246+
excludePatterns: string[],
247+
includeBinary?: boolean
227248
) {
228249
if (!projectPaths || projectPaths.length === 0) {
229250
return
230251
}
231252

232-
const sourceFiles = await collectFiles(
233-
projectPaths,
234-
(useCase === FeatureUseCase.TEST_GENERATION
235-
? [...(vscode.workspace.workspaceFolders ?? [])].sort(
236-
(a, b) => b.uri.fsPath.length - a.uri.fsPath.length
237-
)
238-
: vscode.workspace.workspaceFolders) as CurrentWsFolders,
239-
{
240-
maxTotalSizeBytes: this._payloadByteLimits['project'],
241-
excludePatterns:
242-
useCase === FeatureUseCase.TEST_GENERATION
243-
? [...CodeWhispererConstants.testGenExcludePatterns, ...defaultExcludePatterns]
244-
: defaultExcludePatterns,
245-
}
246-
)
253+
const sourceFiles = await collectFiles(projectPaths, workspaceFolders, {
254+
maxTotalSizeBytes: ZipUtil._payloadByteLimits['project'],
255+
excludePatterns,
256+
})
247257
for (const file of sourceFiles) {
248258
const projectName = path.basename(file.workspaceFolder.uri.fsPath)
249259
const zipEntryPath = this.getZipEntryPath(projectName, file.relativeFilePath)
250260

251-
if (ZipConstants.knownBinaryFileExts.includes(path.extname(file.fileUri.fsPath))) {
252-
if (useCase === FeatureUseCase.TEST_GENERATION) {
253-
continue
254-
}
261+
if (ZipConstants.knownBinaryFileExts.includes(path.extname(file.fileUri.fsPath)) && includeBinary) {
255262
await this.processBinaryFile(zip, file.fileUri, zipEntryPath)
256263
} else {
257264
const isFileOpenAndDirty = this.isFileOpenAndDirty(file.fileUri)
@@ -300,8 +307,8 @@ export class ZipUtil {
300307
const fileSize = Buffer.from(fileContent).length
301308

302309
if (
303-
this.aboveByteLimit(this._totalSize, 'project') ||
304-
this.willReachProjectByteLimit(this._totalSize, fileSize)
310+
ZipUtil.aboveByteLimit(this._totalSize, 'project') ||
311+
ZipUtil.willReachProjectByteLimit(this._totalSize, fileSize)
305312
) {
306313
throw new ProjectSizeExceededError()
307314
}
@@ -317,8 +324,8 @@ export class ZipUtil {
317324
const fileSize = (await fs.stat(uri.fsPath)).size
318325

319326
if (
320-
this.aboveByteLimit(this._totalSize, 'project') ||
321-
this.willReachProjectByteLimit(this._totalSize, fileSize)
327+
ZipUtil.aboveByteLimit(this._totalSize, 'project') ||
328+
ZipUtil.willReachProjectByteLimit(this._totalSize, fileSize)
322329
) {
323330
throw new ProjectSizeExceededError()
324331
}
@@ -342,12 +349,12 @@ export class ZipUtil {
342349

343350
public getZipDirPath(useCase: FeatureUseCase): string {
344351
if (this._zipDir === '') {
345-
const prefix =
346-
useCase === FeatureUseCase.TEST_GENERATION
347-
? CodeWhispererConstants.TestGenerationTruncDirPrefix
348-
: CodeWhispererConstants.codeScanTruncDirPrefix
352+
// const prefix =
353+
// useCase === FeatureUseCase.TEST_GENERATION
354+
// ? CodeWhispererConstants.TestGenerationTruncDirPrefix
355+
// : CodeWhispererConstants.codeScanTruncDirPrefix
349356

350-
this._zipDir = path.join(this._tmpDir, `${prefix}_${this._timestamp}`)
357+
this._zipDir = path.join(this._tmpDir, `${this._zipDirPrefix}_${this._timestamp}`)
351358
}
352359
return this._zipDir
353360
}
@@ -363,7 +370,7 @@ export class ZipUtil {
363370
scope === CodeWhispererConstants.CodeAnalysisScope.FILE_AUTO ||
364371
scope === CodeWhispererConstants.CodeAnalysisScope.FILE_ON_DEMAND
365372
) {
366-
zipFilePath = await this.zipFile(uri, scope)
373+
zipFilePath = await this.zipFile(uri, scope === CodeWhispererConstants.CodeAnalysisScope.FILE_AUTO)
367374
} else if (scope === CodeWhispererConstants.CodeAnalysisScope.PROJECT) {
368375
zipFilePath = await this.zipProject(FeatureUseCase.CODE_SCAN)
369376
} else {
@@ -390,7 +397,6 @@ export class ZipUtil {
390397

391398
public async generateZipTestGen(projectPath: string, initialExecution: boolean): Promise<ZipMetadata> {
392399
try {
393-
// const repoMapFile = await LspClient.instance.getRepoMapJSON()
394400
const zipDirPath = this.getZipDirPath(FeatureUseCase.TEST_GENERATION)
395401

396402
const metadataDir = path.join(zipDirPath, 'utgRequiredArtifactsDir')
@@ -404,11 +410,6 @@ export class ZipUtil {
404410
}
405411
await Promise.all(Object.values(dirs).map((dir) => fs.mkdir(dir)))
406412

407-
// if (await fs.exists(repoMapFile)) {
408-
// await fs.copy(repoMapFile, path.join(dirs.repoMapDir, 'repoMapData.json'))
409-
// await fs.delete(repoMapFile)
410-
// }
411-
412413
if (!initialExecution) {
413414
await this.processTestCoverageFiles(dirs.testCoverageDir)
414415

@@ -456,6 +457,12 @@ interface GitDiffOptions {
456457
scope?: CodeWhispererConstants.CodeAnalysisScope
457458
}
458459

460+
export function getPrefixFromUseCase(useCase: FeatureUseCase) {
461+
return useCase === FeatureUseCase.TEST_GENERATION
462+
? CodeWhispererConstants.TestGenerationTruncDirPrefix
463+
: CodeWhispererConstants.codeScanTruncDirPrefix
464+
}
465+
459466
async function getGitDiffContentForProjects(
460467
projectPaths: string[],
461468
filepath?: string,

0 commit comments

Comments
 (0)