【Java/Go后端手撸原生Agent（第七篇）：Token预算管理 + 滑动窗口上下文裁剪】

系列目录：Phase 1-2 基础框架 → Phase 3 状态机+文件读取 → Phase 4-6 原生FC+Judge校验 → Phase 5 多工具并行+BashTool → Phase 6 流式输出 → Phase 7 Token预算与上下文管理（本文）

上一篇链接：Java/Go后端手撸原生Agent（第六篇）：流式输出让Agent思考过程"看得见"

前言

前六篇我们完成了ReAct智能体的全部核心链路：状态机驱动、原生Function Calling、多工具并行、BashTool执行引擎、LLM-as-Judge校验、流式SSE输出。到这里Agent已经能跑通完整的"提问→思考→调工具→观察→再思考→回答→校验"闭环了。

但跑长对话时会遇到一个所有LLM应用都绕不开的硬约束------上下文窗口（Context Window）是有限的 。你的模型可能是8K、32K、128K窗口，但无论多大，总有上限。当前的ShortMemory是无限增长的：每轮对话加user+assistant+tool消息，BashTool一次find . -name "*.py"可能输出几千字符，对话五六个回合后messages数组就可能膨胀到爆，API返回context_length_exceeded错误直接崩溃。

本文解决这个问题。核心做两件事：

1. Token精确计数：不做中文/1.5这种粗糙估算，用tiktoken精确计算每条消息、tools schema、system prompt的token开销；
2. 滑动窗口裁剪 ：超预算时从最早的消息开始裁剪，但有一个关键约束------不能破坏tool_calls和tool结果消息的配对完整性，否则API会拒绝请求。

前置说明：

• 新增文件：agent/token_counter.py（token计数模块）；
• 改造文件：agent/memory.py（ShortMemory增加token感知和滑动窗口裁剪）；main.py（每轮THINKING前做预算检查）；
• 零改动文件：llm_client.py、tools/下所有工具、Judge校验逻辑------这些模块不需要感知token管理。

一、问题本质：token是LLM世界的"内存"

1.1 粗估为什么不行

很多Demo里"token估算"是这么做的：len(text) // 2（英文）或len(text) // 1.5（中文）。这种估算在只计算content文本时勉强能用，但在Agent场景下严重失准，因为每次API请求的token开销由四部分组成：

总token = system prompt + tools schema + messages历史 + 输出预留

其中tools schema是最容易被忽略的隐性开销。你的Agent有3个工具（calculator/read_file/bash），每个工具的function定义（name/description/parameters JSON schema）就要几百token。如果工具数量增长到10个，tools schema轻松破2000 token。用len(text)//2完全估算不到这部分。

而且OpenAI协议里每条消息有固定开销：<|im_start|>role\n + content + <|im_end|>\n，大约4 token/条。消息数量多了这部分也不能忽略。

1.2 裁剪的真正难点：不是删消息，是保护配对完整性

直觉上"上下文满了从前面删"不就行了？不行。OpenAI协议有硬约束：

• role=tool消息必须紧跟在包含对应tool_calls的assistant消息后面；
• 每条tool消息通过tool_call_id精确配对到assistant消息里的某个tool_call；
• 如果删了assistant(tool_calls)而留下它的tool结果消息，tool消息悬空，API报错；
• 如果删了tool结果而留下assistant(tool_calls)，assistant引用了不存在的工具结果，API同样报错。

举个例子，合法的消息序列：

user: "帮我读文件并计算"
assistant(tool_calls=[{id:call_1, name:read_file}, {id:call_2, name:calculator}]): "我先读文件再计算"
tool(tool_call_id=call_1): "文件内容是..."
tool(tool_call_id=call_2): "计算结果: 42"
user: "[Judge反馈] 你还没回答完整..."
assistant: "根据文件内容和计算结果..."

如果你只删前两条（user+assistant），留下两个tool消息在最前面，这两个tool没有对应的assistant(tool_calls)在前面，协议断裂。如果你只删第一个tool消息，留下assistant(tool_calls=call_1, call_2)和tool(call_2)，call_1的结果丢失，同样断裂。

所以裁剪单位不能是"一条消息"，必须是"一个消息块"。

