fix(inline-completion): potential inline completion failure due to input validation exception of supplemental context

packages/amazonq/.changes/next-release/Bug Fix-bde47093-ee34-4802-a556-d1e5dc526f4d.json

@@ -0,0 +1,4 @@
+{
+	"type": "Bug Fix",
+	"description": "Fix inline completion failure due to context length exceeding the threshold"
+}

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

@@ -7,12 +7,15 @@ import assert from 'assert'
 import * as FakeTimers from '@sinonjs/fake-timers'
 import * as vscode from 'vscode'
 import * as sinon from 'sinon'
+import * as os from 'os'
 import * as crossFile from 'aws-core-vscode/codewhisperer'
 import { TestFolder, assertTabCount, installFakeClock } from 'aws-core-vscode/test'
-import { FeatureConfigProvider } from 'aws-core-vscode/codewhisperer'
+import { CodeWhispererSupplementalContext, FeatureConfigProvider } from 'aws-core-vscode/codewhisperer'
 import { toTextEditor } from 'aws-core-vscode/test'
 import { LspController } from 'aws-core-vscode/amazonq'
 
+const newLine = os.EOL
+
 describe('supplementalContextUtil', function () {
     let testFolder: TestFolder
     let clock: FakeTimers.InstalledClock
@@ -83,4 +86,180 @@ describe('supplementalContextUtil', function () {
             })
         })
     })
+
+    describe('truncation', function () {
+        it('truncate context should do nothing if everything fits in constraint', function () {
+            const chunkA: crossFile.CodeWhispererSupplementalContextItem = {
+                content: 'a',
+                filePath: 'a.java',
+                score: 0,
+            }
+            const chunkB: crossFile.CodeWhispererSupplementalContextItem = {
+                content: 'b',
+                filePath: 'b.java',
+                score: 1,
+            }
+            const chunks = [chunkA, chunkB]
+
+            const supplementalContext: CodeWhispererSupplementalContext = {
+                isUtg: false,
+                isProcessTimeout: false,
+                supplementalContextItems: chunks,
+                contentsLength: 25000,
+                latency: 0,
+                strategy: 'codemap',
+            }
+
+            const actual = crossFile.truncateSuppelementalContext(supplementalContext)
+            assert.strictEqual(actual.supplementalContextItems.length, 2)
+            assert.strictEqual(actual.supplementalContextItems[0].content, 'a')
+            assert.strictEqual(actual.supplementalContextItems[1].content, 'b')
+        })
+
+        it('truncateLineByLine should drop the last line if max length is greater than threshold', function () {
+            const input =
+                repeatString('a', 11) +
+                newLine +
+                repeatString('b', 11) +
+                newLine +
+                repeatString('c', 11) +
+                newLine +
+                repeatString('d', 11) +
+                newLine +
+                repeatString('e', 11)
+
+            assert.ok(input.length > 50)
+            const actual = crossFile.truncateLineByLine(input, 50)
+            assert.ok(actual.length <= 50)
+
+            const input2 = repeatString(`b${newLine}`, 10)
+            const actual2 = crossFile.truncateLineByLine(input2, 8)
+            assert.ok(actual2.length <= 8)
+        })
+
+        it('truncation context should make context length per item lte 10240 cap', function () {
+            const chunkA: crossFile.CodeWhispererSupplementalContextItem = {
+                content: repeatString(`a${newLine}`, 4000),
+                filePath: 'a.java',
+                score: 0,
+            }
+            const chunkB: crossFile.CodeWhispererSupplementalContextItem = {
+                content: repeatString(`b${newLine}`, 6000),
+                filePath: 'b.java',
+                score: 1,
+            }
+            const chunkC: crossFile.CodeWhispererSupplementalContextItem = {
+                content: repeatString(`c${newLine}`, 1000),
+                filePath: 'c.java',
+                score: 2,
+            }
+            const chunkD: crossFile.CodeWhispererSupplementalContextItem = {
+                content: repeatString(`d${newLine}`, 1500),
+                filePath: 'd.java',
+                score: 3,
+            }
+
+            assert.ok(
+                chunkA.content.length + chunkB.content.length + chunkC.content.length + chunkD.content.length > 20480
+            )
+
+            const supplementalContext: CodeWhispererSupplementalContext = {
+                isUtg: false,
+                isProcessTimeout: false,
+                supplementalContextItems: [chunkA, chunkB, chunkC, chunkD],
+                contentsLength: 25000,
+                latency: 0,
+                strategy: 'codemap',
+            }
+
+            const actual = crossFile.truncateSuppelementalContext(supplementalContext)
+            assert.strictEqual(actual.supplementalContextItems.length, 3)
+            assert.ok(actual.contentsLength <= 20480)
+            assert.strictEqual(actual.strategy, 'codemap')
+        })
+
+        it('truncate context should make context items lte 5', function () {
+            const chunkA: crossFile.CodeWhispererSupplementalContextItem = {
+                content: 'a',
+                filePath: 'a.java',
+                score: 0,
+            }
+            const chunkB: crossFile.CodeWhispererSupplementalContextItem = {
+                content: 'b',
+                filePath: 'b.java',
+                score: 1,
+            }
+            const chunkC: crossFile.CodeWhispererSupplementalContextItem = {
+                content: 'c',
+                filePath: 'c.java',
+                score: 2,
+            }
+            const chunkD: crossFile.CodeWhispererSupplementalContextItem = {
+                content: 'd',
+                filePath: 'd.java',
+                score: 3,
+            }
+            const chunkE: crossFile.CodeWhispererSupplementalContextItem = {
+                content: 'e',
+                filePath: 'e.java',
+                score: 4,
+            }
+            const chunkF: crossFile.CodeWhispererSupplementalContextItem = {
+                content: 'f',
+                filePath: 'f.java',
+                score: 5,
+            }
+            const chunkG: crossFile.CodeWhispererSupplementalContextItem = {
+                content: 'g',
+                filePath: 'g.java',
+                score: 6,
+            }
+            const chunks = [chunkA, chunkB, chunkC, chunkD, chunkE, chunkF, chunkG]
+
+            assert.strictEqual(chunks.length, 7)
+
+            const supplementalContext: CodeWhispererSupplementalContext = {
+                isUtg: false,
+                isProcessTimeout: false,
+                supplementalContextItems: chunks,
+                contentsLength: 25000,
+                latency: 0,
+                strategy: 'codemap',
+            }
+
+            const actual = crossFile.truncateSuppelementalContext(supplementalContext)
+            assert.strictEqual(actual.supplementalContextItems.length, 5)
+        })
+
+        describe('truncate line by line', function () {
+            it('should return empty if empty string is provided', function () {
+                const input = ''
+                const actual = crossFile.truncateLineByLine(input, 50)
+                assert.strictEqual(actual, '')
+            })
+
+            it('should return empty if 0 max length is provided', function () {
+                const input = 'aaaaa'
+                const actual = crossFile.truncateLineByLine(input, 0)
+                assert.strictEqual(actual, '')
+            })
+
+            it('should flip the value if negative max length is provided', function () {
+                const input = `aaaaa${newLine}bbbbb`
+                const actual = crossFile.truncateLineByLine(input, -6)
+                const expected = crossFile.truncateLineByLine(input, 6)
+                assert.strictEqual(actual, expected)
+                assert.strictEqual(actual, 'aaaaa')
+            })
+        })
+    })
 })
