tiny-agent 0.1.0

//! Middleware:在 agent 生命周期的关键节点**拦截并改写**行为。
//!
//! 与 `Trajectory`(只读观测)互补——中间件可以改请求、改消息、改工具结果,
//! 甚至在工具执行前短路。钩子按 onion 模型组合:`before_*` 正序执行,
//! `after_*` 逆序执行(先包裹的最后收尾)。
//!
//! 常见用途:
//! - `before_model`:注入/改写系统提示、裁剪历史、设置 max_tokens、做 prompt 缓存标记;
//! - `after_model`:给 assistant 消息打标、做内容审计;
//! - `before_tool`:危险工具护栏(拒绝执行)、命中缓存直接返回;
//! - `after_tool`:敏感信息脱敏、超长输出截断。

use crate::{
    providers::LLMRequest,
    shared::{ContentBlock, Message, ToolCall},
    state::AgentState,
};
use async_trait::async_trait;
use std::sync::Arc;

/// 中间件可见的运行上下文(只读快照)。
pub struct MiddlewareCtx<'a> {
    pub session_id: &'a str,
    pub state: &'a AgentState,
}

/// `before_tool` 的返回:放行,或用给定结果直接短路(工具不会真正执行)。
pub enum ToolDecision {
    /// 继续正常执行该工具。
    Proceed,
    /// 跳过执行,直接把这个结果作为工具输出
    ShortCircuit {
        output: Vec<ContentBlock>,
        is_error: bool,
    },
}

/// 中间件钩子。所有方法都有默认空实现,按需覆盖。
///
/// 返回 `Err` 表示中间件自身出错,会被运行时当作一次失败(`kind = "middleware"`)处理。
#[async_trait]
pub trait Middleware: Send + Sync {
    fn name(&self) -> &str;

    /// 模型调用前:可改写即将发给 provider 的请求。
    async fn before_model(
        &self,
        _ctx: &MiddlewareCtx<'_>,
        _req: &mut LLMRequest,
    ) -> Result<(), String> {
        Ok(())
    }

    /// 模型产出 assistant 消息后、落库前:可改写该消息。
    async fn after_model(
        &self,
        _ctx: &MiddlewareCtx<'_>,
        _msg: &mut Message,
    ) -> Result<(), String> {
        Ok(())
    }

    /// 工具执行前:放行或短路。
    async fn before_tool(
        &self,
        _ctx: &MiddlewareCtx<'_>,
        _call: &ToolCall,
    ) -> Result<ToolDecision, String> {
        Ok(ToolDecision::Proceed)
    }

    /// 工具执行后、结果落库前:可改写该 `ToolResult` 块。
    async fn after_tool(
        &self,
        _ctx: &MiddlewareCtx<'_>,
        _call: &ToolCall,
        _result: &mut ContentBlock,
    ) -> Result<(), String> {
        Ok(())
    }
}

/// 一组中间件,按 onion 顺序驱动。运行时持有它并在各钩子点调用。
#[derive(Clone, Default)]
pub struct MiddlewareChain {
    middlewares: Vec<Arc<dyn Middleware>>,
}

impl MiddlewareChain {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn push(&mut self, middleware: Arc<dyn Middleware>) {
        self.middlewares.push(middleware);
    }

    pub fn with(mut self, middleware: Arc<dyn Middleware>) -> Self {
        self.middlewares.push(middleware);
        self
    }

    pub fn is_empty(&self) -> bool {
        self.middlewares.is_empty()
    }

    /// 正序执行所有 `before_model`。出错时用中间件 `name()` 标注是谁挂的。
    pub async fn before_model(
        &self,
        ctx: &MiddlewareCtx<'_>,
        req: &mut LLMRequest,
    ) -> Result<(), String> {
        for m in &self.middlewares {
            m.before_model(ctx, req)
                .await
                .map_err(|e| format!("[{}] {e}", m.name()))?;
        }
        Ok(())
    }

    /// 逆序执行所有 `after_model`。出错时用中间件 `name()` 标注是谁挂的。
    pub async fn after_model(
        &self,
        ctx: &MiddlewareCtx<'_>,
        msg: &mut Message,
    ) -> Result<(), String> {
        for m in self.middlewares.iter().rev() {
            m.after_model(ctx, msg)
                .await
                .map_err(|e| format!("[{}] {e}", m.name()))?;
        }
        Ok(())
    }

    /// 正序询问 `before_tool`,第一个返回 `ShortCircuit` 的中间件胜出。
    /// 出错时用中间件 `name()` 标注是谁挂的。
    pub async fn before_tool(
        &self,
        ctx: &MiddlewareCtx<'_>,
        call: &ToolCall,
    ) -> Result<ToolDecision, String> {
        for m in &self.middlewares {
            match m
                .before_tool(ctx, call)
                .await
                .map_err(|e| format!("[{}] {e}", m.name()))?
            {
                ToolDecision::Proceed => {}
                short_circuit => return Ok(short_circuit),
            }
        }
        Ok(ToolDecision::Proceed)
    }

    /// 逆序执行所有 `after_tool`。出错时用中间件 `name()` 标注是谁挂的。
    pub async fn after_tool(
        &self,
        ctx: &MiddlewareCtx<'_>,
        call: &ToolCall,
        result: &mut ContentBlock,
    ) -> Result<(), String> {
        for m in self.middlewares.iter().rev() {
            m.after_tool(ctx, call, result)
                .await
                .map_err(|e| format!("[{}] {e}", m.name()))?;
        }
        Ok(())
    }
}