1.3 三种策略的分层选择

上下文管理有三种策略，从轻到重：

策略	原理	可靠性	复杂度
滑动窗口（本文做）	从最早开始删完整消息块，超预算就删	确定性，不调LLM	低
摘要压缩（下一步做）	被删的消息用LLM总结成一段话，塞回历史开头	保留语义但丢失细节，依赖LLM	中
重要性排序/RAG（未来）	按消息与当前问题的相关性选择保留，向量检索	最优但复杂，需要向量数据库	高

本次只做滑动窗口。原因：它是确定性代码、零额外LLM调用、是所有更高级策略的基础（摘要压缩也需要先确定哪些消息要被摘要）。按"先验证再固化"原则，先用最简方案跑通，等真实遇到"裁掉了重要上下文"的场景再加摘要。

二、改造1：Token精确计数模块

2.1 tiktoken：OpenAI官方tokenizer

Token计数不需要自己写规则，用OpenAI官方的tiktoken库：

pip install tiktoken

tiktoken内置了GPT-4/GPT-3.5使用的BPE tokenizer（cl100k_base），大多数兼容OpenAI协议的模型（DeepSeek/Qwen/GLM等）也使用兼容的tokenizer，误差可接受。

2.2 计数规则

参考OpenAI官方cookbook，精确计数需要考虑：

• 每条消息固定4 token（role标记 + 开始/结束符）；
• content、tool_call的name/id/arguments分别序列化后计数；
• tools schema整体序列化为JSON后计数，每个tool额外6 token结构开销；
• 整个请求额外2 token前缀（assistant回复引导token）。

2.3 agent/token_counter.py完整代码

"""
Token计数模块。

为什么需要精确计数？
LLM的上下文窗口是硬限制（token数上限），超出会报context_length_exceeded错误。
粗估（"中文字符/1.5"）对content文本还行，但tool_calls JSON结构和tools schema的
开销不可忽略------10个工具的schema可能就占2000+ tokens。

计数规则参考OpenAI官方cookbook：
- 每条消息固定开销 ~4 tokens（<|im_start|>role\n...<|im_end|>\n）
- tool_calls序列化为JSON后按文本计数
- tools schema同样序列化后计数
- 最后有一个固定开销 ~2 tokens（assistant回复前缀）
"""
from __future__ import annotations

import json
from typing import Any, Optional

import tiktoken

# cl100k_base: GPT-4/GPT-3.5-turbo（0613及之后版本）使用的tokenizer
# 兼容OpenAI协议的大多数模型也用这个（DeepSeek/Qwen等）
_ENCODING = tiktoken.get_encoding("cl100k_base")

# 每条消息的固定开销token数（role标记+开始/结束符）
_TOKENS_PER_MESSAGE = 4
# 整个请求的固定前缀开销（assistant回复引导token）
_TOKENS_PER_REQUEST = 2
# tool_call固定开销（type/function/id字段名+结构开销）
_TOKENS_PER_TOOL_CALL = 8


def count_text(text: str) -> int:
    """计算一段文本的token数"""
    return len(_ENCODING.encode(text))


def count_tool_call(tool_call: dict[str, Any]) -> int:
    """
    计算单个tool_call的token数。
    tool_call是OpenAI格式: {"id": "call_xxx", "type": "function", "function": {"name": "...", "arguments": "..."}}
    """
    tokens = _TOKENS_PER_TOOL_CALL
    fn = tool_call.get("function", {})
    name = fn.get("name", "")
    arguments = fn.get("arguments", "")
    if name:
        tokens += count_text(name)
    if arguments:
        if isinstance(arguments, str):
            tokens += count_text(arguments)
        else:
            tokens += count_text(json.dumps(arguments, ensure_ascii=False))
    call_id = tool_call.get("id", "")
    if call_id:
        tokens += count_text(call_id)
    return tokens


