Skip to content

Commit

Permalink
fix(amazonq): Adding payload manifest for payload limit errors. (#5324)
Browse files Browse the repository at this point in the history 
* Adding payload manifest for payload limit errors for debugging.
  • Loading branch information
@ashishrp-aws
ashishrp-aws authored Feb 5, 2025
1 parent 970fee0 commit 17dcd44
Show file tree
Hide file tree
Showing 9 changed files with 110 additions and 71 deletions.
4 changes: 2 additions & 2 deletions .../software/aws/toolkits/jetbrains/services/amazonqCodeTest/CodeWhispererCodeTestSession.kt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -75,12 +75,12 @@ class CodeWhispererCodeTestSession(val sessionContext: CodeTestSessionContext) {

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

// 2 & 3. CreateUploadURL and upload the context.
Expand Down
4 changes: 4 additions & 0 deletions ...c/software/aws/toolkits/jetbrains/services/amazonqCodeTest/CodeWhispererUTGChatManager.kt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,7 @@ import software.aws.toolkits.jetbrains.services.amazonqCodeTest.session.Session
import software.aws.toolkits.jetbrains.services.amazonqCodeTest.utils.combineBuildAndExecuteLogFiles
import software.aws.toolkits.jetbrains.services.codemodernizer.utils.calculateTotalLatency
import software.aws.toolkits.jetbrains.services.codewhisperer.codetest.CodeTestException
import software.aws.toolkits.jetbrains.services.codewhisperer.codetest.fileTooLarge
import software.aws.toolkits.jetbrains.services.codewhisperer.codetest.sessionconfig.CodeTestSessionConfig
import software.aws.toolkits.jetbrains.services.codewhisperer.codetest.testGenStoppedError
import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererClientAdaptor
Expand Down Expand Up @@ -99,6 +100,9 @@ class CodeWhispererUTGChatManager(val project: Project, private val cs: Coroutin
session.srcZipFileSize = codeTestResponseContext.payloadContext.srcZipFileSize
session.artifactUploadDuration = codeTestResponseContext.serviceInvocationContext.artifactsUploadDuration
val path = codeTestResponseContext.currentFileRelativePath
if (codeTestResponseContext.payloadContext.payloadLimitCrossed == true) {
fileTooLarge()
}

val createUploadUrlResponse = codeTestResponseContext.createUploadUrlResponse ?: return
throwIfCancelled(session)
Expand Down
1 change: 0 additions & 1 deletion ...re/aws/toolkits/jetbrains/services/codewhisperer/codescan/CodeWhispererCodeScanSession.kt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -91,7 +91,6 @@ class CodeWhispererCodeScanSession(val sessionContext: CodeScanSessionContext) {
if (isProjectScope()) {
LOG.debug {
"Total size of source payload in KB: ${payloadContext.srcPayloadSize * 1.0 / TOTAL_BYTES_IN_KB} \n" +
"Total size of build payload in KB: ${(payloadContext.buildPayloadSize ?: 0) * 1.0 / TOTAL_BYTES_IN_KB} \n" +
"Total size of source zip file in KB: ${payloadContext.srcZipFileSize * 1.0 / TOTAL_BYTES_IN_KB} \n" +
"Total number of lines reviewed: ${payloadContext.totalLines} \n" +
"Total number of files included in payload: ${payloadContext.totalFiles} \n" +
Expand Down
5 changes: 4 additions & 1 deletion ...toolkits/jetbrains/services/codewhisperer/codescan/sessionconfig/CodeScanSessionConfig.kt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -334,7 +334,8 @@ data class PayloadContext(
val scannedFiles: List<VirtualFile>,
val srcPayloadSize: Long,
val srcZipFileSize: Long,
val buildPayloadSize: Long? = null,
val payloadManifest: Set<Pair<String, Long>>? = null,
val payloadLimitCrossed: Boolean? = false,
)

data class PayloadMetadata(
Expand All @@ -343,4 +344,6 @@ data class PayloadMetadata(
val linesScanned: Long,
val language: CodewhispererLanguage,
val codeDiff: String? = null,
val payloadManifest: Set<Pair<String, Long>>? = null,
val payloadLimitCrossed: Boolean? = false,
)
2 changes: 1 addition & 1 deletion .../src/software/aws/toolkits/jetbrains/services/codewhisperer/codetest/CodeTestException.kt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@ open class CodeTestException(
internal fun noFileOpenError(): Nothing =
throw CodeTestException(message("codewhisperer.codescan.no_file_open"), "ProjectZipError")

internal fun fileTooLarge(): Nothing =
fun fileTooLarge(): Nothing =
throw CodeTestException(message("codewhisperer.codescan.file_too_large_telemetry"), "ProjectZipError")

internal fun cannotFindFile(errorMessage: String, filepath: String): Nothing =
Expand Down
150 changes: 99 additions & 51 deletions ...toolkits/jetbrains/services/codewhisperer/codetest/sessionconfig/CodeTestSessionConfig.kt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -101,15 +101,17 @@ class CodeTestSessionConfig(
}

// Copy all the included source files to the source zip
val srcZip = zipFiles(payloadMetadata.sourceFiles.map { Path.of(it) })
val srcZip = zipFiles(payloadMetadata)
val payloadContext = PayloadContext(
payloadMetadata.language,
payloadMetadata.linesScanned,
payloadMetadata.sourceFiles.size,
Instant.now().toEpochMilli() - start,
payloadMetadata.sourceFiles.mapNotNull { Path.of(it).toFile().toVirtualFile() },
payloadMetadata.payloadSize,
srcZip.length()
srcZip.length(),
payloadMetadata.payloadManifest,
payloadMetadata.payloadLimitCrossed,
)

return Payload(payloadContext, srcZip)
Expand All @@ -129,16 +131,20 @@ class CodeTestSessionConfig(
}
}

private fun zipFiles(files: List<Path>): File = createTemporaryZipFile {
files.forEach { file ->
try {
val relativePath = file.relativeTo(projectRoot.toNioPath())
val projectBaseName = projectRoot.name
val zipEntryPath = "$projectBaseName/${relativePath.toString().replace("\\", "/")}"
LOG.debug { "Adding file to ZIP: $zipEntryPath" }
it.putNextEntry(zipEntryPath, file)
} catch (e: Exception) {
cannotFindFile("Zipping error: ${e.message}", file.toString())
private fun zipFiles(payloadMetadata: PayloadMetadata): File = createTemporaryZipFile {
if (payloadMetadata.payloadLimitCrossed == false) {
payloadMetadata.sourceFiles.forEach { file ->
Path.of(file).let { path ->
try {
val relativePath = path.relativeTo(projectRoot.toNioPath())
val projectBaseName = projectRoot.name
val zipEntryPath = "$projectBaseName/${relativePath.toString().replace("\\", "/")}"
LOG.debug { "Adding file to ZIP: $zipEntryPath" }
it.putNextEntry(zipEntryPath, path)
} catch (e: Exception) {
cannotFindFile("Zipping error: ${e.message}", path.toString())
}
}
}
}

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

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

// Create a data structure to store file information
val fileInfoList = mutableSetOf<Pair<String, Long>>()
var exceededPayloadLimit = false

// Adding Target File to make sure target file doesn't get filtered out.
selectedFile?.let { selected ->
files.add(selected.path)
currentTotalFileSize += selected.length
currentTotalLines += countLinesInVirtualFile(selected)

// Add selected file info to the list
fileInfoList.add(
Pair(
selected.path,
selected.length
)
)

selected.programmingLanguage().let { language ->
if (language !is CodeWhispererUnknownLanguage) {
languageCounts[language] = (languageCounts[language] ?: 0) + 1
}
}
}

moduleLoop@ for (module in project.modules) {
val changeListManager = ChangeListManager.getInstance(module.project)
module.guessModuleDir()?.let { moduleDir ->
val gitIgnoreFilteringUtil = GitIgnoreFilteringUtil(moduleDir)
stack.push(moduleDir)
while (stack.isNotEmpty()) {
val current = stack.pop()

if (!current.isDirectory) {
if (current.isFile && current.path != selectedFile?.path &&
!changeListManager.isIgnoredFile(current) &&
runBlocking { !gitIgnoreFilteringUtil.ignoreFile(current) } &&
runReadAction { !fileIndex.isInLibrarySource(current) }
) {
if (willExceedPayloadLimit(currentTotalFileSize, current.length)) {
fileTooLarge()
} else {
try {
val language = current.programmingLanguage()
if (language !is CodeWhispererUnknownLanguage) {
languageCounts[language] = (languageCounts[language] ?: 0) + 1
run {
for (module in project.modules) {
val changeListManager = ChangeListManager.getInstance(module.project)
module.guessModuleDir()?.let { moduleDir ->
val gitIgnoreFilteringUtil = GitIgnoreFilteringUtil(moduleDir)
stack.push(moduleDir)
while (stack.isNotEmpty()) {
val current = stack.pop()

if (!current.isDirectory) {
if (current.isFile && current.path != selectedFile?.path &&
!changeListManager.isIgnoredFile(current) &&
runBlocking { !gitIgnoreFilteringUtil.ignoreFile(current) } &&
runReadAction { !fileIndex.isInLibrarySource(current) }
) {
if (willExceedPayloadLimit(currentTotalFileSize, current.length)) {
fileInfoList.add(
Pair(
current.path,
current.length
)
)
exceededPayloadLimit = true
return@run
} else {
try {
val language = current.programmingLanguage()
if (language !is CodeWhispererUnknownLanguage) {
languageCounts[language] = (languageCounts[language] ?: 0) + 1
}
files.add(current.path)
currentTotalFileSize += current.length
currentTotalLines += countLinesInVirtualFile(current)

// Add file info to the list
fileInfoList.add(
Pair(
current.path,
current.length
)
)
} catch (e: Exception) {
LOG.debug { "Error parsing the file: ${current.path} with error: ${e.message}" }
continue
}
files.add(current.path)
currentTotalFileSize += current.length
currentTotalLines += countLinesInVirtualFile(current)
} catch (e: Exception) {
LOG.debug { "Error parsing the file: ${current.path} with error: ${e.message}" }
continue
}
}
}
} else {
// Directory case: only traverse if not ignored
if (!changeListManager.isIgnoredFile(current) &&
runBlocking { !gitIgnoreFilteringUtil.ignoreFile(current) } &&
!traversedDirectories.contains(current) && current.isValid &&
runReadAction { !fileIndex.isInLibrarySource(current) }
) {
for (child in current.children) {
stack.push(child)
} else {
// Directory case: only traverse if not ignored
if (!changeListManager.isIgnoredFile(current) &&
runBlocking { !gitIgnoreFilteringUtil.ignoreFile(current) } &&
!traversedDirectories.contains(current) && current.isValid &&
runReadAction { !fileIndex.isInLibrarySource(current) }
) {
for (child in current.children) {
stack.push(child)
}
}
traversedDirectories.add(current)
}
traversedDirectories.add(current)
}
}
}
Expand All @@ -238,7 +286,7 @@ class CodeTestSessionConfig(
cannotFindValidFile("Amazon Q: doesn't contain valid files to generate tests")
}
programmingLanguage = maxCountLanguage
return PayloadMetadata(files, currentTotalFileSize, currentTotalLines, maxCountLanguage.toTelemetryType())
return PayloadMetadata(files, currentTotalFileSize, currentTotalLines, maxCountLanguage.toTelemetryType(), null, fileInfoList, exceededPayloadLimit)
}

fun getRelativePath(): Path? = try {
Expand Down
2 changes: 0 additions & 2 deletions .../aws/toolkits/jetbrains/services/codewhisperer/telemetry/CodeWhispererTelemetryService.kt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -315,7 +315,6 @@ class CodeWhispererTelemetryService {
"Number of security scan issues with fixes: $issuesWithFixes, \n" +
"Language: ${payloadContext.language}, \n" +
"Uncompressed source payload size in bytes: ${payloadContext.srcPayloadSize}, \n" +
"Uncompressed build payload size in bytes: ${payloadContext.buildPayloadSize}, \n" +
"Compressed source zip file size in bytes: ${payloadContext.srcZipFileSize}, \n" +
"Total project size in bytes: ${codeScanEvent.totalProjectSizeInBytes}, \n" +
"Total duration of the security scan job in milliseconds: ${codeScanEvent.duration}, \n" +
Expand All @@ -334,7 +333,6 @@ class CodeWhispererTelemetryService {
codewhispererCodeScanJobId = codeScanJobId,
codewhispererCodeScanProjectBytes = codeScanEvent.totalProjectSizeInBytes,
codewhispererCodeScanSrcPayloadBytes = payloadContext.srcPayloadSize,
codewhispererCodeScanBuildPayloadBytes = payloadContext.buildPayloadSize,
codewhispererCodeScanSrcZipFileBytes = payloadContext.srcZipFileSize,
codewhispererCodeScanTotalIssues = totalIssues.toLong(),
codewhispererCodeScanIssuesWithFixes = issuesWithFixes.toLong(),
Expand Down
2 changes: 0 additions & 2 deletions ...s/toolkits/jetbrains/services/codewhisperer/telemetry/CodeWhispererTelemetryServiceNew.kt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -299,7 +299,6 @@ class CodeWhispererTelemetryServiceNew {
"Number of security scan issues with fixes: $issuesWithFixes, \n" +
"Language: ${payloadContext.language}, \n" +
"Uncompressed source payload size in bytes: ${payloadContext.srcPayloadSize}, \n" +
"Uncompressed build payload size in bytes: ${payloadContext.buildPayloadSize}, \n" +
"Compressed source zip file size in bytes: ${payloadContext.srcZipFileSize}, \n" +
"Total project size in bytes: ${codeScanEvent.totalProjectSizeInBytes}, \n" +
"Total duration of the security scan job in milliseconds: ${codeScanEvent.duration}, \n" +
Expand All @@ -317,7 +316,6 @@ class CodeWhispererTelemetryServiceNew {
codewhispererCodeScanJobId = codeScanJobId,
codewhispererCodeScanProjectBytes = codeScanEvent.totalProjectSizeInBytes,
codewhispererCodeScanSrcPayloadBytes = payloadContext.srcPayloadSize,
codewhispererCodeScanBuildPayloadBytes = payloadContext.buildPayloadSize,
codewhispererCodeScanSrcZipFileBytes = payloadContext.srcZipFileSize,
codewhispererCodeScanTotalIssues = totalIssues.toLong(),
codewhispererCodeScanIssuesWithFixes = issuesWithFixes.toLong(),
Expand Down
11 changes: 0 additions & 11 deletions ...tware/aws/toolkits/jetbrains/services/codewhisperer/codetest/CodeTestSessionConfigTest.kt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -101,17 +101,6 @@ class CodeTestSessionConfigTest {
}
}

@Test
fun `test createPayload should throw CodeWhispererCodeScanException if project size is more than Payload Limit`() {
codeTestSessionConfig.stub {
onGeneric { getPayloadLimitInBytes() }.thenReturn(1000)
}

assertThrows<CodeTestException> {
codeTestSessionConfig.createPayload()
}
}

private fun setupTestProject() {
val testModule = projectRule.fixture.addModule("testModule")
val testModule2 = projectRule.fixture.addModule("testModule2")
Expand Down

0 comments on commit 17dcd44

Please sign in to comment.