ContextPilot: Teaching Agents for Proactive Context Management via Fine-grained RL
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
| 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 index search readChunk search readChunk readChunk search readChunk 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 index search readChunk note delete search readChunk updateNote delete readNote finish |
| 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. |
| Model | Length | NovelQA | 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 | 32K | 60.24 | 62.45 | 40.60 | - | - |
| RL-MemoryAgent-14B | 32K | 78.86 | 74.24 | 59.00 | - | - |
| ReadAgent-8B | 32K | 16.38 | 24.02 | 0.00 | - | - |
| ReadAgent-14B | 32K | 23.12 | 34.06 | 14.60 | - | - |
| Backbone | Method | BrowseComp | BrowseComp-ZH | GAIA | xBench-DS | Avg. |
| pass@3 | pass@3 | pass@1 | pass@1 | |||
| WebSailor-7B | ReAct | 11.33 1.03 | 25.47 0.98 | 31.07 0.79 | 34.00 0.82 | 25.47 |
| ReAct (w/ truncation) | 12.67 0.62 | 27.91 0.91 | 33.66 0.92 | 35.00 0.82 | 27.31 | |
| ReSum | 15.83 0.85 | 38.99 1.56 | 38.19 1.21 | 35.33 1.25 | 32.09 | |
| SUPO | 18.50 1.08 | 42.68 0.99 | 42.07 1.65 | 42.00 1.41 | 36.31 | |
| OpenSeeker | 16.50 2.12 | 41.87 1.02 | 41.42 1.21 | 43.33 0.94 | 35.78 |
| Tool Design | NovelQA | 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 |
| Training method | NovelQA | 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) |
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 |
| NarrativeQA | 100 | 81 | 778 | |
| RL | LongBench-v2 | 488 | – | – |
| OpenSeeker-v1 (sampled) | 1,000 | – | – |
| 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. |
| Model | Training method | NovelQA | 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 |
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