[Agent refactor] Context management: boundaries, compression, and token budgeting

[is-Xiaoen]

Context

This addresses track 6 of the agent refactor (#1216):

define history / summary / runtime / system prompt boundaries
define compression triggers and strategies
define what belongs to session context and what does not

I've been working in this area through the session persistence track (#732, #1170) and spent time reading the compression and context-building code. Below is what I found, and a proposal for how to clarify these boundaries.

---

Current state

Context management is currently spread across three locations with implicit boundaries between them:

• context.go:BuildMessages() — assembles system prompt (cached static + dynamic) + summary + history + current message into []Message. Runs sanitizeHistoryForProvider() to drop orphaned tool pairs at read time.
• loop.go:maybeSummarize() — checks two conditions after each turn: len(history) > SummarizeMessageThreshold (default 20) or estimateTokens(history) > ContextWindow * SummarizeTokenPercent / 100. If either is true, fires a background goroutine to run summarizeSession().
• loop.go:forceCompression() — called reactively when the LLM returns a context-window error. Drops the oldest 50% of conversation messages, appends an emergency note to the system prompt.

There is no explicit model of how much context space is available, what fills it, or when compression should happen relative to the actual budget.

---

Specific problems

1. ContextWindow defaults to MaxTokens

In instance.go:227:

ContextWindow: maxTokens,

MaxTokens is the max output tokens (default 32768 in defaults.go:33), passed to the LLM as the max_tokens request parameter (loop.go:930). But ContextWindow should represent the model's input capacity — typically 128K+ for modern models.

Setting ContextWindow = maxTokens means:

• maybeSummarize threshold = 32768 * 75 / 100 = 24576 estimated tokens
• History gets summarized far too early, well before the model's actual context limit is reached
• Conversely, if a user raises max_tokens to a large value, summarization never triggers at all

PR #556 identified the same issue.

2. forceCompression can orphan tool pairs

forceCompression() slices conversation at mid = len(conversation) / 2 (loop.go:1355) without checking whether the cut falls between an assistant message with ToolCalls and its matching tool result messages.

The read-path defense (sanitizeHistoryForProvider at context.go:577) catches orphaned pairs at query time, but the stored session history remains corrupted — tool messages without their matching assistant predecessor, or assistant messages with tool_calls but no results following. PR #665 identified this gap.

3. Compression is reactive, not proactive

forceCompression only runs after the LLM already rejected the request with a context-window error (loop.go:1009-1027). This means:

• The user sees "Context window exceeded. Compressing history and retrying..." — disruptive
• The emergency drop is blunt — 50% of messages gone without summarization
• A failed LLM call is wasted (and billed) before we realize the context was too large

A proactive check before the LLM call would prevent this entirely for the common case.

4. Token estimation undercounts

estimateTokens() (loop.go:1691) only counts utf8.RuneCountInString(m.Content). It ignores:

• ToolCalls — function name + JSON arguments can be substantial (complex tool args easily add thousands of tokens)
• The system prompt — built separately in BuildMessages, not included in the history estimate
• Tool definitions — injected by the provider adapter, invisible to the estimator

The summarization threshold check in maybeSummarize compares against ContextWindow using this undercount, so the check is weaker than intended in both directions.

---

Proposal

The goal is to clarify existing implicit boundaries, not introduce new abstractions. This follows the refactor's "minimum concepts" rule — no new types unless the current code cannot be clarified without them.

A. Separate context_window from max_tokens in config

Add context_window as an explicit field in AgentDefaults. Default to 0, meaning "fall back to a safe default" (e.g. 131072). This lets ContextWindow and MaxTokens serve their actual distinct purposes:

• MaxTokens → max output tokens per LLM call
• ContextWindow → total input capacity of the model

A follow-up improvement could auto-detect context window from the provider, but that's not needed for the initial fix.

B. Compute the available history budget explicitly

After building the system prompt, we know the fixed overhead. The available space for history becomes a simple subtraction:

fixed = tokenEstimate(systemPrompt) + tokenEstimate(toolDefinitions)
reserve = maxTokens  // leave room for model output
historyBudget = contextWindow - fixed - reserve

This historyBudget replaces the current ContextWindow * SummarizeTokenPercent / 100 as the compression threshold. No new types — just making the arithmetic explicit and correct.

C. Proactive pre-call check

Before calling the LLM in runLLMIteration, estimate the total token cost of the assembled messages slice. If it exceeds contextWindow - reserve, run summarization before the call.

forceCompression stays as a last-resort fallback for cases where the estimate was too low. But it should stop being the primary compression path.

D. Tool-pair-aware truncation

When forceCompression or summarizeSession truncates history, ensure the cut point does not fall inside a [assistant+tool_calls, tool_result, ...] group. If it does, move the cut before the group start. This gives write-path protection to complement the existing read-path sanitization.

E. Include ToolCalls in token estimation

Extend estimateTokens to account for m.ToolCalls — serialize function name and arguments into the character count. A rough estimate is still better than ignoring them.

---

What this does NOT propose

• No new types or abstractions (no ContextBudget struct, no ContextManager interface)
• No changes to session storage format
• No changes to the SessionStore or BuildMessages API signatures
• No dependency on other refactor tracks ([Agent refactor]what an Agent is #1218 agent abstraction, [Agent refactor] Event-driven agent loop with hooks, interrupts, and steering #1316 event model)
• No changes to summarization prompt or LLM-based compression logic

This is a boundary-clarification and correctness track, not a feature expansion.

---

Relationship to other tracks

Intentionally independent. Context management operates on []providers.Message and integer token counts. It does not depend on the Agent abstraction (track 1), EventBus (track 3), persona assembly (track 4), or capability model (track 5). If the AgentLoop lifecycle changes (track 2/3), the call sites may shift, but the budget logic itself is unaffected.

---

Related issues and PRs

• Meta: Agent refactor #1216 — umbrella issue (track 6)
• [Refactor]: Agent System Refactor / [Refactor]: Agent系统重构 #772 — original refactor proposal (pain points 1–3, 5)
• fix(agent): decouple context_window from max_tokens #556 — context_window ≠ max_tokens (open, same root cause as A above)
• fix(agent): prevent history compression from orphaning tool_call/tool_result pairs #665 — orphaned tool pairs after compression (open, same issue as D above)
• Fix: improve history compression with retry logic and multi-byte character support #1167 — compression retry improvements (merged)
• feat(memory): JSONL-backed session persistence #732 — JSONL session persistence (merged)
• feat(session): integrate JSONL persistence into agent loop #1170 — SessionStore integration into agent loop (merged)

---

I'd be happy to take this on if it fits the refactor direction. My plan would be to start with a working note in docs/agent-refactor/context.md covering the boundary definitions, then follow up with implementation PRs targeting the refactor/agent branch.

[alexhoshina]

或许在此之前我们可以构建一个上下文的结构图，这可以让截断和压缩都更能正确的找到其边界。

---

Perhaps before this, we could construct a structural diagram of the context, which would allow both truncation and compression to more accurately identify their boundaries.

[is-Xiaoen]

Good idea — the boundary rules should be visual before we touch code.

1. Context Window Regions

Everything sent to the LLM consumes a single context window budget. The regions stack in this order:

graph TB
    A["<b>System Prompt</b> (fixed, cached)<br/>identity · rules · skills · memory · dynamic context · summary"]
    B["<b>Conversation History</b> (compressible)<br/>message groups 1 ... N — see diagram below"]
    C["<b>Current User Message</b>"]
    D["<b>Tool Definitions</b> (provider-injected)"]
    E["<b>Output Reserve</b> (max_tokens)"]

    A --> B --> C --> D --> E

    style A fill:#dbeafe,stroke:#2563eb
    style B fill:#fef3c7,stroke:#d97706
    style C fill:#d1fae5,stroke:#059669
    style D fill:#e0e7ff,stroke:#6366f1
    style E fill:#fce4ec,stroke:#e11d48

Loading

history_budget = context_window − system_tokens − current_msg_tokens − tool_def_tokens − max_tokens

Only the yellow region (conversation history) is compressible. Everything else is fixed overhead or reserved.

---

2. Message Groups and Safe Boundaries

In the session history, messages form atomic groups per conversation turn. The session stores all intermediate messages (AddFullMessage in loop.go:1130 for assistant+tool_calls, loop.go:1264 for tool results).

Three group types exist:

graph TB
    subgraph Turn1["Turn 1 — Simple Exchange"]
        U1["user: hello"]:::user --> A1["assistant: hi"]:::asst
    end

    subgraph Turn2["Turn 2 — Tool Call"]
        U2["user: search X"]:::user --> A2["assistant + TC: search"]:::tc
        A2 -->|"atomic"| TR2["tool: result"]:::tool
        TR2 -->|"atomic"| AF2["assistant: found X"]:::asst
    end

    subgraph Turn3["Turn 3 — Chained Tool Calls"]
        U3["user: save and notify"]:::user --> A3["assistant + TC: save"]:::tc
        A3 -->|"atomic"| TR3["tool: saved"]:::tool
        TR3 -->|"atomic"| A3b["assistant + TC: notify"]:::tc
        A3b -->|"atomic"| TR3b["tool: sent"]:::tool
        TR3b -->|"atomic"| AF3["assistant: done"]:::asst
    end

    A1 ==>|"SAFE"| U2
    AF2 ==>|"SAFE"| U3
    AF3 ==>|"SAFE"| MORE["..."]

    classDef user fill:#dbeafe,stroke:#2563eb
    classDef asst fill:#d1fae5,stroke:#059669
    classDef tc fill:#fef3c7,stroke:#d97706
    classDef tool fill:#e5e7eb,stroke:#6b7280

Loading

Boundary rule: a SAFE cut point (thick arrow) exists only after a final assistant (no tool_calls) and before the next user. The atomic edges inside tool sequences must never be broken — splitting them creates orphaned tool_call/tool_result pairs that trigger provider 400 errors (Anthropic, DeepSeek, Zhipu).

To find the nearest safe cut point at a given index:

1. Scan backward — look for a user message whose predecessor is an assistant without tool_calls (or history start)
2. If none found, scan forward with the same rule
3. Cut before the user message at the boundary

---

3. Boundary-Aware Compression Flow

All three compression paths (proactive, emergency, async) should use the same safe-boundary detection:

flowchart TD
    START["Build messages for LLM"] --> EST["Estimate total tokens<br/>(system + history + current + tools)"]
    EST --> CHECK{"total exceeds<br/>context_window − max_tokens ?"}
    CHECK -->|"No"| CALL["Call LLM"]
    CHECK -->|"Yes"| PRO["Proactive: summarize oldest<br/>groups at safe boundary"]
    PRO --> REBUILD["Rebuild messages"]
    REBUILD --> EST

    CALL --> RESP{"Response?"}
    RESP -->|"Success"| SAVE["Save to session"]
    RESP -->|"Context error"| FORCE["forceCompression:<br/>truncate at safe boundary"]
    FORCE --> REBUILD2["Rebuild messages"]
    REBUILD2 --> CALL

    SAVE --> POST{"history_tokens<br/>above soft threshold ?"}
    POST -->|"No"| DONE["Done"]
    POST -->|"Yes"| ASYNC["Async: summarizeSession<br/>split at safe boundary"]
    ASYNC --> DONE

    style PRO fill:#fff3cd,stroke:#ffc107
    style FORCE fill:#f8d7da,stroke:#dc3545
    style ASYNC fill:#d1ecf1,stroke:#0dcaf0

Loading

Key change from the current code: all three boxes (forceCompression, summarizeSession, and the proposed proactive check) use the same boundary-finding logic to select cut points, instead of slicing at arbitrary indices like mid or len-4.

[alexhoshina]

所以我们需要四个上下文Builder吗？抽取一个公共的上下文Builder接口，为系统提示词、对话历史、用户消息、工具定义，四个区块创建各自的具体的Builder。还是全部使用同一个通用的Builder会更好？

---

So do we need four context Builders? Extract a common context Builder interface and create specific Builders for the four blocks: system prompt, conversation history, user message, and tool definitions. Or would it be better to use the same generic Builder for all?

[is-Xiaoen]

I'd lean toward neither — the issue isn't how we build each region, but that we never budget for them. Let me explain.

Construction is already separated

Region	Current code	Where
System Prompt	BuildSystemPromptWithCache() + buildDynamicContext()	context.go:493
History	GetHistory() → sanitizeHistoryForProvider()	context.go:548
User Message	Message{Role:"user", Content:msg}	context.go:563
Tool Definitions	ToolRegistry.ToProviderDefs()	registry.go:261

Tool definitions are already a separate parameter in Chat(ctx, messages, toolDefs, model, opts) at loop.go:950 — they never flow through BuildMessages().

Why 4 Builders adds complexity without proportional benefit

A common interface would need to unify string, []Message, and []ToolDefinition — fundamentally different output types. And only history needs complex logic (compression, boundary detection, summarization). The other three regions are fixed-size pass-through data, meaning 3 out of 4 Builder implementations would be trivial wrappers — adding concepts without matching value.

From the refactor guide: "if a behavior can be expressed without a new abstraction, do not add one".

The actual gap: no budget awareness before calling LLM

runAgentLoop (loop.go:757)
 ├─ BuildMessages(history, ...) → messages       ← no size check
 └─ runLLMIteration (loop.go:876)
     ├─ ToProviderDefs() → toolDefs              ← no size check
     └─ Chat(messages, toolDefs, max_tokens)     ← send blind
          └─ context error → forceCompression    ← reactive fix only

We assemble everything, send it to the LLM, and only compress after the provider rejects it with a 400. What we need is a proactive check.

Proposed: budget-aware functions, no new types

Each region's token cost is computed independently (addressing the "separate concerns" intuition), but as pure functions — not interface methods:

history_budget = context_window
               − estimateTokens(system_prompt)
               − estimateTokens(current_msg)
               − estimateToolDefsTokens(tool_defs)
               − max_tokens

The proactive check slots in between BuildMessages() and Chat():

flowchart TD
    BUILD["BuildMessages + ToProviderDefs"] --> EST["Estimate each region's token cost"]
    EST --> CALC["history_budget =<br/>context_window minus fixed costs"]
    CALC --> CHECK{"history<br/>within budget?"}
    CHECK -->|"Yes"| CALL["Call LLM"]
    CHECK -->|"No"| COMP["Compress at safe boundary"]
    COMP --> REBUILD["Rebuild messages"]
    REBUILD --> CHECK
    CALL --> RESP{"Response?"}
    RESP -->|"OK"| SAVE["Save + async summarize"]
    RESP -->|"Context error"| FORCE["forceCompression fallback"]
    FORCE --> REBUILD2["Rebuild"]
    REBUILD2 --> CALL

    style COMP fill:#fff3cd,stroke:#ffc107
    style FORCE fill:#f8d7da,stroke:#dc3545

Loading

The proactive path (yellow) prevents most context errors before they happen. The emergency path (red) stays as a fallback for edge cases where estimation undershoots.

How this fits the refactor

• Makes the implicit constraint (context must fit) explicit and preventive — "if an existing boundary can be made explicit, do that first"
• Adds no new types or interfaces — "do not introduce a new concept unless it is strictly necessary"
• No restructuring of ContextBuilder or the provider call path — purely additive

One prerequisite: the budget calculation needs a correct context_window value. Currently instance.go:227 sets ContextWindow: maxTokens, which conflates the model's context window with the output generation limit. These should be separated (a config field with sensible default) before the budget check can be meaningful.

[yinwm]

#1490 merged, should close this issue? @is-Xiaoen @alexhoshina