Skip to content

Commit 17dcd44

Browse files
ashishrp-awsashishrp-aws
authored
fix(amazonq): Adding payload manifest for payload limit errors. (#5324)
* Adding payload manifest for payload limit errors for debugging.
1 parent 970fee0 commit 17dcd44

File tree

9 files changed

+110
-71
lines changed

9 files changed

+110
-71
lines changed

‎plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/CodeWhispererCodeTestSession.kt

+2-2
Original file line numberDiff line numberDiff line change
@@ -75,12 +75,12 @@ class CodeWhispererCodeTestSession(val sessionContext: CodeTestSessionContext) {
7575

7676
LOG.debug {
7777
"Total size of source payload in KB: ${payloadContext.srcPayloadSize * 1.0 / TOTAL_BYTES_IN_KB} \n" +
78-
"Total size of build payload in KB: ${(payloadContext.buildPayloadSize ?: 0) * 1.0 / TOTAL_BYTES_IN_KB} \n" +
7978
"Total size of source zip file in KB: ${payloadContext.srcZipFileSize * 1.0 / TOTAL_BYTES_IN_KB} \n" +
8079
"Total number of lines included: ${payloadContext.totalLines} \n" +
8180
"Total number of files included in payload: ${payloadContext.totalFiles} \n" +
8281
"Total time taken for creating payload: ${payloadContext.totalTimeInMilliseconds * 1.0 / TOTAL_MILLIS_IN_SECOND} seconds\n" +
83-
"Payload context language: ${payloadContext.language}"
82+
"Payload context language: ${payloadContext.language}" +
83+
"Payload exceeded the limit: ${payloadContext.payloadLimitCrossed}"
8484
}
8585

8686
// 2 & 3. CreateUploadURL and upload the context.

‎plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/CodeWhispererUTGChatManager.kt

+4
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@ import software.aws.toolkits.jetbrains.services.amazonqCodeTest.session.Session
4040
import software.aws.toolkits.jetbrains.services.amazonqCodeTest.utils.combineBuildAndExecuteLogFiles
4141
import software.aws.toolkits.jetbrains.services.codemodernizer.utils.calculateTotalLatency
4242
import software.aws.toolkits.jetbrains.services.codewhisperer.codetest.CodeTestException
43+
import software.aws.toolkits.jetbrains.services.codewhisperer.codetest.fileTooLarge
4344
import software.aws.toolkits.jetbrains.services.codewhisperer.codetest.sessionconfig.CodeTestSessionConfig
4445
import software.aws.toolkits.jetbrains.services.codewhisperer.codetest.testGenStoppedError
4546
import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererClientAdaptor
@@ -99,6 +100,9 @@ class CodeWhispererUTGChatManager(val project: Project, private val cs: Coroutin
99100
session.srcZipFileSize = codeTestResponseContext.payloadContext.srcZipFileSize
100101
session.artifactUploadDuration = codeTestResponseContext.serviceInvocationContext.artifactsUploadDuration
101102
val path = codeTestResponseContext.currentFileRelativePath
103+
if (codeTestResponseContext.payloadContext.payloadLimitCrossed == true) {
104+
fileTooLarge()
105+
}
102106

103107
val createUploadUrlResponse = codeTestResponseContext.createUploadUrlResponse ?: return
104108
throwIfCancelled(session)

‎plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/codescan/CodeWhispererCodeScanSession.kt

-1
Original file line numberDiff line numberDiff line change
@@ -91,7 +91,6 @@ class CodeWhispererCodeScanSession(val sessionContext: CodeScanSessionContext) {
9191
if (isProjectScope()) {
9292
LOG.debug {
9393
"Total size of source payload in KB: ${payloadContext.srcPayloadSize * 1.0 / TOTAL_BYTES_IN_KB} \n" +
94-
"Total size of build payload in KB: ${(payloadContext.buildPayloadSize ?: 0) * 1.0 / TOTAL_BYTES_IN_KB} \n" +
9594
"Total size of source zip file in KB: ${payloadContext.srcZipFileSize * 1.0 / TOTAL_BYTES_IN_KB} \n" +
9695
"Total number of lines reviewed: ${payloadContext.totalLines} \n" +
9796
"Total number of files included in payload: ${payloadContext.totalFiles} \n" +

‎plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/codescan/sessionconfig/CodeScanSessionConfig.kt

+4-1
Original file line numberDiff line numberDiff line change
@@ -334,7 +334,8 @@ data class PayloadContext(
334334
val scannedFiles: List<VirtualFile>,
335335
val srcPayloadSize: Long,
336336
val srcZipFileSize: Long,
337-
val buildPayloadSize: Long? = null,
337+
val payloadManifest: Set<Pair<String, Long>>? = null,
338+
val payloadLimitCrossed: Boolean? = false,
338339
)
339340

340341
data class PayloadMetadata(
@@ -343,4 +344,6 @@ data class PayloadMetadata(
343344
val linesScanned: Long,
344345
val language: CodewhispererLanguage,
345346
val codeDiff: String? = null,
347+
val payloadManifest: Set<Pair<String, Long>>? = null,
348+
val payloadLimitCrossed: Boolean? = false,
346349
)

‎plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/codetest/CodeTestException.kt

+1-1
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@ open class CodeTestException(
1616
internal fun noFileOpenError(): Nothing =
1717
throw CodeTestException(message("codewhisperer.codescan.no_file_open"), "ProjectZipError")
1818

19-
internal fun fileTooLarge(): Nothing =
19+
fun fileTooLarge(): Nothing =
2020
throw CodeTestException(message("codewhisperer.codescan.file_too_large_telemetry"), "ProjectZipError")
2121

2222
internal fun cannotFindFile(errorMessage: String, filepath: String): Nothing =

‎plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/codetest/sessionconfig/CodeTestSessionConfig.kt

+99-51
Original file line numberDiff line numberDiff line change
@@ -101,15 +101,17 @@ class CodeTestSessionConfig(
101101
}
102102

103103
// Copy all the included source files to the source zip
104-
val srcZip = zipFiles(payloadMetadata.sourceFiles.map { Path.of(it) })
104+
val srcZip = zipFiles(payloadMetadata)
105105
val payloadContext = PayloadContext(
106106
payloadMetadata.language,
107107
payloadMetadata.linesScanned,
108108
payloadMetadata.sourceFiles.size,
109109
Instant.now().toEpochMilli() - start,
110110
payloadMetadata.sourceFiles.mapNotNull { Path.of(it).toFile().toVirtualFile() },
111111
payloadMetadata.payloadSize,
112-
srcZip.length()
112+
srcZip.length(),
113+
payloadMetadata.payloadManifest,
114+
payloadMetadata.payloadLimitCrossed,
113115
)
114116

115117
return Payload(payloadContext, srcZip)
@@ -129,16 +131,20 @@ class CodeTestSessionConfig(
129131
}
130132
}
131133

132-
private fun zipFiles(files: List<Path>): File = createTemporaryZipFile {
133-
files.forEach { file ->
134-
try {
135-
val relativePath = file.relativeTo(projectRoot.toNioPath())
136-
val projectBaseName = projectRoot.name
137-
val zipEntryPath = "$projectBaseName/${relativePath.toString().replace("\\", "/")}"
138-
LOG.debug { "Adding file to ZIP: $zipEntryPath" }
139-
it.putNextEntry(zipEntryPath, file)
140-
} catch (e: Exception) {
141-
cannotFindFile("Zipping error: ${e.message}", file.toString())
134+
private fun zipFiles(payloadMetadata: PayloadMetadata): File = createTemporaryZipFile {
135+
if (payloadMetadata.payloadLimitCrossed == false) {
136+
payloadMetadata.sourceFiles.forEach { file ->
137+
Path.of(file).let { path ->
138+
try {
139+
val relativePath = path.relativeTo(projectRoot.toNioPath())
140+
val projectBaseName = projectRoot.name
141+
val zipEntryPath = "$projectBaseName/${relativePath.toString().replace("\\", "/")}"
142+
LOG.debug { "Adding file to ZIP: $zipEntryPath" }
143+
it.putNextEntry(zipEntryPath, path)
144+
} catch (e: Exception) {
145+
cannotFindFile("Zipping error: ${e.message}", path.toString())
146+
}
147+
}
142148
}
143149
}
144150

@@ -160,6 +166,18 @@ class CodeTestSessionConfig(
160166
if (buildAndExecuteLogFile != null) {
161167
it.putNextEntry(Path.of(utgDir, buildAndExecuteLogDir, "buildAndExecuteLog").name, buildAndExecuteLogFile.inputStream)
162168
}
169+
if (payloadMetadata.payloadLimitCrossed == true) {
170+
// write payloadMetadata.payloadManifest into a json file in repoMapData directory. manifest is Set<Pair<String, Long>> write into file first or ina byte stream
171+
val payloadManifestPath = Path.of(utgDir, "repoMapData", "payloadManifest.json")
172+
LOG.debug { "Adding payload manifest to ZIP: $payloadManifestPath" }
173+
// Create ZIP entry with the relative path
174+
val zipEntry = ZipEntry(payloadManifestPath.toString())
175+
it.putNextEntry(zipEntry)
176+
// Write the content once using byte array
177+
payloadMetadata.payloadManifest.toString().toByteArray().inputStream().use { inputStream ->
178+
inputStream.copyTo(it)
179+
}
180+
}
163181
}.toFile()
164182

165183
fun getProjectPayloadMetadata(): PayloadMetadata {
@@ -170,61 +188,91 @@ class CodeTestSessionConfig(
170188
var currentTotalLines = 0L
171189
val languageCounts = mutableMapOf<CodeWhispererProgrammingLanguage, Int>()
172190

191+
// Create a data structure to store file information
192+
val fileInfoList = mutableSetOf<Pair<String, Long>>()
193+
var exceededPayloadLimit = false
194+
173195
// Adding Target File to make sure target file doesn't get filtered out.
174196
selectedFile?.let { selected ->
175197
files.add(selected.path)
176198
currentTotalFileSize += selected.length
177199
currentTotalLines += countLinesInVirtualFile(selected)
200+
201+
// Add selected file info to the list
202+
fileInfoList.add(
203+
Pair(
204+
selected.path,
205+
selected.length
206+
)
207+
)
208+
178209
selected.programmingLanguage().let { language ->
179210
if (language !is CodeWhispererUnknownLanguage) {
180211
languageCounts[language] = (languageCounts[language] ?: 0) + 1
181212
}
182213
}
183214
}
184215

185-
moduleLoop@ for (module in project.modules) {
186-
val changeListManager = ChangeListManager.getInstance(module.project)
187-
module.guessModuleDir()?.let { moduleDir ->
188-
val gitIgnoreFilteringUtil = GitIgnoreFilteringUtil(moduleDir)
189-
stack.push(moduleDir)
190-
while (stack.isNotEmpty()) {
191-
val current = stack.pop()
192-
193-
if (!current.isDirectory) {
194-
if (current.isFile && current.path != selectedFile?.path &&
195-
!changeListManager.isIgnoredFile(current) &&
196-
runBlocking { !gitIgnoreFilteringUtil.ignoreFile(current) } &&
197-
runReadAction { !fileIndex.isInLibrarySource(current) }
198-
) {
199-
if (willExceedPayloadLimit(currentTotalFileSize, current.length)) {
200-
fileTooLarge()
201-
} else {
202-
try {
203-
val language = current.programmingLanguage()
204-
if (language !is CodeWhispererUnknownLanguage) {
205-
languageCounts[language] = (languageCounts[language] ?: 0) + 1
216+
run {
217+
for (module in project.modules) {
218+
val changeListManager = ChangeListManager.getInstance(module.project)
219+
module.guessModuleDir()?.let { moduleDir ->
220+
val gitIgnoreFilteringUtil = GitIgnoreFilteringUtil(moduleDir)
221+
stack.push(moduleDir)
222+
while (stack.isNotEmpty()) {
223+
val current = stack.pop()
224+
225+
if (!current.isDirectory) {
226+
if (current.isFile && current.path != selectedFile?.path &&
227+
!changeListManager.isIgnoredFile(current) &&
228+
runBlocking { !gitIgnoreFilteringUtil.ignoreFile(current) } &&
229+
runReadAction { !fileIndex.isInLibrarySource(current) }
230+
) {
231+
if (willExceedPayloadLimit(currentTotalFileSize, current.length)) {
232+
fileInfoList.add(
233+
Pair(
234+
current.path,
235+
current.length
236+
)
237+
)
238+
exceededPayloadLimit = true
239+
return@run
240+
} else {
241+
try {
242+
val language = current.programmingLanguage()
243+
if (language !is CodeWhispererUnknownLanguage) {
244+
languageCounts[language] = (languageCounts[language] ?: 0) + 1
245+
}
246+
files.add(current.path)
247+
currentTotalFileSize += current.length
248+
currentTotalLines += countLinesInVirtualFile(current)
249+
250+
// Add file info to the list
251+
fileInfoList.add(
252+
Pair(
253+
current.path,
254+
current.length
255+
)
256+
)
257+
} catch (e: Exception) {
258+
LOG.debug { "Error parsing the file: ${current.path} with error: ${e.message}" }
259+
continue
206260
}
207-
files.add(current.path)
208-
currentTotalFileSize += current.length
209-
currentTotalLines += countLinesInVirtualFile(current)
210-
} catch (e: Exception) {
211-
LOG.debug { "Error parsing the file: ${current.path} with error: ${e.message}" }
212-
continue
213261
}
214262
}
215-
}
216-
} else {
217-
// Directory case: only traverse if not ignored
218-
if (!changeListManager.isIgnoredFile(current) &&
219-
runBlocking { !gitIgnoreFilteringUtil.ignoreFile(current) } &&
220-
!traversedDirectories.contains(current) && current.isValid &&
221-
runReadAction { !fileIndex.isInLibrarySource(current) }
222-
) {
223-
for (child in current.children) {
224-
stack.push(child)
263+
} else {
264+
// Directory case: only traverse if not ignored
265+
if (!changeListManager.isIgnoredFile(current) &&
266+
runBlocking { !gitIgnoreFilteringUtil.ignoreFile(current) } &&
267+
!traversedDirectories.contains(current) && current.isValid &&
268+
runReadAction { !fileIndex.isInLibrarySource(current) }
269+
) {
270+
for (child in current.children) {
271+
stack.push(child)
272+
}
225273
}
274+
traversedDirectories.add(current)
226275
}
227-
traversedDirectories.add(current)
228276
}
229277
}
230278
}
@@ -238,7 +286,7 @@ class CodeTestSessionConfig(
238286
cannotFindValidFile("Amazon Q: doesn't contain valid files to generate tests")
239287
}
240288
programmingLanguage = maxCountLanguage
241-
return PayloadMetadata(files, currentTotalFileSize, currentTotalLines, maxCountLanguage.toTelemetryType())
289+
return PayloadMetadata(files, currentTotalFileSize, currentTotalLines, maxCountLanguage.toTelemetryType(), null, fileInfoList, exceededPayloadLimit)
242290
}
243291

244292
fun getRelativePath(): Path? = try {

‎plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/telemetry/CodeWhispererTelemetryService.kt

-2
Original file line numberDiff line numberDiff line change
@@ -315,7 +315,6 @@ class CodeWhispererTelemetryService {
315315
"Number of security scan issues with fixes: $issuesWithFixes, \n" +
316316
"Language: ${payloadContext.language}, \n" +
317317
"Uncompressed source payload size in bytes: ${payloadContext.srcPayloadSize}, \n" +
318-
"Uncompressed build payload size in bytes: ${payloadContext.buildPayloadSize}, \n" +
319318
"Compressed source zip file size in bytes: ${payloadContext.srcZipFileSize}, \n" +
320319
"Total project size in bytes: ${codeScanEvent.totalProjectSizeInBytes}, \n" +
321320
"Total duration of the security scan job in milliseconds: ${codeScanEvent.duration}, \n" +
@@ -334,7 +333,6 @@ class CodeWhispererTelemetryService {
334333
codewhispererCodeScanJobId = codeScanJobId,
335334
codewhispererCodeScanProjectBytes = codeScanEvent.totalProjectSizeInBytes,
336335
codewhispererCodeScanSrcPayloadBytes = payloadContext.srcPayloadSize,
337-
codewhispererCodeScanBuildPayloadBytes = payloadContext.buildPayloadSize,
338336
codewhispererCodeScanSrcZipFileBytes = payloadContext.srcZipFileSize,
339337
codewhispererCodeScanTotalIssues = totalIssues.toLong(),
340338
codewhispererCodeScanIssuesWithFixes = issuesWithFixes.toLong(),

‎plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/telemetry/CodeWhispererTelemetryServiceNew.kt

-2
Original file line numberDiff line numberDiff line change
@@ -299,7 +299,6 @@ class CodeWhispererTelemetryServiceNew {
299299
"Number of security scan issues with fixes: $issuesWithFixes, \n" +
300300
"Language: ${payloadContext.language}, \n" +
301301
"Uncompressed source payload size in bytes: ${payloadContext.srcPayloadSize}, \n" +
302-
"Uncompressed build payload size in bytes: ${payloadContext.buildPayloadSize}, \n" +
303302
"Compressed source zip file size in bytes: ${payloadContext.srcZipFileSize}, \n" +
304303
"Total project size in bytes: ${codeScanEvent.totalProjectSizeInBytes}, \n" +
305304
"Total duration of the security scan job in milliseconds: ${codeScanEvent.duration}, \n" +
@@ -317,7 +316,6 @@ class CodeWhispererTelemetryServiceNew {
317316
codewhispererCodeScanJobId = codeScanJobId,
318317
codewhispererCodeScanProjectBytes = codeScanEvent.totalProjectSizeInBytes,
319318
codewhispererCodeScanSrcPayloadBytes = payloadContext.srcPayloadSize,
320-
codewhispererCodeScanBuildPayloadBytes = payloadContext.buildPayloadSize,
321319
codewhispererCodeScanSrcZipFileBytes = payloadContext.srcZipFileSize,
322320
codewhispererCodeScanTotalIssues = totalIssues.toLong(),
323321
codewhispererCodeScanIssuesWithFixes = issuesWithFixes.toLong(),

‎plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/codetest/CodeTestSessionConfigTest.kt

-11
Original file line numberDiff line numberDiff line change
@@ -101,17 +101,6 @@ class CodeTestSessionConfigTest {
101101
}
102102
}
103103

104-
@Test
105-
fun `test createPayload should throw CodeWhispererCodeScanException if project size is more than Payload Limit`() {
106-
codeTestSessionConfig.stub {
107-
onGeneric { getPayloadLimitInBytes() }.thenReturn(1000)
108-
}
109-
110-
assertThrows<CodeTestException> {
111-
codeTestSessionConfig.createPayload()
112-
}
113-
}
114-
115104
private fun setupTestProject() {
116105
val testModule = projectRule.fixture.addModule("testModule")
117106
val testModule2 = projectRule.fixture.addModule("testModule2")

0 commit comments

Comments
 (0)