diff --git a/packages/agent-runtime/src/system-prompt/__tests__/truncate-file-tree.test.ts b/packages/agent-runtime/src/system-prompt/__tests__/truncate-file-tree.test.ts new file mode 100644 index 0000000000..a04bef60d8 --- /dev/null +++ b/packages/agent-runtime/src/system-prompt/__tests__/truncate-file-tree.test.ts @@ -0,0 +1,178 @@ +import { describe, expect, test } from 'bun:test' +import { truncateFileTreeBasedOnTokenBudget } from '../truncate-file-tree' +import type { + FileTreeNode, + ProjectFileContext, +} from '@codebuff/common/util/file' + +const mockLogger = { + debug: () => {}, + info: () => {}, + warn: () => {}, + error: () => {}, +} as any + +describe('truncateFileTreeBasedOnTokenBudget', () => { + test('returns none truncation level when within budget', () => { + const fileTree: FileTreeNode[] = [ + { + name: 'index.ts', + type: 'file', + filePath: 'src/index.ts', + lastReadTime: 0, + }, + { + name: 'util.ts', + type: 'file', + filePath: 'src/util.ts', + lastReadTime: 0, + }, + ] + + const fileContext = { + fileTree, + fileTokenScores: { + 'src/index.ts': { main: 10 }, + }, + } as unknown as ProjectFileContext + + const result = truncateFileTreeBasedOnTokenBudget({ + fileContext, + tokenBudget: 5000, + logger: mockLogger, + }) + + expect(result.truncationLevel).toBe('none') + expect(result.printedTree).toContain('index.ts') + expect(result.printedTree).toContain('util.ts') + expect(result.tokenCount).toBeGreaterThan(0) + expect(result.tokenCount).toBeLessThanOrEqual(5000) + }) + + test('filters out unimportant build directories and files', () => { + const fileTree: FileTreeNode[] = [ + { + name: 'src', + type: 'directory', + filePath: '/project/src/', + children: [ + { + name: 'main.ts', + type: 'file', + filePath: '/project/src/main.ts', + lastReadTime: 0, + }, + { + name: 'bundle.min.js', + type: 'file', + filePath: '/project/src/bundle.min.js', + lastReadTime: 0, + }, + ], + }, + { + name: 'dist', + type: 'directory', + filePath: '/project/dist/', + children: [ + { + name: 'out.js', + type: 'file', + filePath: '/project/dist/out.js', + lastReadTime: 0, + }, + ], + }, + ] + + const fileContext = { + fileTree, + fileTokenScores: {}, + } as unknown as ProjectFileContext + + const result = truncateFileTreeBasedOnTokenBudget({ + fileContext, + tokenBudget: 5000, + logger: mockLogger, + }) + + expect(result.printedTree).toContain('main.ts') + expect(result.printedTree).not.toContain('bundle.min.js') + expect(result.printedTree).not.toContain('dist') + }) + + test('truncates depth-based when token budget is very small', () => { + const fileTree: FileTreeNode[] = Array.from({ length: 100 }, (_, i) => ({ + name: `file_${i}.ts`, + type: 'file', + filePath: `src/deep/nested/sub/path/file_${i}.ts`, + lastReadTime: 0, + })) + + const fileContext = { + fileTree, + fileTokenScores: {}, + } as unknown as ProjectFileContext + + const result = truncateFileTreeBasedOnTokenBudget({ + fileContext, + tokenBudget: 50, + logger: mockLogger, + }) + + expect(result.tokenCount).toBeLessThanOrEqual(150) + expect(result.truncationLevel).toBe('depth-based') + }) + + test('filters out compiled binary and library files including .so', () => { + const fileTree: FileTreeNode[] = [ + { + name: 'app.exe', + type: 'file', + filePath: '/project/bin/app.exe', + lastReadTime: 0, + }, + { + name: 'libfoo.so', + type: 'file', + filePath: '/project/lib/libfoo.so', + lastReadTime: 0, + }, + { + name: 'libbar.dll', + type: 'file', + filePath: '/project/lib/libbar.dll', + lastReadTime: 0, + }, + { + name: 'libbaz.lib', + type: 'file', + filePath: '/project/lib/libbaz.lib', + lastReadTime: 0, + }, + { + name: 'valid.ts', + type: 'file', + filePath: '/project/src/valid.ts', + lastReadTime: 0, + }, + ] + + const fileContext = { + fileTree, + fileTokenScores: {}, + } as unknown as ProjectFileContext + + const result = truncateFileTreeBasedOnTokenBudget({ + fileContext, + tokenBudget: 5000, + logger: mockLogger, + }) + + expect(result.printedTree).toContain('valid.ts') + expect(result.printedTree).not.toContain('libfoo.so') + expect(result.printedTree).not.toContain('app.exe') + expect(result.printedTree).not.toContain('libbar.dll') + expect(result.printedTree).not.toContain('libbaz.lib') + }) +}) diff --git a/packages/agent-runtime/src/system-prompt/truncate-file-tree.ts b/packages/agent-runtime/src/system-prompt/truncate-file-tree.ts index 01e26b6f79..e6f4ba0e34 100644 --- a/packages/agent-runtime/src/system-prompt/truncate-file-tree.ts +++ b/packages/agent-runtime/src/system-prompt/truncate-file-tree.ts @@ -4,7 +4,7 @@ import { } from '@codebuff/common/util/file' import { sampleSizeWithSeed } from '@codebuff/common/util/random' -import { countTokens, countTokensJson } from '../util/token-counter' +import { countTokens } from '../util/token-counter' import type { Logger } from '@codebuff/common/types/contracts/logger' import type { @@ -32,7 +32,7 @@ export const truncateFileTreeBasedOnTokenBudget = (params: { const filteredTree = removeUnimportantFiles(fileTree) const treeWithTokens = printFileTreeWithTokens(filteredTree, fileTokenScores) - const treeWithTokensCount = countTokensJson(treeWithTokens) + const treeWithTokensCount = countTokens(treeWithTokens) if (treeWithTokensCount <= tokenBudget) { return { @@ -43,14 +43,14 @@ export const truncateFileTreeBasedOnTokenBudget = (params: { } const printedFilteredTree = printFileTree(filteredTree) - const filteredTreeNoTokensCount = countTokensJson(printedFilteredTree) + const filteredTreeNoTokensCount = countTokens(printedFilteredTree) if (filteredTreeNoTokensCount <= tokenBudget) { const filteredTreeWithTokens = printFileTreeWithTokens( filteredTree, fileTokenScores, ) - const filteredTreeWithTokensCount = countTokensJson(filteredTreeWithTokens) + const filteredTreeWithTokensCount = countTokens(filteredTreeWithTokens) if (filteredTreeWithTokensCount <= tokenBudget) { if (DEBUG) { logger.debug( @@ -120,11 +120,8 @@ export const truncateFileTreeBasedOnTokenBudget = (params: { // Sample 30 random files and count their tokens together const sampleCount = Math.min(30, sortedFiles.length) - const sampleFiles = sampleSizeWithSeed( - sortedFiles, - sampleCount, - JSON.stringify(sortedFiles) + JSON.stringify(sampleCount), - ) + const sampleSeed = `${sortedFiles.length}:${sampleCount}:${sortedFiles[0]?.path ?? ''}:${sortedFiles[sortedFiles.length - 1]?.path ?? ''}` + const sampleFiles = sampleSizeWithSeed(sortedFiles, sampleCount, sampleSeed) const sampleText = sampleFiles.map((f) => f.node.name).join(' ') const sampleTokens = countTokens(sampleText) @@ -168,7 +165,7 @@ export const truncateFileTreeBasedOnTokenBudget = (params: { .filter((n): n is FileTreeNode => n !== null) currentPrintedTree = printFileTree(currentTree) - currentTokenCount = countTokensJson(currentPrintedTree) + currentTokenCount = countTokens(currentPrintedTree) // Safety check - if we're not making progress, break if (currentTokenCount >= previousTokenCount) { @@ -241,7 +238,7 @@ function pruneFileTokenScores(params: { .sort((a, b) => a.score - b.score) let printedTree = printFileTreeWithTokens(fileTree, fileTokenScores) - let totalTokens = countTokensJson(printedTree) + let totalTokens = countTokens(printedTree) if (totalTokens <= tokenBudget) { return { pruned: fileTokenScores, printedTree, tokenCount: totalTokens } @@ -263,7 +260,7 @@ function pruneFileTokenScores(params: { let index = initialKeepIndex printedTree = printFileTreeWithTokens(fileTree, pruned) - totalTokens = countTokensJson(printedTree) + totalTokens = countTokens(printedTree) while (totalTokens > tokenBudget && index < sortedTokens.length) { const remainingToRemove = totalTokens - tokenBudget @@ -282,7 +279,7 @@ function pruneFileTokenScores(params: { // Note: The below function can take a while, so we optimized to have few loop iterations. printedTree = printFileTreeWithTokens(fileTree, pruned) - totalTokens = countTokensJson(printedTree) + totalTokens = countTokens(printedTree) index += batchSize } @@ -309,9 +306,8 @@ const removeUnimportantFiles = (fileTree: FileTreeNode[]): FileTreeNode[] => { if (node.type === 'directory') { // Filter out common build/cache directories const dirPath = node.filePath.toLowerCase() - const isUnimportantDir = unimportantExtensions.some( - (ext) => - ext.startsWith('/') && ext.endsWith('/') && dirPath.includes(ext), + const isUnimportantDir = UNIMPORTANT_DIR_PATTERNS.some((dir) => + dirPath.includes(dir), ) if (isUnimportantDir) { return false @@ -323,15 +319,26 @@ const removeUnimportantFiles = (fileTree: FileTreeNode[]): FileTreeNode[] => { } const filePath = node.filePath.toLowerCase() - return !unimportantExtensions.some( - (ext) => !ext.startsWith('/') && filePath.endsWith(ext), - ) + return !UNIMPORTANT_EXTENSIONS.some((ext) => filePath.endsWith(ext)) } return fileTree.filter(shouldKeepFile) } -const unimportantExtensions = [ +const UNIMPORTANT_DIR_PATTERNS = [ + // Build output directories + '/dist/', + '/build/', + '/out/', + '/target/', + + // Package manager directories + '/node_modules/', + '/.venv/', + '/vendor/', +] as const + +const UNIMPORTANT_EXTENSIONS = [ // Generated JavaScript/TypeScript files '.min.js', '.min.css', @@ -356,17 +363,6 @@ const unimportantExtensions = [ '.gem', '.rbc', - // Build output directories - '/dist/', - '/build/', - '/out/', - '/target/', - - // Package manager directories - '/node_modules/', - '/.venv/', - '/vendor/', - // Logs and temporary files '.log', '.tmp', @@ -411,4 +407,4 @@ const unimportantExtensions = [ '.tiff', '.tif', '.webp', -] +] as const