ContextPilot: Teaching Agents for Proactive Context Management via Fine-grained RL

Zhuoshi Pan; Qizhi Pei; Junru Lu; Honglin Lin; H. Vicky Zhao; Di Yin; Xing Sun

cs.CLAug 28, 2026

ContextPilot: Teaching Agents for Proactive Context Management via Fine-grained RL

Authors: Zhuoshi Pan, Qizhi Pei, Junru Lu, Honglin Lin, H. Vicky Zhao, Di Yin, Xing Sun

Organizations: Tsinghua University · Tencent Youtu Lab · Shanghai AI Lab

Abstract

Long-horizon agentic tasks require large language models (LLMs) to iteratively retrieve, integrate, and maintain dispersed information across multi-turn interactions, but preserving all interaction histories leads to a continuously growing working context. Recent proactive context management methods allow models to edit their own working context with specialized tools, yet they still face three key limitations: (1) a limited toolset restricted to search, deletion, and summarization, with no support for global planning, long-term memory, and adaptive compression; (2) inefficient exploration that treats context management actions uniformly despite their heterogeneous impacts on final outcomes; and (3) coarse-grained credit assignment that assigns the final trajectory-level reward to all intermediate context editing actions during RL. To bridge these gaps, we introduce ContextPilot, a proactive context management framework for long-horizon agentic reasoning. Our approach systematically augments the toolset with planning, long-term memory, and soft context offloading tools. We further propose an RL method tailored for context management, which uses context and entropy variation to identify critical editing decisions for branch sampling and estimates action-level advantages from all branched trajectories that pass through the corresponding context editing action. Experiments on long-context QA and deep search tasks show that ContextPilot achieves stronger performance with a more compact working context, consistently outperforming existing baselines across various base models and benchmarks. Code is available at https://github.com/Tencent/ContextPilot.

Figures & tables

Figure 1: Overview of ContextPilot : (a) Extending the context management tools to support planning, long-term memory, and soft context offloading; (b) Improved RL training algorithm with context-aware partial rollout and (c) fine-grained snapshot-level credit assignment.

Figure 2: Impact of different tools on final outcomes: for each tool, we branch from it, sample

10

continuation rollouts, and report the standard deviation of their task success rates. Results are evaluated with Qwen3-8B on NovelQA [ 31 ] .

Correct Answer Weak Context Management Q: Name the reviews editor of a mid-20th-century university journal, given an article they wrote and their Dec-2023 role. Gold: Robert D. McBain Answer: Robert D. McBain ✓ Tool chain: analyze

\rightarrow

index

\rightarrow

\times 3

\rightarrow

readChunk

\rightarrow

\times 2

\rightarrow

readChunk

\rightarrow\cdots\rightarrow

readChunk

\rightarrow

\times 2

\rightarrow

readChunk

\rightarrow

finish

Wrong Answer Strong Context Management Q: Find the guitarist who, in 70s high school, founded a band known for conscious multilingual music, and died in 2003. Gold: Frank Leepa Answer: Ty Longley ✗ Tool chain: analyze

\rightarrow

index

\rightarrow

\rightarrow

readChunk

\rightarrow

note

\rightarrow

delete

\times 2

\rightarrow

\times 2

\rightarrow

readChunk

\rightarrow

updateNote

\rightarrow

delete

\times 2

\rightarrow

readNote

\rightarrow

finish

Figure 3: Case study on the mismatch between task correctness and context management quality. Cases are sampled from the traces of StateLM-8B [ 16 ] on BrowseComp+ [ 2 ] . Questions are condensed to save space.

Category	Tool Name	Function Description
Perception & Planning	analyzeText	Calculate context length.
	checkBudget	Check remaining token budget.
	plan	Propose a concise plan.
Information Retrieval	buildIndex	Build a searchable index.
	searchContext	Search for relevant content.
	readChunk	Load a specific context chunk.

Table 1: Context management tools of ContextPilot . Tools highlighted in purple are newly introduced.

Model	Length	NovelQA	$\infty$ Bench	LongMemEval-S	BrowseComp+	Avg.
Qwen3.5-397B-A17B (w/o tools)	256K	88.77	90.39	81.00	62.05	80.55
Qwen3.5-397B-A17B (w/ tools)	32K	91.94	92.13	83.60	80.96	87.16
RL-MemoryAgent-7B $\dagger$	32K	60.24	62.45	40.60	-	-
RL-MemoryAgent-14B $\dagger$	32K	78.86	74.24	59.00	-	-
ReadAgent-8B $\dagger$	32K	16.38	24.02	0.00	-	-
ReadAgent-14B $\dagger$	32K	23.12	34.06	14.60	-	-

