Middleware 管道系统分析

class AgentRunContext(SerializationMixin):
    INJECTABLE: ClassVar[set[str]] = {"agent", "thread", "result"}
    
    def __init__(
        self,
        agent: "AgentProtocol",
        messages: list[ChatMessage],
        thread: "AgentThread | None" = None,
        is_streaming: bool = False,
        metadata: dict[str, Any] | None = None,
        result: AgentResponse | AsyncIterable[AgentResponseUpdate] | None = None,
        terminate: bool = False,
        kwargs: dict[str, Any] | None = None,
    ) -> None:

class FunctionInvocationContext(SerializationMixin):
    INJECTABLE: ClassVar[set[str]] = {"function", "arguments", "result"}
    
    def __init__(
        self,
        function: "FunctionTool[Any, Any]",
        arguments: "BaseModel",  # Pydantic 验证后的参数
        metadata: dict[str, Any] | None = None,
        result: Any = None,
        terminate: bool = False,
        kwargs: dict[str, Any] | None = None,
    ) -> None:

class ChatContext(SerializationMixin):
    INJECTABLE: ClassVar[set[str]] = {"chat_client", "result"}
    
    def __init__(
        self,
        chat_client: "ChatClientProtocol",
        messages: "MutableSequence[ChatMessage]",  # 可修改的消息列表
        options: Mapping[str, Any] | None,  # LLM 参数
        is_streaming: bool = False,
        metadata: dict[str, Any] | None = None,
        result: "ChatResponse | AsyncIterable[ChatResponseUpdate] | None" = None,
        terminate: bool = False,
        kwargs: dict[str, Any] | None = None,
    ) -> None:

class AgentMiddleware(ABC):
    @abstractmethod
    async def process(
        self,
        context: AgentRunContext,
        next: Callable[[AgentRunContext], Awaitable[None]],
    ) -> None:
        """
        关键设计：不返回值，所有数据通过 context 流转
        - 设置 context.result 可覆盖执行结果
        - 调用 next(context) 继续执行链
        """

async def process(self, context: TContext, next: NextFn) -> None
# 其中 NextFn = Callable[[TContext], Awaitable[None]]

# 类型别名定义简洁的函数签名
AgentMiddlewareCallable = Callable[
    [AgentRunContext, Callable[[AgentRunContext], Awaitable[None]]], 
    Awaitable[None]
]

class MiddlewareWrapper(Generic[TContext]):
    """将纯函数包装为类协议"""
    def __init__(self, func: Callable[..., Awaitable[None]]) -> None:
        self.func = func
    
    async def process(self, context: T, next: NextFn) -> None:
        await self.func(context, next)

# 类方式（有状态）
class LoggingMiddleware(AgentMiddleware):
    def __init__(self, logger):
        self.logger = logger
    
    async def process(self, context, next):
        self.logger.info(f"Agent: {context.agent.name}")
        await next(context)

# 函数方式（无状态/简洁）
@agent_middleware
async def logging_middleware(context: AgentRunContext, next):
    print(f"Before: {context.agent.name}")
    await next(context)
    print(f"After: {context.result}")

def agent_middleware(func: AgentMiddlewareCallable) -> AgentMiddlewareCallable:
    # 通过属性标记中间件类型，用于自动分类
    func._middleware_type: MiddlewareType = MiddlewareType.AGENT
    return func

def function_middleware(func: FunctionMiddlewareCallable) -> FunctionMiddlewareCallable:
    func._middleware_type: MiddlewareType = MiddlewareType.FUNCTION
    return func

def chat_middleware(func: ChatMiddlewareCallable) -> ChatMiddlewareCallable:
    func._middleware_type: MiddlewareType = MiddlewareType.CHAT
    return func

