p

Author: Will-ShaoHua

plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererConstants.kt

@@ -189,6 +189,7 @@ object CodeWhispererConstants {
         const val NUMBER_OF_CHUNK_TO_FETCH = 3
         const val MAX_TOTAL_LENGTH = 20480
         const val MAX_LENGTH_PER_CHUNK = 10240
+        const val MAX_CONTEXT_COUNT = 5
     }
 
     object Utg {

plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererFileContextProvider.kt

@@ -342,6 +342,10 @@ class DefaultCodeWhispererFileContextProvider(private val project: Project) : Fi
             }
         }
 
+        if (c.size > CodeWhispererConstants.CrossFile.MAX_CONTEXT_COUNT) {
+            c = c.subList(0, CodeWhispererConstants.CrossFile.MAX_CONTEXT_COUNT)
+        }
+
         var curTotalLength = c.sumOf { it.content.length }
         while (curTotalLength >= CodeWhispererConstants.CrossFile.MAX_TOTAL_LENGTH) {
             val last = c.last()

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

@@ -529,6 +529,29 @@ class CodeWhispererFileContextProviderTest {
         assertThat(r.targetFileName).isEqualTo("foo")
     }
 
+    @Test
+    fun `truncate context should make context item lte 5`() {
+        val supplementalContext = SupplementalContextInfo(
+            isUtg = false,
+            contents = listOf(
+                Chunk(content = "a", path = "a.java"),
+                Chunk(content = "b", path = "b.java"),
+                Chunk(content = "c", path = "c.java"),
+                Chunk(content = "d", path = "d.java"),
+                Chunk(content = "e", path = "e.java"),
+                Chunk(content = "f", path = "e.java"),
+                Chunk(content = "g", path = "e.java"),
+            ),
+            targetFileName = "foo",
+            strategy = CrossFileStrategy.Codemap
+        )
+
+        val r = sut.truncateContext(supplementalContext)
+        assertThat(r.contents).hasSize(5)
+        assertThat(r.strategy).isEqualTo(CrossFileStrategy.Codemap)
+        assertThat(r.targetFileName).isEqualTo("foo")
+    }
+
     @Test
     fun `truncate context should make context length per item fit in 10240 cap`() {
         val chunkA = Chunk(content = "a\n".repeat(4000), path = "a.java")