Table 2: Performance comparison of ContextPilot against baseline methods on long-context QA tasks. We run each method three times and report the mean and standard deviation. Results with

\dagger

are from Liu et al. [16] .

Backbone	Method	BrowseComp	BrowseComp-ZH	GAIA	xBench-DS	Avg.
		pass@3	pass@3	pass@1	pass@1
WebSailor-7B	ReAct	11.33 $\pm$ 1.03	25.47 $\pm$ 0.98	31.07 $\pm$ 0.79	34.00 $\pm$ 0.82	25.47
	ReAct (w/ truncation)	12.67 $\pm$ 0.62	27.91 $\pm$ 0.91	33.66 $\pm$ 0.92	35.00 $\pm$ 0.82	27.31
	ReSum	15.83 $\pm$ 0.85	38.99 $\pm$ 1.56	38.19 $\pm$ 1.21	35.33 $\pm$ 1.25	32.09
	SUPO	18.50 $\pm$ 1.08	42.68 $\pm$ 0.99	42.07 $\pm$ 1.65	42.00 $\pm$ 1.41	36.31
	OpenSeeker	16.50 $\pm$ 2.12	41.87 $\pm$ 1.02	41.42 $\pm$ 1.21	43.33 $\pm$ 0.94	35.78

Table 3: Performance comparison of ContextPilot against baseline methods on deep search tasks. We run each method three times and report the mean and standard deviation.

Figure 4: Token usage per turn on BrowseComp.

Panel of Figure 5 — Figure 5: Tool category breakdown during Qwen3-8B RL training on $\infty$ Bench and NovelQA.

Panel of Figure 6 — Figure 6: Tool invocation failure rates and task success rate on $\infty$ Bench and NovelQA during RL training.

Tool Design	NovelQA	$\infty$ Bench	LME-S	BC+	Avg.
Original tools	88.90	85.15	74.00	63.49	77.89
+ Planning	89.76	87.23	78.50	65.66	80.29
+ Soft offloading	91.28	89.63	80.20	71.20	83.08
+ Long-term memory	91.94	92.13	83.60	80.96	87.16

Table 4: Ablation results of tool design using Qwen3.5-397B-A17B. LME-S and BC+ denote LongMemEval-S and BrowseComp+, respectively.

Training method	NovelQA	$\infty$ Bench	LME-S	BC+
SFT	82.56	71.03	60.67	48.84
GRPO	83.53 (+0.97)	72.78 (+1.75)	60.07 (-0.60)	50.96 (+2.12)
+ Entropy.	82.52 (-1.01)	73.07 (+0.29)	62.13 (+2.06)	49.64 (-1.32)
+ Context.	83.05 (+0.53)	73.94 (+0.87)	61.40 (-0.73)	51.08 (+1.44)
+ Fine-grained.	83.88 (+0.83)	75.25 (+1.31)	64.27 (+2.87)	54.18 (+3.10)

Table 5: Ablation results of the RL algorithm using Qwen3-8B. Results for Gemma4-E4B-it are in Table 8 (Appendix). “+ Entropy.”, “+ Context.”, and “+ Fine-grained.” denote RL training with entropy-based partial rollout, context-aware partial rollout, and fine-grained credit assignment. Numbers in parentheses indicate changes relative to the previous line.

Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

Stage	Source	Questions	Trajectories	Snapshots
SFT	NovelQA	3,096	2,987	50,691
SFT	NarrativeQA	100	81	778
RL	LongBench-v2	488	–	–
RL	OpenSeeker-v1 (sampled)	1,000	–	–

Appendix

Table 6: Statistics of the training data, where NovelQA, NarrativeQA, and LongBench-v2 are used for long-context QA and OpenSeeker-v1 is used for deep search.

Panel of Figure 7 — Figure 7: System prompt used by ContextPilot for long-context document QA.

Panel of Figure 8 — Figure 8: System prompt used by ContextPilot for deep search tasks.

Panel of Figure 9 — Figure 9: Prompt template used by the LLM-as-a-Judge for evaluating open-ended questions, following StateLM [ 16 ] .