def _create_handler_chain(
    self,
    final_handler: Callable[[Any], Awaitable[Any]],
    result_container: dict[str, Any],
    result_key: str = "result",
) -> Callable[[Any], Awaitable[None]]:
    """
    递归构建洋葱式调用链
    
    核心思想：
    - 从最后一个中间件向前构建
    - 每个中间件包装下一个处理函数
    - 最终包装实际业务逻辑
    """
    def create_next_handler(index: int) -> Callable[[Any], Awaitable[None]]:
        if index >= len(self._middleware):
            # 链尾：实际执行逻辑
            async def final_wrapper(c: Any) -> None:
                result = await final_handler(c)
                result_container[result_key] = result
                c.result = result  # 结果回写 Context
            return final_wrapper
        
        # 递归构建：当前中间件包装下一个 handler
        middleware = self._middleware[index]
        next_handler = create_next_handler(index + 1)  # 递归！
        
        async def current_handler(c: Any) -> None:
            await middleware.process(c, next_handler)
        
        return current_handler
    
    return create_next_handler(0)  # 从第一个中间件开始

def _create_streaming_handler_chain(
    self,
    final_handler: Callable[[Any], Any],  # 返回 AsyncIterable
    result_container: dict[str, Any],
    result_key: str = "result_stream",
) -> Callable[[Any], Awaitable[None]]:
    
    def create_next_handler(index: int) -> Callable[[Any], Awaitable[None]]:
        if index >= len(self._middleware):
            async def final_wrapper(c: Any) -> None:
                if c.terminate:
                    return  # 提前终止
                
                # 处理非协程情况（如生成器函数）
                try:
                    result = await final_handler(c)
                except TypeError:
                    result = final_handler(c)  # 同步生成器
                
                result_container[result_key] = result
                c.result = result
            return final_wrapper
        
        middleware = self._middleware[index]
        next_handler = create_next_handler(index + 1)
        
        async def current_handler(c: Any) -> None:
            await middleware.process(c, next_handler)
            if c.terminate:  # 终止检查
                return
        
        return current_handler

# Agent 非流式执行（第 776-826 行）
async def execute(
    self,
    agent: "AgentProtocol",
    messages: list[ChatMessage],
    context: AgentRunContext,
    final_handler: Callable[[AgentRunContext], Awaitable[AgentResponse]],
) -> AgentResponse | None:
    # 1. 更新 Context
    context.agent = agent
    context.messages = messages
    context.is_streaming = False
    
    # 2. 短路：无中间件时直接执行
    if not self._middleware:
        return await final_handler(context)
    
    # 3. 构建并执行链
    result_container = {"result": None}
    first_handler = self._create_handler_chain(agent_final_handler, result_container, "result")
    await first_handler(context)
    
    # 4. 结果处理（支持 override）
    if context.result is not None:
        return context.result  # 中间件覆盖的结果
    return result_container.get("result") or AgentResponse()

def use_agent_middleware(agent_class: type[TAgent]) -> type[TAgent]:
    # 1. 保存原始方法
    original_run = agent_class.run
    original_run_stream = agent_class.run_stream
    
    def _build_middleware_pipelines(
        agent_level_middlewares: Sequence[Middleware] | None,
        run_level_middlewares: Sequence[Middleware] | None = None,
    ) -> tuple[AgentMiddlewarePipeline, FunctionMiddlewarePipeline, list[ChatMiddleware]]:
        """
        双层中间件架构：
        - Agent 级：实例化时注册，全局生效
        - Run 级：每次调用时传入，仅本次生效
        """
        middleware = categorize_middleware(
            *(agent_level_middlewares or ()), 
            *(run_level_middlewares or ())
        )
        return (
            AgentMiddlewarePipeline(middleware["agent"]),
            FunctionMiddlewarePipeline(middleware["function"]),
            middleware["chat"],
        )
    
    async def middleware_enabled_run(self, messages, *, thread=None, middleware=None, **kwargs):
        """被装饰后的 run 方法"""
        # 获取实例级中间件
        agent_middleware = getattr(self, "middleware", None)
        
        # 构建管道（合并实例级 + 调用级）
        agent_pipeline, function_pipeline, chat_middlewares = _build_middleware_pipelines(
            agent_middleware, middleware
        )
        
        # 传递 Function 中间件到下层（通过 kwargs 隐式传递）
        if function_pipeline.has_middlewares:
            kwargs["_function_middleware_pipeline"] = function_pipeline
        
        # 传递 Chat 中间件到下层
        if chat_middlewares:
            kwargs["middleware"] = chat_middlewares
        
        # 标准化消息格式
        normalized_messages = normalize_messages(messages)
        
        # 执行中间件管道或短路
        if agent_pipeline.has_middlewares:
            context = AgentRunContext(
                agent=self,
                messages=normalized_messages,
                thread=thread,
                is_streaming=False,
                kwargs=kwargs,
            )
            
            async def _execute_handler(ctx):
                return await original_run(self, ctx.messages, thread=thread, **ctx.kwargs)
            
            return await agent_pipeline.execute(self, normalized_messages, context, _execute_handler)
        
        # 无中间件：直接调用原始方法
        return await original_run(self, normalized_messages, thread=thread, **kwargs)

