🤖 AI Summary
This work addresses the computational overhead induced by long contexts in tool-augmented language agents by proposing an embedded pruning method that leverages contextual relevance signals implicitly captured in the internal representations generated by an encoding agent during tool output processing. For the first time, this approach performs row-wise pruning of tool outputs directly within the agent, eliminating the need for external classifiers. The method introduces a line-count-aware embedding mechanism coupled with a lightweight pruning head to enable fine-grained, row-level decisions on retention or removal. Experimental results demonstrate that the proposed technique reduces prompt and generation tokens by up to 39% across multiple benchmarks, while improving the SWE-Bench Verified pass rate by 3.8% on MiMo-V2-Flash and boosting accuracy by 2.2 points on the Oolong long-context benchmark.
📝 Abstract
Pruning long context for coding agents has been a vital technology for efficient context management. While existing context pruning methods such as SWE-Pruner realize this by attaching a separate code classifier, we find the agent itself encodes internal representations indicating the relevance of code context when reading tool output. Based on this finding, we propose SWE-Pruner Pro, which prunes tool outputs directly inside the agent. Concretely, a small head turns the agent's own internal representations into a keep-or-prune label for each line, with a length-aware embedding keyed to each tool output's line count. Across two open-weight backbones and four multi-turn benchmarks, SWE-Pruner Pro saves up to 39% of prompt and completion tokens while preserving task quality, with bounded inference overhead. Notably, on MiMo-V2-Flash SWE-Pruner Pro additionally raises the SWE-Bench Verified resolve rate by +3.8% and the long-context Oolong accuracy by +2.2 points.