diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/context/additionalContextProvider.test.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/context/additionalContextProvider.test.ts index 3be50a289e..4e124d2fa5 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/context/additionalContextProvider.test.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/context/additionalContextProvider.test.ts @@ -9,7 +9,7 @@ import { getInitialContextInfo, getUserPromptsDirectory } from './contextUtils' import { LocalProjectContextController } from '../../../shared/localProjectContextController' import { workspaceUtils } from '@aws/lsp-core' import { ChatDatabase } from '../tools/chatDb/chatDb' -import { TriggerContext } from './agenticChatTriggerContext' +import { TriggerContext, workspaceChunkMaxSize } from './agenticChatTriggerContext' import { expect } from 'chai' describe('AdditionalContextProvider', () => { @@ -897,6 +897,53 @@ describe('AdditionalContextProvider', () => { assert.strictEqual(result[0].innerContext, 'Content from indexing library') }) + it('should cap oversized context content and mark it as truncated', async () => { + const mockWorkspaceFolder = { + uri: URI.file('/workspace').toString(), + name: 'test', + } + sinon.stub(workspaceUtils, 'getWorkspaceFolderPaths').returns(['/workspace']) + const triggerContext: TriggerContext = { + workspaceFolder: mockWorkspaceFolder, + } + + fsExistsStub.callsFake((pathStr: string) => { + if (pathStr.includes(path.join('.amazonq', 'rules'))) { + return Promise.resolve(true) + } + return Promise.resolve(false) + }) + fsReadDirStub.resolves([{ name: 'rule1.md', isFile: () => true, isDirectory: () => false }]) + + const largeContent = 'line of text\n'.repeat(10_000) + assert.ok(largeContent.length > workspaceChunkMaxSize) + + getContextCommandPromptStub + .onFirstCall() + .resolves([]) + .onSecondCall() + .resolves([ + { + name: 'Large Rule', + description: 'Test Description', + content: largeContent, + filePath: '/workspace/.amazonq/rules/rule1.md', + relativePath: '.amazonq/rules/rule1.md', + startLine: 1, + endLine: 10, + }, + ]) + + const result = await provider.getAdditionalContext(triggerContext, '') + + assert.strictEqual(result.length, 1) + const innerContext = result[0].innerContext ?? '' + assert.strictEqual(innerContext.length, workspaceChunkMaxSize) + assert.ok(innerContext.startsWith('line of text\n')) + assert.ok(innerContext.includes('[Content truncated:')) + assert.ok(innerContext.includes(`${largeContent.length} characters`)) + }) + it('should handle filesystem read errors gracefully in fallback', async () => { const mockWorkspaceFolder = { uri: URI.file('/workspace').toString(), diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/context/additionalContextProvider.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/context/additionalContextProvider.ts index 60a880a956..fa3f38d39e 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/context/additionalContextProvider.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/context/additionalContextProvider.ts @@ -26,6 +26,7 @@ import { getInitialContextInfo, promptFileExtension, getCodeSymbolDescription, + truncateContextContent, } from './contextUtils' import { LocalProjectContextController } from '../../../shared/localProjectContextController' import { Features } from '../../types' @@ -567,7 +568,7 @@ export class AdditionalContextProvider { const entry = { name: prompt.name.substring(0, additionalContentNameLimit), description: '', - innerContext: prompt.content.substring(0, workspaceChunkMaxSize), + innerContext: truncateContextContent(prompt.content, workspaceChunkMaxSize), type: contextType, path: prompt.filePath, relativePath: relativePath, diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/context/contextUtils.test.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/context/contextUtils.test.ts index b3c0a339e3..aa0d575dd1 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/context/contextUtils.test.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/context/contextUtils.test.ts @@ -9,6 +9,7 @@ import { mergeRelevantTextDocuments, mergeFileLists, getCodeSymbolDescription, + truncateContextContent, } from './contextUtils' import * as pathUtils from '@aws/lsp-core/out/util/path' import { sanitizeFilename } from '@aws/lsp-core/out/util/text' @@ -428,4 +429,38 @@ describe('contextUtils', () => { expect(result).to.equal(`Interface, ${path.join('workspace', 'src', 'models.ts')}`) }) }) + + describe('truncateContextContent', () => { + it('returns content unchanged when it fits within the limit', () => { + const content = 'short content' + expect(truncateContextContent(content, 100)).to.equal(content) + }) + + it('returns content unchanged when it is exactly the limit', () => { + const content = 'a'.repeat(50) + expect(truncateContextContent(content, 50)).to.equal(content) + }) + + it('appends a truncation marker and never exceeds the limit', () => { + const lines = Array.from({ length: 1000 }, (_, i) => `This is line ${i + 1} of a large file.`) + const content = lines.join('\n') + const maxLength = 4096 + + const result = truncateContextContent(content, maxLength) + + expect(result.length).to.equal(maxLength) + expect(result.startsWith('This is line 1 of a large file.')).to.equal(true) + expect(result).to.include('[Content truncated:') + expect(result).to.include(`${content.length} characters`) + // The marker must be the tail of the result so the model sees it after the excerpt + expect(result.endsWith('Read the file directly to see the rest.]')).to.equal(true) + }) + + it('falls back to a plain cut when the limit is too small for the marker', () => { + const content = 'x'.repeat(500) + const result = truncateContextContent(content, 10) + + expect(result).to.equal('x'.repeat(10)) + }) + }) }) diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/context/contextUtils.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/context/contextUtils.ts index 6fa6cee098..c4d40a2835 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/context/contextUtils.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/context/contextUtils.ts @@ -65,6 +65,31 @@ export const getUserPromptsDirectory = (): string => { return path.join(getUserHomeDir(), '.aws', 'amazonq', 'prompts') } +/** + * Truncates context content (e.g. an `@file` or pinned file) to `maxLength` characters. + * + * When the content is cut, an explicit marker is appended so the model knows it is only + * seeing the beginning of the file and should read the rest via tools instead of assuming + * the file ends where the excerpt ends. The returned string never exceeds `maxLength`. + * + * @param content - Full content of the context item + * @param maxLength - Maximum number of characters allowed for the returned string + * @returns The original content if it fits, otherwise a truncated prefix followed by a marker + */ +export function truncateContextContent(content: string, maxLength: number): string { + if (content.length <= maxLength) { + return content + } + + const marker = `\n\n[Content truncated: this file has ${content.length} characters, only the beginning is included above. Read the file directly to see the rest.]` + + if (marker.length >= maxLength) { + return content.substring(0, maxLength) + } + + return content.substring(0, maxLength - marker.length) + marker +} + /** * Creates a secure file path for a new prompt file. *