Squeez: Task-Conditioned Tool-Output Pruning for Coding Agents

Ádám Kovács

cs.SEApr 4, 2026

Squeez: Task-Conditioned Tool-Output Pruning for Coding Agents

Authors: Ádám Kovács

Organizations: KR Labs

Abstract

Coding agents repeatedly consume long tool observations even though only a small fraction of each observation matters for the next step. We study task-conditioned tool-output pruning: given a focused query and one tool output, return the smallest verbatim evidence block the agent should inspect next. We introduce a benchmark of 11,477 examples built from SWE-bench repository interactions and synthetic multi-ecosystem tool outputs, with a manually curated 618-example test set. We fine-tune Qwen 3.5 2B with LoRA and compare it against larger zero-shot models and heuristic pruning baselines. Our model reaches 0.86 recall and 0.80 F1 while removing 92% of input tokens, outperforming zero-shot Qwen 3.5 35B A3B by 11 recall points and all heuristic baselines by a wide margin.

Figures & tables

Panel of Figure 1 — Figure 1: System overview. The benchmark stores grounded spans over raw tool output, then derives a generative training target. At evaluation, predicted lines are compared to gold spans using recall, F 1 , and compression.

Source	Raw inputs	Released rows
SWE-derived	10,713	9,205
Synthetic positives	2,039	1,697
Synthetic negatives	—	575
Total	12,752	11,477

Table 1: Source composition of the released benchmark. Synthetic negatives are constructed by pairing mismatched queries and tool outputs.

Tool	Rows	Tool	Rows	Tool	Rows
read_file	3768	ls	347	build_output	168
grep	1330	type_check	317	docker_build	151
git_log	720	git_blame	291	make_cmake	133
python	698	npm_build	230	kubectl	123
test_output	546	tsc	229	cargo_build	122
curl	493	npm_install	227	go_build	120

Table 2: Tool-family counts in the released benchmark.

Tool	Rows	Avg. input	Avg. gold
read_file	3768	1677	84
grep	1330	779	19
git_log	720	161	11
python	698	60	28
curl	493	723	68
pip_install	441	438	79

Table 3: Largest tool families in the dataset. Token counts are averages over released examples.

Model	Prec.	Recall	Strict F 1	Exact	F 1	Compression
Squeez-2B	0.80	0.86	0.79	0.49	0.80	0.92
Qwen 3.5 35B A3B	0.74	0.75	0.70	0.39	0.73	0.92
Kimi K2	0.61	0.53	0.53	0.30	0.68	0.94
Qwen 3.5 2B (base)	0.42	0.53	0.41	0.19	0.55	0.82
BM25 (10%)	0.13	0.22	0.13	0.01	0.23	0.90
First-N (10%)	0.07	0.14	0.08	0.02	0.16	0.91

Table 4: Held-out test results. Compression is the fraction of input removed. Squeez-2B is LoRA-tuned; other generative models use the same prompt zero-shot.

Pattern	Query / Observation	Squeez-2B	Baseline error
Precise selection in structured output	git_log , 21 lines. Query asks for the commit relevant to the dimension-order change in xr.polyval .	Selects the single correct commit entry.	Qwen 35B chooses a plausible but wrong transpose-related commit; Qwen 2B base over-selects several polyval commits.
Failure-block extraction	service log , 176 lines. Query asks for the TLS handshake failure affecting the health-check request.	Returns the correct 5-line health-check failure block.	Qwen 35B selects a later payment-request TLS failure; Kimi K2 keeps only part of the correct block.
Correct empty prediction	docker_logs , 316 lines. Query asks for a numpy version conflict that is not present.	Returns empty output.	Qwen 35B generates explanatory text (“No relevant lines found…”); Qwen 2B base returns unrelated database errors.
Adjacent over-selection	build_output , 110 lines. Query asks for the Dockerfile syntax error on line 12.	Finds the Dockerfile error but also includes a nearby Python SyntaxError .	Qwen 35B misses the Dockerfile error entirely and selects only the Python block.

Table 5: Representative held-out examples. The fine-tuned model is strongest on precise hits, compact failure blocks, and true negatives; its residual errors are usually adjacent over-selection.

AI Coding Agents

AI coding agents use language models and software tools to plan, edit, test, and maintain code with varying levels of autonomy.

See all 327 papers

SWE-Pruner Pro: The Coder LLM Already Knows What to Prune
Jul 20, 2026Yuhang Wang, Yuling Shi, Shaoqiu Zhang +6LLM Pruning AI Coding Agents
LLM4SE Lab, Shanghai Jiao Tong University · Douyin Group
Paritok-4B: Intent-Conditioned Context Compression for Coding Agents
Aug 25, 2026Jiayu Shi, Luzhuo ChenContext Compaction for LLM Agents Context Compression
Paritok
Qwen3-Coder-Next Technical Report
Feb 28, 2026Ruisheng Cao, Mouxiang Chen, Jiawei Chen +17Code Language Models AI Coding Agents

Jul 20, 2026J/K move · Enter open · S save

SWE-Pruner Pro: The Coder LLM Already Knows What to Prune

Yuhang Wang, Yuling Shi, Shaoqiu Zhang, Jialiang Liang+5

LLM4SE Lab, Shanghai Jiao Tong University · Douyin Group

LLM Pruning AI Coding Agents LLM Context Management

Open paper AskPDF arXiv

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.

Abstract

AI Coding Agents

Explore similar work

SWE-Pruner Pro: The Coder LLM Already Knows What to Prune

Paritok-4B: Intent-Conditioned Context Compression for Coding Agents

Qwen3-Coder-Next Technical Report