def count_message(msg: dict[str, Any]) -> int:
    """计算单条消息的token数（含固定开销）"""
    tokens = _TOKENS_PER_MESSAGE

    role = msg.get("role", "")
    content = msg.get("content")
    if content:
        tokens += count_text(str(content))

    if role == "assistant":
        tool_calls = msg.get("tool_calls")
        if tool_calls:
            for tc in tool_calls:
                tokens += count_tool_call(tc)

    if role == "tool":
        tool_call_id = msg.get("tool_call_id", "")
        if tool_call_id:
            tokens += count_text(tool_call_id)

    return tokens


def count_messages(messages: list[dict[str, Any]]) -> int:
    """计算一组消息的总token数（不含tools schema和request前缀）"""
    total = 0
    for msg in messages:
        total += count_message(msg)
    return total


def count_tools_schema(tools: Optional[list[dict[str, Any]]]) -> int:
    """计算tools参数（Function Schema数组）的token数"""
    if not tools:
        return 0
    total = 0
    for tool in tools:
        total += count_text(json.dumps(tool, ensure_ascii=False))
        total += 6  # 每个tool的固定结构开销
    return total


def count_request(
    system_prompt: Optional[str],
    messages: list[dict[str, Any]],
    tools: Optional[list[dict[str, Any]]] = None,
) -> int:
    """
    计算整个API请求的token数（不含输出token）。
    总token = system_prompt开销 + messages开销 + tools schema开销 + 请求前缀开销
    """
    total = 0
    if system_prompt:
        total += _TOKENS_PER_MESSAGE
        total += count_text(system_prompt)
    total += count_messages(messages)
    total += count_tools_schema(tools)
    total += _TOKENS_PER_REQUEST
    return total

模块职责单一：只做计数，不做裁剪。每个函数粒度足够细，count_text/count_message/count_messages/count_tools_schema分别对应不同层级的计数需求，调用方按需组合。

三、改造2：ShortMemory增加滑动窗口裁剪

3.1 消息块的定义

这是核心设计。裁剪单位是"消息块"，块的规则：

第一条消息的类型	块组成
user消息	单条为一个块
assistant消息（无tool_calls）	单条为一个块
assistant消息（有tool_calls）	assistant + 紧随其后、tool_call_id匹配的所有tool消息 = 一个块
tool消息（异常）	单条删除（不应该出现在块首，属于断裂数据）

为什么assistant(tool_calls)后面的tool消息可能不是全部匹配？在正常流程中，assistant(tool_calls)后面紧跟的tool消息就是它的全部结果，且tool_call_id一一对应。但考虑到边缘情况（比如之前有异常删除导致历史中间出现断裂），算法只收集tool_call_id匹配的tool消息，遇到非tool消息或所有ID匹配完毕就停止。

3.2 _pop_oldest_block算法

def _pop_oldest_block(self) -> int:
    """删除最早的一个完整消息块，返回删除的token数"""
    if not self.history:
        return 0

    first = self.history[0]

    # 独立消息块：user 或 assistant（无tool_calls）
    if first["role"] == "user" or (first["role"] == "assistant" and not first.get("tool_calls")):
        self.history.pop(0)
        return count_message(first)

    # 工具调用块：assistant(tool_calls) + 其对应的tool结果消息
    if first["role"] == "assistant" and first.get("tool_calls"):
        tc_ids = {tc["id"] for tc in first["tool_calls"]}
        self.history.pop(0)  # 先删assistant
        removed = count_message(first)

        # 然后删除紧随其后、tool_call_id匹配的tool消息
        i = 0
        while i < len(self.history) and self.history[i]["role"] == "tool":
            msg_tc_id = self.history[i].get("tool_call_id")
            if msg_tc_id in tc_ids:
                removed_msg = self.history.pop(i)  # pop后下一个元素移到i位置，不要递增i
                removed += count_message(removed_msg)
                tc_ids.discard(msg_tc_id)
                if not tc_ids:
                    break  # 所有配对的tool消息都删完了
            else:
                i += 1
        return removed

    # 异常：tool消息在最前面（配对断裂），单独删除
    removed = self.history.pop(0)
    return count_message(removed)

注意pop(i)后不要递增i------pop之后原来的i+1位置元素移到了i位置，需要继续检查i位置。这和遍历list时删除元素的经典坑一样。

3.3 agent/memory.py完整代码