Panel of Figure 10 — Figure 10: Token usage per turn on BrowseComp and BrowseComp-ZH.

Tool	Failure Reason
plan	The predicted summary is empty or unparseable.
searchContext	The search query is empty or malformed, or the model searches before a valid index has been built.
readChunk	The requested chunk id is out of range, deleted, or not associated with the current index.
readMultiChunks	Any requested chunk id is invalid, duplicated, deleted, or exceeds the allowed batch size.
note	The note content is empty, unparseable, or does not specify the value to be stored.
updateNote	The target note does not exist, or the update content is empty or unparseable.

Appendix

Table 7: Tool execution failure reason analysis.

Model	Training method	NovelQA	$\infty$ Bench	LongMemEval-S	BrowseComp+	Avg.
Qwen3-8B	SFT	82.56	71.03	60.67	48.84	65.78
	GRPO	83.53 (+0.97)	72.78 (+1.75)	60.07 (-0.60)	50.96 (+2.12)	66.84 (+1.06)
	+ Entropy-based partial rollout	82.52 (-1.01)	73.07 (+0.29)	62.13 (+2.06)	49.64 (-1.32)	66.84 (+0.00)
	+ Context-aware partial rollout	83.05 (+0.53)	73.94 (+0.87)	61.40 (-0.73)	51.08 (+1.44)	67.37 (+0.53)
	+ Fine-grained credit assignment	83.88 (+0.83)	75.25 (+1.31)	64.27 (+2.87)	54.18 (+3.10)	69.40 (+2.03)
Gemma4-E4B-it	SFT	66.80	55.02	55.07	42.05	54.74

Appendix

Table 8: Complete ablation results of the RL algorithm for Qwen3-8B and Gemma4-E4B-it. Numbers in parentheses indicate changes relative to the previous line.

LLM Context Management

Large language model context management studies how to select, compress, retrieve, and organize information within bounded model inputs.

See all 244 papers

Learning Agent-Compatible Context Management for Long-Horizon Tasks
May 29, 2026Lu Yi, Runlin Lei, Liuyi Yao +6Long-Horizon Agent Tasks LLM Context Management
1Renmin University of China · Work done during internship at Tongyi Lab, Alibaba Group · 2Tongyi Lab, Alibaba Group +2
LongSeeker: Elastic Context Orchestration for Long-Horizon Search Agents
May 6, 2026Yijun Lu, Rui Ye, Yuwen Du +3Web Search Agents Tool-Using Agents
Shanghai Jiao Tong University
Continuous Context Management
Sep 28, 2026William Hoy, Jingxuan Fan, Nurcin Celik +1Context Distillation Long-Horizon LLM Agents
University of Miami · Harvard University

May 29, 2026J/K move · Enter open · S save

Learning Agent-Compatible Context Management for Long-Horizon Tasks

Lu Yi, Runlin Lei, Liuyi Yao, Yuexiang Xie+5

1Renmin University of China · Work done during internship at Tongyi Lab, Alibaba Group · 2Tongyi Lab, Alibaba Group+2

Long-Horizon Agent Tasks LLM Context Management RL for LLM Agents

Open paper AskPDF arXiv

LLM agents increasingly face long-horizon tasks such as web search and deep research in real-world applications, where accumulated context can cause long-context degradation and reasoning failures. Prior work mitigates this through context management with agent-side context control or fixed strategies such as summarization, which require training the agent itself for adaptation - making it impractical for closed-source agents and ignoring that different agents may require different strategies. We introduce Adaptive Context Management (AdaCoM), which trains an external LLM to manage the context of a frozen agent through flexible modification actions and end-to-end reinforcement learning. Across diverse agents on web search and deep research benchmarks, AdaCoM substantially improves performance by preserving task constraints and progress while pruning stale content. The learned strategies reveal a Fidelity-Reliability Trade-off: agents with higher vanilla ReAct performance benefit from higher-fidelity context preservation, whereas lower-performing agents require more aggressive compression to stay within a reliable reasoning regime. Transfer experiments show that AdaCoM generalizes most effectively across agents with similar capability (measured by vanilla ReAct performance), suggesting a practical path toward reusable context managers for agent systems.

Abstract

LLM Context Management

Explore similar work

Learning Agent-Compatible Context Management for Long-Horizon Tasks

LongSeeker: Elastic Context Orchestration for Long-Horizon Search Agents

Continuous Context Management