Patrick Plate
92 lines
3.2 KiB
TypeScript
92 lines
3.2 KiB
TypeScript
import { parametersSchema as z, defineCustomTool } from "@roo-code/types"
|
|
import { statSync, readdirSync } from "fs"
|
|
import { join, relative } from "path"
|
|
|
|
export default defineCustomTool({
|
|
name: "context_budget",
|
|
description: "Estimate token cost of reading files or directories. Returns file count, total lines, chars, and estimated tokens. Helps decide whether to load full files or use targeted reads.",
|
|
parameters: z.object({
|
|
paths: z.array(z.string()).describe("File or directory paths to analyze"),
|
|
recursive: z.boolean().optional().describe("Recurse into directories. Default: true"),
|
|
extensions: z.array(z.string()).optional().describe("Filter by file extensions, e.g. ['.java', '.ts']. Default: all files"),
|
|
}),
|
|
async execute({ paths, recursive = true, extensions }) {
|
|
let totalFiles = 0
|
|
let totalLines = 0
|
|
let totalChars = 0
|
|
const breakdown: Array<{ path: string; lines: number; chars: number; tokens: number }> = []
|
|
|
|
function processFile(filePath: string) {
|
|
try {
|
|
const stat = statSync(filePath)
|
|
if (!stat.isFile()) return
|
|
if (extensions && !extensions.some(ext => filePath.endsWith(ext))) return
|
|
|
|
const { readFileSync } = require("fs")
|
|
const content = readFileSync(filePath, "utf-8")
|
|
const lines = content.split("\n").length
|
|
const chars = content.length
|
|
const tokens = Math.ceil(chars / 4)
|
|
|
|
totalFiles++
|
|
totalLines += lines
|
|
totalChars += chars
|
|
breakdown.push({ path: filePath, lines, chars, tokens })
|
|
} catch (e) {
|
|
// skip unreadable files
|
|
}
|
|
}
|
|
|
|
function processDir(dirPath: string, recurse: boolean) {
|
|
try {
|
|
const entries = readdirSync(dirPath, { withFileTypes: true })
|
|
for (const entry of entries) {
|
|
if (entry.name.startsWith(".") || entry.name === "node_modules" || entry.name === "target" || entry.name === ".git") continue
|
|
const fullPath = join(dirPath, entry.name)
|
|
if (entry.isFile()) {
|
|
processFile(fullPath)
|
|
} else if (entry.isDirectory() && recurse) {
|
|
processDir(fullPath, recurse)
|
|
}
|
|
}
|
|
} catch (e) {
|
|
// skip unreadable dirs
|
|
}
|
|
}
|
|
|
|
for (const p of paths) {
|
|
try {
|
|
const stat = statSync(p)
|
|
if (stat.isFile()) {
|
|
processFile(p)
|
|
} else if (stat.isDirectory()) {
|
|
processDir(p, recursive !== false)
|
|
}
|
|
} catch (e) {
|
|
breakdown.push({ path: p, lines: 0, chars: 0, tokens: 0 })
|
|
}
|
|
}
|
|
|
|
const totalTokens = Math.ceil(totalChars / 4)
|
|
|
|
// Sort by tokens descending, take top 15
|
|
breakdown.sort((a, b) => b.tokens - a.tokens)
|
|
const topFiles = breakdown.slice(0, 15)
|
|
|
|
const result = {
|
|
summary: {
|
|
files: totalFiles,
|
|
totalLines,
|
|
totalChars,
|
|
estimatedTokens: totalTokens,
|
|
warning: totalTokens > 50000 ? "⚠️ LARGE — will consume significant context budget" :
|
|
totalTokens > 20000 ? "⚠️ MEDIUM — consider targeted reads" :
|
|
"✅ Fits comfortably in context",
|
|
},
|
|
largestFiles: topFiles,
|
|
}
|
|
|
|
return JSON.stringify(result, null, 2)
|
|
},
|
|
})
|