Optimize truncateFileTree token estimation and sampling seed

Author: nordicnode

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')
+  })
+})

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