feat(tool-loop): add token-budget awareness and tool output pruning #61

New Issue

2026-06-23T11:23:09Z

CoreRasurae commented

2026-06-23 11:23:09 +00:00

Metadata

Commit Message: feat(tool-loop): add token-budget awareness and tool output pruning
Branch: feature/m2-tool-budget

Background and Context

The multi-turn tool-call loop (implemented in #59/#60) accumulates all AIMessages and ToolMessages into a single messages list across every round of the loop. Every subsequent ainvoke() call sends the entire accumulated conversation. This accumulation is architecturally correct and essential for effective LLM reasoning: the model must see its prior reasoning (AIMessage content), prior tool decisions (tool_calls), and prior tool results (ToolMessages) to build coherent multi-step chains.

However, the current implementation has zero context management. The loop blindly accumulates messages until it hits _TOOL_MAX_ROUNDS (default 20) or the model stops making tool calls. Two distinct failure modes arise:

Failure Mode 1: Context overflow from irrelevant tool output

The LLM reads a large file (e.g., 50KB log) via file_read but only needs 3 specific lines. The full 50KB enters context as a ToolMessage, crowding out room for the actual relevant data it still needs to gather from other files. The loop terminates prematurely not because the model made too many rounds, but because low-value tool output displaced space for high-value data. The final answer is degraded.

Failure Mode 2: Silent context window exhaustion

When accumulated context exceeds the model's context window, the provider either rejects the request with a context-length error or silently truncates from the beginning, causing the model to forget its system prompt or early reasoning.

Why Not Automatic Compaction?

Automatic compaction is harmful because the compressor cannot know what the LLM still needs. The correct approach is LLM-directed extraction: let the LLM extract only the relevant content from each tool output before it enters the main conversation context.

Architecture Decision Record

ADR-2031 (docs/adr/ADR-2031-tool-loop-token-budget-and-pruning.md, status: accepted) documents all specification extensions. Key decisions:

D-1: New config field token_budget_percent (number, default 0.85). Budget ceiling = token_budget_percent x model_context_window. At 75% consumption a warning is emitted; at 100% the budget-exhausted synthesis flow triggers.
D-2: New config fields allow_tool_output_pruning (boolean, default false) and pruning_model (string, optional). Initial pruning scope is file_read only — the tool most likely to produce context-busting output. Mechanism extends to other tools in future.
D-3: Tool output pruning pass with output_prune (boolean, default true) and output_prune_context (string, optional) meta-arguments injected into tool schemas. The main model can opt-out per call (output_prune: false) or provide a focused search directive (output_prune_context) that replaces the broader task context in the extraction prompt. Both meta-arguments are stripped before tool execution.
D-4: Budget-exhausted synthesis flow — injects synthesis prompt, permits one final tool-call round, then forces text-only response.
D-7: Normative marker format [PRUNE_INFO_START] / [PRUNE_INFO_END] and [PRUNE_OUTPUT_START] / [PRUNE_OUTPUT_END] for the pruning model's response. Marker strings reserved under §15.

Proposed Solution: Two Complementary Mechanisms

1. Token-Budget Awareness

Track estimated token consumption and detect when the context window is approaching its limit:

Token estimation: Before each ainvoke() call, estimate current token count using chat_model.get_num_tokens() when available, falling back to len(content) // 4 heuristic.
Budget warning: At 75% budget consumption, log a warning with current estimate vs. budget.
Budget exhaustion: When budget is exceeded, inject a synthesis prompt, allow one final tool-call round, strip tool definitions, then force a text-only response.

2. Tool Output Pruning

An implicit extraction pass scoped to file_read calls, gated behind allow_tool_output_pruning: true (opt-in, backward compatible).

Meta-arguments injected into tool schemas:

output_prune (boolean, default true): set to false to skip pruning and receive the full raw output.
output_prune_context (string, optional): a focused description from the main model of what it is looking for in the tool output. When provided, it replaces the broader task context in the extraction prompt.

Pruning response format:

[PRUNE_INFO_START]
<describe what was pruned, character counts, removed sections>
[PRUNE_INFO_END]
[PRUNE_OUTPUT_START]
<extracted relevant content>
[PRUNE_OUTPUT_END]

Design principles:

Only raw tool output passes through the pruning pass. System messages, user messages, conversation history, and AIMessages are sacrosanct.
The pruning model is a lightweight, stateless, single-turn call (by default using the same model; pruning_model config allows a lighter model).
On pruning call failure, the raw tool output is used as-is (graceful degradation).
No retroactive surgery on the messages list.

Acceptance Criteria

When token_budget_percent is configured and accumulated token estimate exceeds the budget ceiling, a warning is logged at 75% consumption and the budget-exhausted flow triggers.
When allow_tool_output_pruning: true, each file_read tool output is passed through the extraction LLM call before becoming a ToolMessage.
When a 50KB file_read result is pruned and the LLM needs only 3 lines, the resulting ToolMessage contains ~200 chars with pruning markers.
SystemMessage, HumanMessage, and AIMessage content is never passed through the pruning pass.
If the pruning LLM call fails, the raw tool output is used as-is and a warning is logged.
The main model can opt-out of pruning per call via output_prune: false in the tool call arguments.
The main model can provide a search directive via output_prune_context to focus the pruning model.
When budget is exhausted and no more space can be freed, the synthesis flow triggers.
All existing tool-calling behavior is preserved when no new config keys are provided (no regressions).
Tests (Behave BDD) cover: config validation, schema augmentation, marker parsing, pruning pass, budget exhaustion, output_prune opt-out, context window lookup, token estimation.
Integration tests (Robot Framework) cover: budget exhaustion, pruning-enabled agent, output_prune=false opt-out, budget config acceptance, backward compatibility.
ASV benchmarks cover: token estimation at multiple scales, marker parsing, schema augmentation, full pruning pass.

Supporting Information

Commit: fbf62145301168bae04411c07f1ade875935609f (original multi-turn tool-call loop)
ADR: docs/adr/ADR-2031-tool-loop-token-budget-and-pruning.md (this feature)
PR: feat(tool-loop): add token-budget awareness and tool output pruning (#62)
Relevant code:
- src/cleveractors/agents/llm.py — new methods: _get_model_context_window, _estimate_token_count, _augment_tool_schemas_for_pruning, _parse_pruning_response, _run_pruning_pass
- src/cleveractors/agents/llm.py — tool loop modified with budget check, pruning pass, schema augmentation, meta-argument stripping
Model context windows (reference): GPT-4o = 128K, Claude 3.5 Sonnet = 200K, Gemini 1.5 Pro = 1M+ tokens

Subtasks

Add token_budget_percent, allow_tool_output_pruning, and pruning_model config keys to LLMAgent with validation
Implement _estimate_token_count(messages) using LangChain chat_model.get_num_tokens() or len(content) // 4 fallback
Implement _get_model_context_window() for known models with 128K default
Insert token-budget check before each ainvoke() in the tool loop
Implement _augment_tool_schemas_for_pruning() to inject output_prune and output_prune_context meta-arguments
Implement _run_pruning_pass() extraction pass with build_chat_model
Implement _parse_pruning_response() marker parsing
Implement pruning failure fallback: use raw output if pruning call fails, log warning
Strip output_prune and output_prune_context from args before tool execution
Implement budget-exhausted synthesis flow with one final tool-call round
Log warnings at 75% budget consumption, info on each pruning pass, errors at budget exhaustion
Tests (Behave): 21 new scenarios in features/llm_agent_tool_calling.feature
Tests (Robot): 5 new integration tests in robot/llm_token_budget.robot
Benchmark (ASV): 16 benchmarks in benchmarks/token_budget_benchmark.py
Coverage verified at 96.8% via nox -s coverage_report
All nox sessions pass: lint, typecheck, unit_tests, integration_tests, security_scan, dead_code
Update CHANGELOG.md
ADR: docs/adr/ADR-2031-tool-loop-token-budget-and-pruning.md

Definition of Done

This issue is complete when:

All subtasks above are completed and checked off.
A Git commit is created where the first line of the commit message matches the Commit Message in Metadata exactly, followed by a blank line, then additional lines providing relevant details about the implementation.
The commit is pushed to the remote on the branch matching the Branch in Metadata exactly.
The commit is submitted as a pull request to master, reviewed, and merged before this issue is marked done.

## Metadata - **Commit Message**: `feat(tool-loop): add token-budget awareness and tool output pruning` - **Branch**: `feature/m2-tool-budget` ## Background and Context The multi-turn tool-call loop (implemented in #59/#60) accumulates all AIMessages and ToolMessages into a single `messages` list across every round of the loop. Every subsequent `ainvoke()` call sends the **entire** accumulated conversation. This accumulation is architecturally correct and essential for effective LLM reasoning: the model must see its prior reasoning (AIMessage content), prior tool decisions (`tool_calls`), and prior tool results (ToolMessages) to build coherent multi-step chains. However, the current implementation has **zero context management**. The loop blindly accumulates messages until it hits `_TOOL_MAX_ROUNDS` (default 20) or the model stops making tool calls. Two distinct failure modes arise: ### Failure Mode 1: Context overflow from irrelevant tool output The LLM reads a large file (e.g., 50KB log) via `file_read` but only needs 3 specific lines. The full 50KB enters context as a ToolMessage, crowding out room for the *actual* relevant data it still needs to gather from other files. The loop terminates prematurely not because the model made too many rounds, but because **low-value tool output displaced space for high-value data**. The final answer is degraded. ### Failure Mode 2: Silent context window exhaustion When accumulated context exceeds the model's context window, the provider either rejects the request with a context-length error or silently truncates from the beginning, causing the model to forget its system prompt or early reasoning. ## Why Not Automatic Compaction? Automatic compaction is harmful because the compressor cannot know what the LLM still needs. The correct approach is **LLM-directed extraction**: let the LLM extract only the relevant content from each tool output before it enters the main conversation context. ## Architecture Decision Record **ADR-2031** (`docs/adr/ADR-2031-tool-loop-token-budget-and-pruning.md`, status: accepted) documents all specification extensions. Key decisions: - **D-1**: New config field `token_budget_percent` (number, default 0.85). Budget ceiling = `token_budget_percent x model_context_window`. At 75% consumption a warning is emitted; at 100% the budget-exhausted synthesis flow triggers. - **D-2**: New config fields `allow_tool_output_pruning` (boolean, default false) and `pruning_model` (string, optional). Initial pruning scope is **`file_read` only** — the tool most likely to produce context-busting output. Mechanism extends to other tools in future. - **D-3**: Tool output pruning pass with **`output_prune`** (boolean, default true) and **`output_prune_context`** (string, optional) meta-arguments injected into tool schemas. The main model can opt-out per call (`output_prune: false`) or provide a focused search directive (`output_prune_context`) that replaces the broader task context in the extraction prompt. Both meta-arguments are stripped before tool execution. - **D-4**: Budget-exhausted synthesis flow — injects synthesis prompt, permits one final tool-call round, then forces text-only response. - **D-7**: Normative marker format `[PRUNE_INFO_START] / [PRUNE_INFO_END]` and `[PRUNE_OUTPUT_START] / [PRUNE_OUTPUT_END]` for the pruning model's response. Marker strings reserved under §15. ## Proposed Solution: Two Complementary Mechanisms ### 1. Token-Budget Awareness Track estimated token consumption and detect when the context window is approaching its limit: - **Token estimation**: Before each `ainvoke()` call, estimate current token count using `chat_model.get_num_tokens()` when available, falling back to `len(content) // 4` heuristic. - **Budget warning**: At 75% budget consumption, log a warning with current estimate vs. budget. - **Budget exhaustion**: When budget is exceeded, inject a synthesis prompt, allow one final tool-call round, strip tool definitions, then force a text-only response. ### 2. Tool Output Pruning An **implicit extraction pass** scoped to `file_read` calls, gated behind `allow_tool_output_pruning: true` (opt-in, backward compatible). **Meta-arguments injected into tool schemas:** - `output_prune` (boolean, default true): set to `false` to skip pruning and receive the full raw output. - `output_prune_context` (string, optional): a focused description from the main model of what it is looking for in the tool output. When provided, it replaces the broader task context in the extraction prompt. **Pruning response format:** ``` [PRUNE_INFO_START] <describe what was pruned, character counts, removed sections> [PRUNE_INFO_END] [PRUNE_OUTPUT_START] <extracted relevant content> [PRUNE_OUTPUT_END] ``` **Design principles:** - Only raw tool output passes through the pruning pass. System messages, user messages, conversation history, and AIMessages are sacrosanct. - The pruning model is a lightweight, stateless, single-turn call (by default using the same model; `pruning_model` config allows a lighter model). - On pruning call failure, the raw tool output is used as-is (graceful degradation). - No retroactive surgery on the `messages` list. ## Acceptance Criteria 1. When `token_budget_percent` is configured and accumulated token estimate exceeds the budget ceiling, a warning is logged at 75% consumption and the budget-exhausted flow triggers. 2. When `allow_tool_output_pruning: true`, each `file_read` tool output is passed through the extraction LLM call before becoming a ToolMessage. 3. When a 50KB `file_read` result is pruned and the LLM needs only 3 lines, the resulting ToolMessage contains ~200 chars with pruning markers. 4. SystemMessage, HumanMessage, and AIMessage content is never passed through the pruning pass. 5. If the pruning LLM call fails, the raw tool output is used as-is and a warning is logged. 6. The main model can opt-out of pruning per call via `output_prune: false` in the tool call arguments. 7. The main model can provide a search directive via `output_prune_context` to focus the pruning model. 8. When budget is exhausted and no more space can be freed, the synthesis flow triggers. 9. All existing tool-calling behavior is preserved when no new config keys are provided (no regressions). 10. Tests (Behave BDD) cover: config validation, schema augmentation, marker parsing, pruning pass, budget exhaustion, output_prune opt-out, context window lookup, token estimation. 11. Integration tests (Robot Framework) cover: budget exhaustion, pruning-enabled agent, output_prune=false opt-out, budget config acceptance, backward compatibility. 12. ASV benchmarks cover: token estimation at multiple scales, marker parsing, schema augmentation, full pruning pass. ## Supporting Information - **Commit**: `fbf62145301168bae04411c07f1ade875935609f` (original multi-turn tool-call loop) - **ADR**: `docs/adr/ADR-2031-tool-loop-token-budget-and-pruning.md` (this feature) - **PR**: #62 - **Relevant code**: - `src/cleveractors/agents/llm.py` — new methods: `_get_model_context_window`, `_estimate_token_count`, `_augment_tool_schemas_for_pruning`, `_parse_pruning_response`, `_run_pruning_pass` - `src/cleveractors/agents/llm.py` — tool loop modified with budget check, pruning pass, schema augmentation, meta-argument stripping - **Model context windows** (reference): GPT-4o = 128K, Claude 3.5 Sonnet = 200K, Gemini 1.5 Pro = 1M+ tokens ## Subtasks - [x] Add `token_budget_percent`, `allow_tool_output_pruning`, and `pruning_model` config keys to LLMAgent with validation - [x] Implement `_estimate_token_count(messages)` using LangChain `chat_model.get_num_tokens()` or `len(content) // 4` fallback - [x] Implement `_get_model_context_window()` for known models with 128K default - [x] Insert token-budget check before each `ainvoke()` in the tool loop - [x] Implement `_augment_tool_schemas_for_pruning()` to inject `output_prune` and `output_prune_context` meta-arguments - [x] Implement `_run_pruning_pass()` extraction pass with `build_chat_model` - [x] Implement `_parse_pruning_response()` marker parsing - [x] Implement pruning failure fallback: use raw output if pruning call fails, log warning - [x] Strip `output_prune` and `output_prune_context` from args before tool execution - [x] Implement budget-exhausted synthesis flow with one final tool-call round - [x] Log warnings at 75% budget consumption, info on each pruning pass, errors at budget exhaustion - [x] Tests (Behave): 21 new scenarios in `features/llm_agent_tool_calling.feature` - [x] Tests (Robot): 5 new integration tests in `robot/llm_token_budget.robot` - [x] Benchmark (ASV): 16 benchmarks in `benchmarks/token_budget_benchmark.py` - [x] Coverage verified at 96.8% via `nox -s coverage_report` - [x] All nox sessions pass: lint, typecheck, unit_tests, integration_tests, security_scan, dead_code - [x] Update CHANGELOG.md - [x] ADR: `docs/adr/ADR-2031-tool-loop-token-budget-and-pruning.md` ## Definition of Done This issue is complete when: - All subtasks above are completed and checked off. - A Git commit is created where the **first line** of the commit message matches the Commit Message in Metadata exactly, followed by a blank line, then additional lines providing relevant details about the implementation. - The commit is pushed to the remote on the branch matching the **Branch** in Metadata exactly. - The commit is submitted as a **pull request** to `master`, reviewed, and **merged** before this issue is marked done.

CoreRasurae added this to the v2.1.0 milestone 2026-06-23 11:23:09 +00:00

CoreRasurae added the

labels 2026-06-23 11:23:59 +00:00

CoreRasurae changed title from ~~feat(tool-loop): add token-budget-aware context management in multi-turn tool-call loop~~ to feat(tool-loop): add token-budget awareness and LLM-directed context pruning

2026-06-23 12:47:52 +00:00

CoreRasurae changed title from ~~feat(tool-loop): add token-budget awareness and LLM-directed context pruning~~ to feat(tool-loop): add token-budget awareness and tool output pruning

2026-06-23 12:53:31 +00:00

CoreRasurae referenced this issue from a commit

2026-06-23 22:58:32 +00:00

feat(tool-loop): add token-budget awareness and tool output pruning

CoreRasurae referenced a pull request that will close this issue

2026-06-23 22:59:49 +00:00

feat(tool-loop): add token-budget awareness and tool output pruning #62

CoreRasurae added a new dependency 2026-06-23 23:01:27 +00:00

#62 feat(tool-loop): add token-budget awareness and tool output pruning

CoreRasurae referenced this issue from a commit

2026-06-24 13:43:22 +00:00

feat(tool-loop): add token-budget awareness and tool output pruning