fix(amazonq): max 20k char for inline supplementalContext

packages/amazonq/.changes/next-release/Bug Fix-72b26a2d-3647-4d07-b0ef-92238e3f0050.json

@@ -0,0 +1,4 @@
+{
+	"type": "Bug Fix",
+	"description": "Fix inline completion supplementalContext length exceeding maximum in certain cases"
+}

packages/amazonq/test/unit/codewhisperer/util/crossFileContextUtil.test.ts

@@ -8,7 +8,12 @@ import * as FakeTimers from '@sinonjs/fake-timers'
 import * as vscode from 'vscode'
 import * as sinon from 'sinon'
 import * as crossFile from 'aws-core-vscode/codewhisperer'
-import { aStringWithLineCount, createMockTextEditor, installFakeClock } from 'aws-core-vscode/test'
+import {
+    aLongStringWithLineCount,
+    aStringWithLineCount,
+    createMockTextEditor,
+    installFakeClock,
+} from 'aws-core-vscode/test'
 import { FeatureConfigProvider, crossFileContextConfig } from 'aws-core-vscode/codewhisperer'
 import {
     assertTabCount,
@@ -71,8 +76,8 @@ describe('crossFileContextUtil', function () {
             assert.strictEqual(actual.supplementalContextItems[2].content.split('\n').length, 50)
         })
 
-        it('for t1 group, should return repomap + opentabs context', async function () {
-            await toTextEditor(aStringWithLineCount(200), 'CrossFile.java', tempFolder, { preview: false })
+        it('for t1 group, should return repomap + opentabs context, should not exceed 20k total length', async function () {
+            await toTextEditor(aLongStringWithLineCount(200), 'CrossFile.java', tempFolder, { preview: false })
             const myCurrentEditor = await toTextEditor('', 'TargetFile.java', tempFolder, {
                 preview: false,
             })
@@ -85,25 +90,23 @@ describe('crossFileContextUtil', function () {
                 .withArgs(sinon.match.any, sinon.match.any, 'codemap')
                 .resolves([
                     {
-                        content: 'foo',
+                        content: 'foo'.repeat(3000),
                         score: 0,
                         filePath: 'q-inline',
                     },
                 ])
 
             const actual = await crossFile.fetchSupplementalContextForSrc(myCurrentEditor, fakeCancellationToken)
             assert.ok(actual)
-            assert.strictEqual(actual.supplementalContextItems.length, 4)
+            assert.strictEqual(actual.supplementalContextItems.length, 3)
             assert.strictEqual(actual?.strategy, 'codemap')
             assert.deepEqual(actual?.supplementalContextItems[0], {
-                content: 'foo',
+                content: 'foo'.repeat(3000),
                 score: 0,
                 filePath: 'q-inline',
             })
-
             assert.strictEqual(actual.supplementalContextItems[1].content.split('\n').length, 50)
             assert.strictEqual(actual.supplementalContextItems[2].content.split('\n').length, 50)
-            assert.strictEqual(actual.supplementalContextItems[3].content.split('\n').length, 50)
         })
 
         it.skip('for t2 group, should return global bm25 context and no repomap', async function () {

packages/core/src/codewhisperer/models/constants.ts

@@ -84,6 +84,7 @@ export const lineBreakWin = '\r\n'
 
 export const supplementalContextTimeoutInMs = 100
 
+export const supplementalContextMaxTotalLength = 20480
 /**
  * Ux of recommendations
  */

packages/core/src/codewhisperer/util/supplementalContext/crossFileContextUtil.ts

@@ -7,7 +7,11 @@ import * as vscode from 'vscode'
 import { FeatureConfigProvider, fs } from '../../../shared'
 import path = require('path')
 import { BM25Document, BM25Okapi } from './rankBm25'
-import { crossFileContextConfig, supplementalContextTimeoutInMs } from '../../models/constants'
+import {
+    crossFileContextConfig,
+    supplementalContextTimeoutInMs,
+    supplementalContextMaxTotalLength,
+} from '../../models/constants'
 import { isTestFile } from './codeParsingUtil'
 import { getFileDistance } from '../../../shared/filesystemUtilities'
 import { getOpenFilesInWindow } from '../../../shared/utilities/editorUtilities'
@@ -99,13 +103,22 @@ export async function fetchSupplementalContextForSrc(
                 const opentabsContext = await fetchOpentabsContext(editor, cancellationToken)
                 const codemap = await fetchProjectContext(editor, 'codemap')
 
+                function addToResult(items: CodeWhispererSupplementalContextItem[]) {
+                    for (const item of items) {
+                        const curLen = result.reduce((acc, i) => acc + i.content.length, 0)
+                        if (curLen + item.content.length < supplementalContextMaxTotalLength) {
+                            result.push(item)
+                        }
+                    }
+                }
+
                 if (codemap && codemap.length > 0) {
-                    result.push(...codemap)
+                    addToResult(codemap)
                     hasCodemap = true
                 }
 
                 if (opentabsContext && opentabsContext.length > 0) {
-                    result.push(...opentabsContext)
+                    addToResult(opentabsContext)
                     hasOpentabs = true
                 }
 

packages/core/src/test/codewhisperer/testUtil.ts

@@ -329,3 +329,12 @@ export function aStringWithLineCount(lineCount: number, start: number = 0): stri
 
     return s.trimEnd()
 }
+
+export function aLongStringWithLineCount(lineCount: number, start: number = 0): string {
+    let s = ''
+    for (let i = start; i < start + lineCount; i++) {
+        s += `a`.repeat(100) + `line${i}\n`
+    }
+
+    return s.trimEnd()
+}