def categorize_middleware(
    *middleware_sources: Middleware | None,
) -> MiddlewareDict:
    """
    智能分类逻辑：
    1. 类实例：通过 isinstance 判断
    2. 函数：通过 _middleware_type 标记或类型注解推断
    """
    result: MiddlewareDict = {"agent": [], "function": [], "chat": []}
    
    for middleware in all_middleware:
        if isinstance(middleware, AgentMiddleware):
            result["agent"].append(middleware)
        elif isinstance(middleware, FunctionMiddleware):
            result["function"].append(middleware)
        elif isinstance(middleware, ChatMiddleware):
            result["chat"].append(middleware)
        elif callable(middleware):
            # 函数类型：标记 > 注解 > 异常
            middleware_type = _determine_middleware_type(middleware)
            # ... 按类型分发

def _determine_middleware_type(middleware: Any) -> MiddlewareType:
    """
    双重验证机制：
    1. 装饰器标记（显式）
    2. 参数类型注解（隐式）
    """
    decorator_type = getattr(middleware, "_middleware_type", None)
    
    # 通过签名推断
    sig = inspect.signature(middleware)
    params = list(sig.parameters.values())
    first_param = params[0]
    
    if first_param.annotation.__name__ == "AgentRunContext":
        param_type = MiddlewareType.AGENT
    # ... 其他类型判断
    
    # 双重验证：装饰器和注解必须一致
    if decorator_type and param_type:
        if decorator_type != param_type:
            raise MiddlewareException("Type mismatch...")
        return decorator_type
    
    return decorator_type or param_type  # 优先级：显式 > 隐式

class ValidationMiddleware(FunctionMiddleware):
    async def process(self, context: FunctionInvocationContext, next):
        # 参数验证失败，直接返回错误，不执行实际函数
        if not self.validate(context.arguments):
            context.result = {"error": "Validation failed"}
            context.terminate = True  # 终止后续执行
            return  # 不调用 next()
        
        await next(context)

# 优雅的递归构建洋葱模型
next_handler = create_next_handler(index + 1)  # 递归构建内层
async def current_handler(c):
    await middleware.process(c, next_handler)  # 包装当前层

def create_next_handler(index: int) -> Callable[[Any], Awaitable[None]]:
    if index >= len(self._middleware):
        # 基准情况：最终处理器
        async def final_wrapper(c: Any) -> None:
            result = await final_handler(c)
            result_container[result_key] = result
            c.result = result
        return final_wrapper
    
    # 递归步骤：当前中间件包装下一个
    middleware = self._middleware[index]
    next_handler = create_next_handler(index + 1)  # 递归调用
    
    async def current_handler(c: Any) -> None:
        await middleware.process(c, next_handler)
    
    return current_handler

return create_next_handler(0)  # 从链头开始

def _build_middleware_pipelines(
    agent_level_middlewares: Sequence[Middleware] | None,
    run_level_middlewares: Sequence[Middleware] | None = None,
) -> tuple[...]:
    """
    设计巧思：使用 None 作为哨兵值，避免空列表的歧义
    *(agent_level_middlewares or ())  # 解包语法确保参数展开正确
    """
    middleware = categorize_middleware(
        *(agent_level_middlewares or ()),
        *(run_level_middlewares or ())
    )
    # ...