"""
短期对话记忆模块（Phase 9：token感知+滑动窗口裁剪）。

核心概念：
- 消息role严格对齐OpenAI Chat Completions API协议（user/assistant/tool）
- system不存入memory，每轮THINKING动态拼接
- 每条消息都有token成本，LLM上下文窗口是硬限制
- 裁剪单位是"消息块"而非单条消息，保护tool_calls↔tool_result的配对完整性

配对完整性约束（OpenAI协议硬要求）：
- assistant消息如果有tool_calls，后面必须紧跟对应tool_call_id的tool结果消息
- 不能只删assistant(tool_calls)而留下它的tool结果（tool消息会悬空）
- 也不能只删tool结果而留下assistant(tool_calls)（assistant会引用不存在的工具结果）
- 所以裁剪时要么删整个块[assistant(tc) + tool*N]，要么删独立消息（user/assistant无tc）
"""
from __future__ import annotations

from typing import Any, Optional

from agent.token_counter import count_message, count_messages


class ShortMemory:
    """短期对话记忆，消息格式直接对齐OpenAI API，带token预算管理"""

    def __init__(self, max_tokens: int = 6000):
        """
        :param max_tokens: memory中messages部分的token上限
            （不含system prompt和tools schema，那些由调用方在外部预算中扣除）
        """
        self.history: list[dict[str, Any]] = []
        self.max_tokens = max_tokens
        self.truncated_blocks = 0

    def add_user(self, content: str):
        self.history.append({"role": "user", "content": content})

    def add_assistant(self, content: Optional[str], tool_calls: Optional[list[dict]] = None):
        msg: dict[str, Any] = {"role": "assistant", "content": content}
        if tool_calls:
            msg["tool_calls"] = tool_calls
        self.history.append(msg)

    def add_tool_result(self, tool_call_id: str, content: str):
        self.history.append({
            "role": "tool",
            "tool_call_id": tool_call_id,
            "content": content,
        })

    def get_messages(self) -> list[dict]:
        return [m.copy() for m in self.history]

    def remove_last(self) -> Optional[dict]:
        if self.history:
            return self.history.pop()
        return None

    def token_count(self) -> int:
        """当前memory中所有消息的token总数"""
        return count_messages(self.history)

    def _pop_oldest_block(self) -> int:
        """
        删除最早的一个完整消息块，返回删除的token数。
        块定义见文件头注释。
        """
        if not self.history:
            return 0

        first = self.history[0]

        # 独立消息块：user 或 assistant（无tool_calls）
        if first["role"] == "user" or (first["role"] == "assistant" and not first.get("tool_calls")):
            self.history.pop(0)
            return count_message(first)

        # 工具调用块：assistant(tool_calls) + 其对应的tool结果消息
        if first["role"] == "assistant" and first.get("tool_calls"):
            tc_ids = {tc["id"] for tc in first["tool_calls"]}
            self.history.pop(0)
            removed = count_message(first)

            i = 0
            while i < len(self.history) and self.history[i]["role"] == "tool":
                msg_tc_id = self.history[i].get("tool_call_id")
                if msg_tc_id in tc_ids:
                    removed_msg = self.history.pop(i)
                    removed += count_message(removed_msg)
                    tc_ids.discard(msg_tc_id)
                    if not tc_ids:
                        break
                else:
                    i += 1
            return removed

        # 异常：tool消息在最前面（配对断裂），单独删除
        removed = self.history.pop(0)
        return count_message(removed)

    def truncate(self, reserve_tokens: int = 0) -> bool:
        """
        裁剪最早的消息块，直到memory的token数在max_tokens - reserve_tokens以内。

        :param reserve_tokens: 需要为外部内容预留的token数（system prompt + tools schema等）
        :return: 是否发生了裁剪
        """
        budget = self.max_tokens - reserve_tokens
        if budget < 0:
            budget = 0

        truncated = False
        while self.history and count_messages(self.history) > budget:
            removed = self._pop_oldest_block()
            if removed == 0:
                break
            truncated = True
            self.truncated_blocks += 1

        return truncated

几个设计决策说明：

