From d98656ca06084c7d7a85d0ab760f8f1883cca041 Mon Sep 17 00:00:00 2001 From: Keyjey101 Date: Mon, 31 Aug 2026 23:58:00 +0500 Subject: [PATCH] Add context pruning with a token budget --- README.md | 3 ++ agent/agent.js | 9 +++- agent/config.template.js | 3 ++ agent/context.js | 105 +++++++++++++++++++++++++++++++++++++++ agent/usage.md | 2 + ide/binds.js | 1 + ide/chat.js | 6 +++ ide/ide.js | 3 ++ start.js | 1 + 9 files changed, 132 insertions(+), 1 deletion(-) create mode 100644 agent/context.js diff --git a/README.md b/README.md index bffa55e..3e3ce4c 100644 --- a/README.md +++ b/README.md @@ -52,6 +52,9 @@ tools/{name}/ {name}.json + {name}.js - In a git repo, in-project write/edit/bash calls are auto-approved; outer paths still prompt - Step limit and tool-output truncation +- Context budget: when the conversation grows past + `CONTEXT_TOKEN_BUDGET` estimated tokens, older tool results are + pruned to stubs while the recent ones stay intact ## Safety diff --git a/agent/agent.js b/agent/agent.js index f587a28..93f31e1 100644 --- a/agent/agent.js +++ b/agent/agent.js @@ -5,6 +5,7 @@ const path = require('node:path'); const { isError, isHashObject, jsonParse } = require('metautil'); +const { DEFAULT_BUDGET, pruneMessages } = require('./context.js'); const { registry } = require('./tools.js'); const INSTRUCTIONS_FILE = path.join(__dirname, 'instructions.md'); @@ -96,6 +97,7 @@ const runAgent = async (options) => { const { task, provider, permissions } = options; const { maxSteps = 30, instructions = INSTRUCTIONS } = options; const { onEvent, priorMessages } = options; + const { contextBudget = DEFAULT_BUDGET } = options; const emit = async (type, data = {}) => { await onEvent?.({ type, ...data }); @@ -110,7 +112,12 @@ const runAgent = async (options) => { const loaded = [...registry.values()]; const tools = loaded.map((tool) => tool.definition); - const response = await provider.respond({ messages, tools }); + const request = pruneMessages(messages, contextBudget); + if (request.pruned > 0) await emit('prune', { pruned: request.pruned }); + const response = await provider.respond({ + messages: request.messages, + tools, + }); const message = response.choices?.[0]?.message; if (!message) throw new Error('Model returned no message.'); diff --git a/agent/config.template.js b/agent/config.template.js index 097f9ac..01102ce 100644 --- a/agent/config.template.js +++ b/agent/config.template.js @@ -7,4 +7,7 @@ module.exports = { BASE_URL: 'https://generativelanguage.googleapis.com/v1beta/openai/', MODEL: 'gemini-3.5-flash-lite', FALLBACK_MODELS: ['gemini-3.5-flash', 'gemini-3.6-flash'], + // Estimated token budget for the agent context. Older tool results + // are pruned to stubs when the conversation grows past it. + CONTEXT_TOKEN_BUDGET: 80_000, }; diff --git a/agent/context.js b/agent/context.js new file mode 100644 index 0000000..5bd375a --- /dev/null +++ b/agent/context.js @@ -0,0 +1,105 @@ +'use strict'; + +const CHARS_PER_TOKEN = 4; +const DEFAULT_BUDGET = 80_000; +const KEEP_RECENT_RESULTS = 4; +const LABEL_CHARS = 60; + +const estimateTokens = (text) => { + const chars = text?.length ?? 0; + return Math.ceil(chars / CHARS_PER_TOKEN); +}; + +const clipLabel = (text) => { + const flat = text.replaceAll('\n', ' ').trim(); + if (flat.length <= LABEL_CHARS) return flat; + return `${flat.slice(0, LABEL_CHARS - 1)}…`; +}; + +const argLabel = (argsText) => { + try { + const args = JSON.parse(argsText ?? '{}'); + for (const value of Object.values(args)) { + if (typeof value === 'string' && value !== '') return clipLabel(value); + } + } catch { + // keep the stub short when arguments are not valid JSON + } + return ''; +}; + +const contentText = (content) => { + if (typeof content === 'string') return content; + return JSON.stringify(content ?? ''); +}; + +const messageTokens = (message) => { + let text = contentText(message?.content); + for (const call of message?.tool_calls ?? []) { + text += call.function?.arguments ?? ''; + } + return estimateTokens(text); +}; + +const contextTokens = (messages) => + messages.reduce((sum, message) => sum + messageTokens(message), 0); + +const toolCallInfo = (messages) => { + const info = new Map(); + for (const message of messages) { + for (const call of message?.tool_calls ?? []) { + info.set(call.id, { + name: call.function?.name ?? 'tool', + label: argLabel(call.function?.arguments), + }); + } + } + return info; +}; + +const stubText = (info) => { + const source = info ? `${info.name} ${info.label}`.trim() : 'tool'; + return ( + `[Pruned ${source}: the older tool result was removed to stay within ` + + 'the context budget. Run the tool again if you still need it.]' + ); +}; + +const prunableIndexes = (messages) => { + const toolIndexes = []; + messages.forEach((message, index) => { + if (message?.role === 'tool') toolIndexes.push(index); + }); + const keepFrom = Math.max(0, toolIndexes.length - KEEP_RECENT_RESULTS); + return toolIndexes.slice(0, keepFrom); +}; + +const pruneMessages = (messages, budget = DEFAULT_BUDGET) => { + let total = contextTokens(messages); + const candidates = prunableIndexes(messages); + if (total <= budget || candidates.length === 0) { + return { messages, pruned: 0, tokens: total }; + } + + const info = toolCallInfo(messages); + const result = [...messages]; + let pruned = 0; + for (const index of candidates) { + if (total <= budget) break; + const original = messages[index]; + const stub = stubText(info.get(original.tool_call_id)); + result[index] = { ...original, content: stub }; + total += estimateTokens(stub) - messageTokens(original); + pruned += 1; + } + return { messages: result, pruned, tokens: total }; +}; + +module.exports = { + CHARS_PER_TOKEN, + DEFAULT_BUDGET, + KEEP_RECENT_RESULTS, + estimateTokens, + contextTokens, + pruneMessages, +}; diff --git a/agent/usage.md b/agent/usage.md index d194141..26875ae 100644 --- a/agent/usage.md +++ b/agent/usage.md @@ -36,4 +36,6 @@ API_KEY Required: OpenAI-compatible API key BASE_URL Required: Chat Completions endpoint MODEL Required: default model FALLBACK_MODELS Optional list used when the primary model returns 503 +CONTEXT_TOKEN_BUDGET Optional estimated token budget for the agent +context; older tool results are pruned to stubs past it (default 80000) See README.md. diff --git a/ide/binds.js b/ide/binds.js index 8358a24..a80f5f6 100644 --- a/ide/binds.js +++ b/ide/binds.js @@ -32,6 +32,7 @@ const EVENT_HANDLERS = { assistant: (ide, event) => ide.chat.addAgent(event.text), tool: (ide, event) => ide.onToolStart(event), result: (ide, event) => ide.onToolResult(event), + prune: (ide, event) => ide.chat.addPrune(event), }; const FOCUS_KEYS = { diff --git a/ide/chat.js b/ide/chat.js index 2eb4cb9..24ad0e7 100644 --- a/ide/chat.js +++ b/ide/chat.js @@ -77,6 +77,12 @@ class Chat { this.push({ kind: 'step', text: label, tone: 'faint' }); } + addPrune({ pruned }) { + const noun = pruned === 1 ? 'result' : 'results'; + const label = `pruned ${pruned} older tool ${noun}`; + this.push({ kind: 'step', text: label, tone: 'faint' }); + } + addTool(name, args, argsText) { const label = toolLabel(name, args, argsText); this.push({ diff --git a/ide/ide.js b/ide/ide.js index 254c66c..407312b 100644 --- a/ide/ide.js +++ b/ide/ide.js @@ -36,6 +36,7 @@ class Ide { ask: (description) => this.requestApproval(description), }); this.maxSteps = options.maxSteps; + this.contextBudget = options.contextBudget; this.tree = new FileTree(workspace.root); this.editor = new Editor(); this.terminal = new Terminal(workspace.root); @@ -534,6 +535,7 @@ class Ide { const provider = this.provider; const permissions = this.permissions; const maxSteps = this.maxSteps; + const contextBudget = this.contextBudget; const priorMessages = this.messages; const onEvent = (event) => this.handleEvent(event); const result = await runAgent({ @@ -541,6 +543,7 @@ class Ide { provider, permissions, maxSteps, + contextBudget, onEvent, priorMessages, }); diff --git a/start.js b/start.js index 09edc7a..f021211 100755 --- a/start.js +++ b/start.js @@ -63,6 +63,7 @@ const launchIde = async (root) => { await startIde({ maxSteps: DEFAULT_MAX_STEPS, model: config.MODEL, + contextBudget: config.CONTEXT_TOKEN_BUDGET, }); };