Skip to content

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

Description

@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

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


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.

Activity

  1. alexhoshina commented on Mar 12, 2026

    @alexhoshina
    Collaborator

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


    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.

  2. is-Xiaoen commented on Mar 12, 2026

    @is-Xiaoen
    ContributorAuthor

    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.

  3. alexhoshina commented on Mar 12, 2026

    @alexhoshina
    Collaborator

    所以我们需要四个上下文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?

  4. is-Xiaoen commented on Mar 12, 2026

    @is-Xiaoen
    ContributorAuthor

    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.

  5. is-Xiaoen commented on Mar 19, 2026

    @is-Xiaoen
    ContributorAuthor

    Yeah, the core proposals (boundary separation, proactive budget check, Turn-based compression) are covered by #1490 now. Closing this one.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Projects

    No projects

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions