From 663335fccc911003d604e7e8e9ff1c14f1df5f35 Mon Sep 17 00:00:00 2001 From: Mikey Date: Wed, 2 Sep 2026 14:34:45 -0700 Subject: [PATCH 1/2] Optimize truncateFileTree token estimation and sampling seed --- .../__tests__/truncate-file-tree.test.ts | 123 ++++++++++++++++++ .../src/system-prompt/truncate-file-tree.ts | 57 ++++---- 2 files changed, 151 insertions(+), 29 deletions(-) create mode 100644 packages/agent-runtime/src/system-prompt/__tests__/truncate-file-tree.test.ts 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..0cf5f7f80f --- /dev/null +++ b/packages/agent-runtime/src/system-prompt/__tests__/truncate-file-tree.test.ts @@ -0,0 +1,123 @@ +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') + }) +}) 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..1ae750ddb2 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,10 +120,11 @@ export const truncateFileTreeBasedOnTokenBudget = (params: { // Sample 30 random files and count their tokens together const sampleCount = Math.min(30, sortedFiles.length) + const sampleSeed = `${sortedFiles.length}:${sampleCount}:${sortedFiles[0]?.path ?? ''}:${sortedFiles[sortedFiles.length - 1]?.path ?? ''}` const sampleFiles = sampleSizeWithSeed( sortedFiles, sampleCount, - JSON.stringify(sortedFiles) + JSON.stringify(sampleCount), + sampleSeed, ) const sampleText = sampleFiles.map((f) => f.node.name).join(' ') const sampleTokens = countTokens(sampleText) @@ -168,7 +169,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 +242,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 +264,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 +283,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 +310,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 +323,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 +367,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', @@ -394,7 +394,6 @@ const unimportantExtensions = [ '.exe', '.dll', '.lib', - '.so', // Media and binary files '.jpg', @@ -411,4 +410,4 @@ const unimportantExtensions = [ '.tiff', '.tif', '.webp', -] +] as const From 53971a306baa4015434f65db2070318b74aa8fab Mon Sep 17 00:00:00 2001 From: Mikey Date: Thu, 3 Sep 2026 09:56:39 -0700 Subject: [PATCH 2/2] Restore .so in UNIMPORTANT_EXTENSIONS and add binary filtering test --- .../__tests__/truncate-file-tree.test.ts | 57 ++++++++++++++++++- .../src/system-prompt/truncate-file-tree.ts | 7 +-- 2 files changed, 58 insertions(+), 6 deletions(-) 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 index 0cf5f7f80f..a04bef60d8 100644 --- 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 @@ -1,6 +1,9 @@ import { describe, expect, test } from 'bun:test' import { truncateFileTreeBasedOnTokenBudget } from '../truncate-file-tree' -import type { FileTreeNode, ProjectFileContext } from '@codebuff/common/util/file' +import type { + FileTreeNode, + ProjectFileContext, +} from '@codebuff/common/util/file' const mockLogger = { debug: () => {}, @@ -120,4 +123,56 @@ describe('truncateFileTreeBasedOnTokenBudget', () => { 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 1ae750ddb2..e6f4ba0e34 100644 --- a/packages/agent-runtime/src/system-prompt/truncate-file-tree.ts +++ b/packages/agent-runtime/src/system-prompt/truncate-file-tree.ts @@ -121,11 +121,7 @@ export const truncateFileTreeBasedOnTokenBudget = (params: { // Sample 30 random files and count their tokens together const sampleCount = Math.min(30, sortedFiles.length) const sampleSeed = `${sortedFiles.length}:${sampleCount}:${sortedFiles[0]?.path ?? ''}:${sortedFiles[sortedFiles.length - 1]?.path ?? ''}` - const sampleFiles = sampleSizeWithSeed( - sortedFiles, - sampleCount, - sampleSeed, - ) + const sampleFiles = sampleSizeWithSeed(sortedFiles, sampleCount, sampleSeed) const sampleText = sampleFiles.map((f) => f.node.name).join(' ') const sampleTokens = countTokens(sampleText) @@ -394,6 +390,7 @@ const UNIMPORTANT_EXTENSIONS = [ '.exe', '.dll', '.lib', + '.so', // Media and binary files '.jpg',