tiny_agent/middleware/
mod.rs

1//! Middleware:在 agent 生命周期的关键节点**拦截并改写**行为。
2//!
3//! 中间件可以改请求、改消息、改工具结果,
4//! 甚至在工具执行前短路。钩子按 onion 模型组合:`before_*` 正序执行,
5//! `after_*` 逆序执行(先包裹的最后收尾)。
6//!
7//! 常见用途:
8//! - `before_model`:注入/改写系统提示、裁剪历史、设置 max_tokens、做 prompt 缓存标记;
9//! - `after_model`:给 assistant 消息打标、做内容审计;
10//! - `before_tool`:危险工具护栏(拒绝执行)、命中缓存直接返回;
11//! - `after_tool`:敏感信息脱敏、超长输出截断。
12
13use crate::{
14    providers::LLMRequest,
15    shared::{ContentBlock, Message, ToolCall},
16    state::AgentState,
17};
18use async_trait::async_trait;
19use std::sync::Arc;
20
21/// 中间件可见的运行上下文(只读快照)。
22pub struct MiddlewareCtx<'a> {
23    pub session_id: &'a str,
24    pub state: &'a AgentState,
25}
26
27/// `before_tool` 的返回:放行,或用给定结果直接短路(工具不会真正执行)。
28pub enum ToolDecision {
29    /// 继续正常执行该工具。
30    Proceed,
31    /// 跳过执行,直接把这个结果作为工具输出
32    ShortCircuit {
33        output: Vec<ContentBlock>,
34        is_error: bool,
35    },
36}
37
38/// 中间件钩子。所有方法都有默认空实现,按需覆盖。
39///
40/// 返回 `Err` 表示中间件自身出错,会被运行时当作一次失败(`kind = "middleware"`)处理。
41#[async_trait]
42pub trait Middleware: Send + Sync {
43    fn name(&self) -> &str;
44
45    /// 模型调用前:可改写即将发给 provider 的请求。
46    async fn before_model(
47        &self,
48        _ctx: &MiddlewareCtx<'_>,
49        _req: &mut LLMRequest,
50    ) -> Result<(), String> {
51        Ok(())
52    }
53
54    /// 模型产出 assistant 消息后、落库前:可改写该消息。
55    async fn after_model(
56        &self,
57        _ctx: &MiddlewareCtx<'_>,
58        _msg: &mut Message,
59    ) -> Result<(), String> {
60        Ok(())
61    }
62
63    /// 工具执行前:放行或短路。
64    async fn before_tool(
65        &self,
66        _ctx: &MiddlewareCtx<'_>,
67        _call: &ToolCall,
68    ) -> Result<ToolDecision, String> {
69        Ok(ToolDecision::Proceed)
70    }
71
72    /// 工具执行后、结果落库前:可改写该 `ToolResult` 块。
73    async fn after_tool(
74        &self,
75        _ctx: &MiddlewareCtx<'_>,
76        _call: &ToolCall,
77        _result: &mut ContentBlock,
78    ) -> Result<(), String> {
79        Ok(())
80    }
81}
82
83/// 一组中间件,按 onion 顺序驱动。运行时持有它并在各钩子点调用。
84#[derive(Clone, Default)]
85pub struct MiddlewareChain {
86    middlewares: Vec<Arc<dyn Middleware>>,
87}
88
89impl MiddlewareChain {
90    pub fn new() -> Self {
91        Self::default()
92    }
93
94    pub fn push(&mut self, middleware: Arc<dyn Middleware>) {
95        self.middlewares.push(middleware);
96    }
97
98    pub fn with(mut self, middleware: Arc<dyn Middleware>) -> Self {
99        self.middlewares.push(middleware);
100        self
101    }
102
103    pub fn is_empty(&self) -> bool {
104        self.middlewares.is_empty()
105    }
106
107    /// 正序执行所有 `before_model`。出错时用中间件 `name()` 标注是谁挂的。
108    pub async fn before_model(
109        &self,
110        ctx: &MiddlewareCtx<'_>,
111        req: &mut LLMRequest,
112    ) -> Result<(), String> {
113        for m in &self.middlewares {
114            m.before_model(ctx, req)
115                .await
116                .map_err(|e| format!("[{}] {e}", m.name()))?;
117        }
118        Ok(())
119    }
120
121    /// 逆序执行所有 `after_model`。出错时用中间件 `name()` 标注是谁挂的。
122    pub async fn after_model(
123        &self,
124        ctx: &MiddlewareCtx<'_>,
125        msg: &mut Message,
126    ) -> Result<(), String> {
127        for m in self.middlewares.iter().rev() {
128            m.after_model(ctx, msg)
129                .await
130                .map_err(|e| format!("[{}] {e}", m.name()))?;
131        }
132        Ok(())
133    }
134
135    /// 正序询问 `before_tool`,第一个返回 `ShortCircuit` 的中间件胜出。
136    /// 出错时用中间件 `name()` 标注是谁挂的。
137    pub async fn before_tool(
138        &self,
139        ctx: &MiddlewareCtx<'_>,
140        call: &ToolCall,
141    ) -> Result<ToolDecision, String> {
142        for m in &self.middlewares {
143            match m
144                .before_tool(ctx, call)
145                .await
146                .map_err(|e| format!("[{}] {e}", m.name()))?
147            {
148                ToolDecision::Proceed => {}
149                short_circuit => return Ok(short_circuit),
150            }
151        }
152        Ok(ToolDecision::Proceed)
153    }
154
155    /// 逆序执行所有 `after_tool`。出错时用中间件 `name()` 标注是谁挂的。
156    pub async fn after_tool(
157        &self,
158        ctx: &MiddlewareCtx<'_>,
159        call: &ToolCall,
160        result: &mut ContentBlock,
161    ) -> Result<(), String> {
162        for m in self.middlewares.iter().rev() {
163            m.after_tool(ctx, call, result)
164                .await
165                .map_err(|e| format!("[{}] {e}", m.name()))?;
166        }
167        Ok(())
168    }
169}
tiny_agent/middleware/mod.rs

tiny_agent/middleware/
mod.rs
