fix(amazonq): Adding payload manifest for payload limit errors.

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

@@ -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.

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

@@ -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
@@ -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)

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

@@ -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" +

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

@@ -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(
@@ -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,
 )

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

@@ -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 =

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

@@ -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)
@@ -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())
+                    }
+                }
             }
         }
 
@@ -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 {
@@ -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)
                     }
                 }
             }
@@ -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 {

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

@@ -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" +
@@ -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(),

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

@@ -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" +
@@ -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(),

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

@@ -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")