中间件

每个AgentExtension都继承MiddlewareHook。当扩展需要包装跨越模型提供商或本地工具处理函数边界的操作时，就覆盖这些方法。中间件没有单独的注册表，也没有单独的构造函数参数。

两个边界

on_model_request(context, request, call_next)

包装每一次主要的模型提供商调用。由于模型输出是流式的，它本身是一个异步生成器。把原始请求或替换后的数据类传给 call_next(request)，并在提前退出时关闭返回的流。

on_tool_call(context, call, call_next)

包装一次本地工具调用真正使用的处理函数。它可以在 await call_next() 之前修改 call.arguments，转换原始的ToolResult，或者直接返回以短路该处理函数。异常会走普通的工具错误处理流程，并被记录为一条失败的 ToolMessage。

模型提供商托管的服务端工具不会经过 on_tool_call，因为提供商是在 on_model_request 内部执行它们的。

顺序

中间件把扩展已有的优先级顺序当作洋葱模型。给定 extensions=[outer, inner]时，顺序为 outer before -> inner before -> operation -> inner after -> outer after。优先级越低越靠外层；优先级相同时保持注册顺序。

from contextlib import aclosing
from dataclasses import replace
from time import monotonic

from zett_agent.extensions.base import AgentExtension
from zett_agent.messages import SystemMessage

class TraceExtension(AgentExtension):
    async def on_model_request(self, context, request, call_next):
        request = replace(
            request,
            messages=(*request.messages, SystemMessage(content="Trace this request.")),
        )
        async with aclosing(call_next(request)) as events:
            async for event in events:
                yield event

    async def on_tool_call(self, context, call, call_next):
        started = monotonic()
        try:
            return await call_next()
        finally:
            record_duration(call.id, monotonic() - started)

agent = await Agent.create(
    model,
    config=config,
    extensions=[TraceExtension()],
)

并发与取消

同一个扩展实例可能被不同会话同时进入。并行的本地工具也会在每个任务中执行一条独立的完整中间件链。请把调用状态保存在方法内的局部变量中，或对共享的可变状态加同步。current_tool_call_id() 在整个工具中间件和处理函数期间都保持在任务本地。

用 try/finally 包住 call_next 来做清理。模型中间件应在 contextlib.aclosing 下消费内层流，这样客户端取消或消费方提前退出才能及时传到模型提供商和所有内层中间件。