• max_tokens是messages部分的预算，不是总预算。system prompt和tools schema的开销由调用方在truncate时通过reserve_tokens传入。memory不感知外部开销，只管理自己管的messages；
• add/add_assistant/add_tool_result不触发自动裁剪。因为添加消息时还不知道这一轮的外部开销（system prompt有没有变化？tools schema这一轮带不带？），裁剪时机在THINKING发起LLM调用前，所有外部信息确定后统一检查；
• truncated_blocks计数器：记录累计裁剪了多少块，用于日志展示，让用户知道"历史对话已被裁剪"。

四、改造3：main.py集成预算检查

4.1 预算分配模型

在run_agent入口定义预算常量：

MODEL_CONTEXT_WINDOW = 8192    # 模型上下文窗口大小，按实际模型调整
OUTPUT_RESERVE = 1024          # 为模型输出预留的token数
memory_budget = MODEL_CONTEXT_WINDOW - OUTPUT_RESERVE
memory = ShortMemory(max_tokens=memory_budget)

总预算8192是保守默认值（覆盖GPT-3.5/DeepSeek-chat等8K模型）。如果你的模型是32K或128K窗口，改这个常量即可。

4.2 THINKING前的预算检查

每轮THINKING调用LLM前，计算外部固定开销，检查总token是否超限：

if state == AgentTaskState.THINKING:
    # ---- Token预算检查：裁剪最早的消息块以确保不超上下文窗口 ----
    # 外部固定开销 = system prompt + tools schema + 请求前缀 + 输出预留
    external_reserve = (
        count_text(SYSTEM_PROMPT) + 4   # system message固定开销4 tokens
        + count_tools_schema(tools_schema)
        + 2                              # 请求前缀开销
        + OUTPUT_RESERVE                 # 输出预留
    )
    total_estimate = memory.token_count() + external_reserve

    if total_estimate > MODEL_CONTEXT_WINDOW:
        print(f"【上下文管理】当前约{total_estimate} tokens，超出窗口{MODEL_CONTEXT_WINDOW}，裁剪最早对话...")
        memory.truncate(reserve_tokens=external_reserve)
        print(f"【上下文管理】已裁剪{memory.truncated_blocks}个消息块，剩余约{memory.token_count()} tokens")

    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    messages.extend(memory.get_messages())
    # ...后续流式调用逻辑不变...

为什么用reserve_tokens参数传给memory.truncate？因为memory的max_tokens是messages部分的预算，但外部开销（system+tools+输出预留）也要占总窗口，所以messages实际可用预算 = max_tokens - external_reserve。memory.truncate内部会计算这个有效预算，然后循环删块直到messages token数在预算内。

4.3 为什么不在add时自动裁剪？

你可能会问：为什么不在add_user/add_assistant/add_tool_result里自动检查裁剪？原因有二：

1. 添加消息时外部开销还不确定 。tools_schema是在run_agent入口构造的（tools_schema = [t.to_openai_tool_schema() for t in tool_list]），add方法不感知它；
2. 裁剪时机有语义。裁剪是"要调LLM了，先确保请求合法"，不是"加了消息就可能超"。加消息和发起LLM调用之间可能有多次添加（比如TOOL_EXECUTING连续加多条tool结果），不应该中间裁剪。

这和Go里GC的触发时机类似------不是每次分配内存都GC，而是到达安全点（分配后、函数返回前）才检查。

五、验证：配对完整性测试

滑动窗口最关键的不是"能删消息"，而是"删完之后消息序列仍然是OpenAI协议合法的"。写一个验证脚本：

from agent.memory import ShortMemory
from agent.token_counter import count_messages