+
+function repeatString(s: string, n: number): string {
+    let output = ''
+    for (let i = 0; i < n; i++) {
+        output += s
+    }
+
+    return output
+}

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

@@ -856,6 +856,8 @@ export const crossFileContextConfig = {
     topK: 3,
     numberOfLinesEachChunk: 50,
     maximumTotalLength: 20480,
+    maxLengthEachChunk: 10240,
+    maxContextCount: 5,
 }
 
 export const utgConfig = {

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

@@ -11,6 +11,8 @@ import { CancellationError } from '../../../shared/utilities/timeoutUtils'
 import { ToolkitError } from '../../../shared/errors'
 import { getLogger } from '../../../shared/logger/logger'
 import { CodeWhispererSupplementalContext } from '../../models/model'
+import * as os from 'os'
+import { crossFileContextConfig } from '../../models/constants'
 
 export async function fetchSupplementalContext(
     editor: vscode.TextEditor,
@@ -36,7 +38,7 @@ export async function fetchSupplementalContext(
     return supplementalContextPromise
         .then((value) => {
             if (value) {
-                return {
+                const resBeforeTruncation = {
                     isUtg: isUtg,
                     isProcessTimeout: false,
                     supplementalContextItems: value.supplementalContextItems.filter(
@@ -46,6 +48,8 @@ export async function fetchSupplementalContext(
                     latency: performance.now() - timesBeforeFetching,
                     strategy: value.strategy,
                 }
+
+                return truncateSuppelementalContext(resBeforeTruncation)
             } else {
                 return undefined
             }
@@ -68,3 +72,66 @@ export async function fetchSupplementalContext(
             }
         })
 }
+
+/**
+ * Requirement
+ * - Maximum 5 supplemental context.
+ * - Each chunk can't exceed 10240 characters
+ * - Sum of all chunks can't exceed 20480 characters
+ */
+export function truncateSuppelementalContext(
+    context: CodeWhispererSupplementalContext
+): CodeWhispererSupplementalContext {
+    let c = context.supplementalContextItems.map((item) => {
+        if (item.content.length > crossFileContextConfig.maxLengthEachChunk) {
+            return {
+                ...item,
+                content: truncateLineByLine(item.content, crossFileContextConfig.maxLengthEachChunk),
+            }
+        } else {
+            return item
+        }
+    })
+
+    if (c.length > crossFileContextConfig.maxContextCount) {
+        c = c.slice(0, crossFileContextConfig.maxContextCount)
+    }
+
+    let curTotalLength = c.reduce((acc, cur) => {
+        return acc + cur.content.length
+    }, 0)
+    while (curTotalLength >= 20480 && c.length - 1 >= 0) {
+        const last = c[c.length - 1]
+        c = c.slice(0, -1)
+        curTotalLength -= last.content.length
+    }
+
+    return {
+        ...context,
+        supplementalContextItems: c,
+        contentsLength: curTotalLength,
+    }
+}
+
+export function truncateLineByLine(input: string, l: number): string {
+    const maxLength = l > 0 ? l : -1 * l
+    if (input.length === 0) {
+        return ''
+    }
+
+    const shouldAddNewLineBack = input.endsWith(os.EOL)
+    let lines = input.trim().split(os.EOL)
+    let curLen = input.length
+    while (curLen > maxLength && lines.length - 1 >= 0) {
+        const last = lines[lines.length - 1]
+        lines = lines.slice(0, -1)
+        curLen -= last.length + 1
+    }
+
+    const r = lines.join(os.EOL)
+    if (shouldAddNewLineBack) {
+        return r + os.EOL
+    } else {
+        return r
+    }
+}