【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错误直接崩溃。

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

前置说明：

• 新增文件：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调用、是所有更高级策略的基础（摘要压缩也需要先确定哪些消息要被摘要）。按"先验证再固化"原则，先用最简方案跑通，等真实遇到"裁掉了重要上下文"的场景再加摘要。

二、改造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 消息块的定义

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

为什么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里自动检查裁剪？原因有二：

这和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返回内容的长度）。

七、小结

本阶段做的事情：

关键设计原则回顾：

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