def test_pairing_integrity():
    """验证裁剪后tool_calls↔tool_result配对完整性"""
    m = ShortMemory(max_tokens=50)
    m.add_user('第一轮对话')
    m.add_assistant('好的', None)
    m.add_user('第二轮：帮我读文件')
    tc = [{'id': 'call_a', 'type': 'function', 'function': {'name': 'read_file', 'arguments': '{"path":"test.py"}'}}]
    m.add_assistant('我来读文件', tc)
    m.add_tool_result('call_a', 'file content here')
    m.add_user('第三轮')
    tc2 = [{'id': 'call_b', 'type': 'function', 'function': {'name': 'calculator', 'arguments': '{"expr":"2+3"}'}}]
    m.add_assistant('计算', tc2)
    m.add_tool_result('call_b', '5')

    print(f'裁剪前消息数: {len(m.history)}')
    print(f'裁剪前roles: {[msg["role"] for msg in m.history]}')

    m.truncate(reserve_tokens=0)

    print(f'裁剪后消息数: {len(m.history)}')
    print(f'裁剪后roles: {[msg["role"] for msg in m.history]}')

    # 验证：没有悬空的tool消息
    for i, msg in enumerate(m.history):
        if msg['role'] == 'tool':
            tc_id = msg['tool_call_id']
            found = False
            for j in range(i-1, -1, -1):
                if m.history[j]['role'] == 'assistant' and m.history[j].get('tool_calls'):
                    if tc_id in {tc['id'] for tc in m.history[j]['tool_calls']}:
                        found = True
                        break
                elif m.history[j]['role'] in ('user', 'assistant') and not m.history[j].get('tool_calls'):
                    break
            assert found, f'配对断裂！tool消息在index={i}没有对应的assistant(tool_calls)'

    print('配对完整性：OK')

test_pairing_integrity()

核心验证逻辑：遍历裁剪后的消息序列，对每条tool消息，向前找最近的带tool_calls的assistant消息，确认它的tool_call_id在那个assistant的tool_calls数组里。如果某条tool消息找不到对应的assistant，就是配对断裂，assert失败。

六、运行效果

配置8K窗口跑一个多轮对话：

=== 第1轮 THINKING ===
用户需要分析base_tool.py文件，我先用bash统计行数，再用read_file读取...
【工具调用意图】bash
【工具调用】bash({"command": "wc -l tools/base_tool.py"})
【工具返回】[退出码] 0\n[执行目录] /.../native-agent-demo\n[stdout]\n56 tools/base_tool.py
【工具调用意图】read_file
【工具调用】read_file({"path": "tools/base_tool.py", "offset": 0, "limit": 56})
【工具返回】[文件: tools/base_tool.py] 共56行...
=== 第2轮 THINKING ===
现在我需要读取execute方法和to_openai_tool_schema方法的内容...
【工具调用意图】calculator
【工具调用】calculator({"expr": "(18+20)*4/2"})
【工具返回】计算结果: (18+20)*4/2 = 76.0
=== 第3轮 THINKING ===
根据工具返回结果...
【校验】正在Judge回答完整性...
【校验通过】回答完整

当对话继续多轮后（比如持续对话10+轮，读了多个大文件），触发裁剪：

【上下文管理】当前约8437 tokens，超出窗口8192，裁剪最早对话...
【上下文管理】已裁剪2个消息块，剩余约6123 tokens
=== 第8轮 THINKING ===
...

裁剪后Agent继续正常工作------最早的对话块被移除，近期上下文保留。在8K窗口下，通常能保留5-8轮完整对话（具体取决于每轮tool返回内容的长度）。

七、小结

本阶段做的事情：

1. 新增token_counter.py：用tiktoken精确计算文本/消息/tool_call/tools schema/完整请求的token数，遵循OpenAI官方计数规则；
2. 改造memory.py：ShortMemory增加token感知（token_count/truncate/truncated_blocks），核心算法是_pop_oldest_block按消息块删除，严格保护tool_calls↔tool_result配对完整性；
3. 改造main.py：每轮THINKING前计算总token预算（system+tools+messages+输出预留），超窗口自动触发裁剪并打印日志；
4. 更新requirements.txt：新增tiktoken依赖。

关键设计原则回顾：

• 展示层和数据层分离：日志打印说"裁剪了N个块"是给人看的，memory里history的完整性是给LLM用的，两者独立；
• 确定性优先：滑动窗口是纯算法（不调LLM），可靠可预测，是所有更高级策略（摘要压缩/RAG）的基础；
• 外部预算reserve_tokens模式：memory只管理messages的token，system prompt/tools schema/输出预留由调用方计算后通过reserve_tokens传入，保持memory职责单一；
• 裁剪时机在"发请求前"：不在add时自动裁剪，在THINKING状态即将调LLM、所有外部信息确定后统一检查。