agent_base/engine/
safety.rs

1use async_trait::async_trait;
2
3use crate::engine::middleware::{Middleware, PostLlmCtx};
4use crate::types::{AgentResult, SafetyConfig};
5
6/// Middleware that enforces a hard limit on tool calls per turn.
7///
8/// Mounted on `on_post_llm` hook. When the turn's tool call count reaches
9/// `max_tool_calls_per_turn`, this middleware:
10/// 1. Clears `ctx.tool_calls` (discards pending tool calls)
11/// 2. Injects a `follow_up_message` forcing the LLM to summarize
12///
13/// This is a **hard constraint** — unlike prompt rules, the model cannot bypass it.
14pub struct TurnToolLimitMiddleware {
15    max_tool_calls_per_turn: usize,
16}
17
18impl TurnToolLimitMiddleware {
19    pub fn new(max_tool_calls_per_turn: usize) -> Self {
20        Self {
21            max_tool_calls_per_turn,
22        }
23    }
24
25    pub fn from_config(config: &SafetyConfig) -> Self {
26        Self::new(config.max_tool_calls_per_turn)
27    }
28}
29
30#[async_trait]
31impl Middleware for TurnToolLimitMiddleware {
32    async fn on_post_llm(&self, ctx: &mut PostLlmCtx) -> AgentResult<()> {
33        // Only intercept when the LLM is trying to call tools
34        if !ctx.is_tool_call || ctx.tool_calls.is_empty() {
35            return Ok(());
36        }
37
38        // Check if this turn has already hit the limit
39        if ctx.turn_tool_calls >= self.max_tool_calls_per_turn {
40            tracing::warn!(
41                turn_tool_calls = ctx.turn_tool_calls,
42                max = self.max_tool_calls_per_turn,
43                pending_calls = ctx.tool_calls.len(),
44                "TurnToolLimit: blocking tool calls — limit reached"
45            );
46
47            // Discard all pending tool calls
48            ctx.tool_calls.clear();
49            ctx.is_tool_call = false;
50
51            // Force LLM to summarize
52            ctx.follow_up_message = Some(
53                "本轮工具调用已达上限。请根据已有结果总结并向用户报告。不要再调用工具。"
54                    .to_string(),
55            );
56        }
57
58        Ok(())
59    }
60}
61
62#[cfg(test)]
63mod tests {
64    use super::*;
65    use crate::types::SessionId;
66
67    fn make_ctx(turn_tool_calls: usize, pending_calls: usize) -> PostLlmCtx {
68        let tool_calls = (0..pending_calls)
69            .map(|i| {
70                (
71                    format!("call_{}", i),
72                    "test_tool".to_string(),
73                    "{}".to_string(),
74                )
75            })
76            .collect();
77        PostLlmCtx {
78            session_id: SessionId::new(1),
79            full_text: String::new(),
80            is_tool_call: pending_calls > 0,
81            tool_calls,
82            available_tools: vec!["test_tool".to_string()],
83            turn_count: 1,
84            total_tool_calls: 0,
85            nudge_count: 0,
86            turn_tool_calls,
87            skip_push: false,
88            follow_up_message: None,
89        }
90    }
91
92    #[tokio::test]
93    async fn allows_calls_under_limit() {
94        let mw = TurnToolLimitMiddleware::new(8);
95        let mut ctx = make_ctx(3, 2);
96        mw.on_post_llm(&mut ctx).await.unwrap();
97        assert_eq!(ctx.tool_calls.len(), 2);
98        assert!(ctx.follow_up_message.is_none());
99    }
100
101    #[tokio::test]
102    async fn blocks_calls_at_limit() {
103        let mw = TurnToolLimitMiddleware::new(8);
104        let mut ctx = make_ctx(8, 3);
105        mw.on_post_llm(&mut ctx).await.unwrap();
106        assert!(ctx.tool_calls.is_empty());
107        assert!(!ctx.is_tool_call);
108        assert!(ctx.follow_up_message.is_some());
109    }
110
111    #[tokio::test]
112    async fn blocks_calls_over_limit() {
113        let mw = TurnToolLimitMiddleware::new(8);
114        let mut ctx = make_ctx(12, 2);
115        mw.on_post_llm(&mut ctx).await.unwrap();
116        assert!(ctx.tool_calls.is_empty());
117        assert!(ctx.follow_up_message.is_some());
118    }
119
120    #[tokio::test]
121    async fn ignores_text_only_responses() {
122        let mw = TurnToolLimitMiddleware::new(8);
123        let mut ctx = make_ctx(10, 0); // no tool calls
124        ctx.is_tool_call = false;
125        ctx.full_text = "just text".to_string();
126        mw.on_post_llm(&mut ctx).await.unwrap();
127        assert!(ctx.follow_up_message.is_none());
128    }
129
130    #[tokio::test]
131    async fn custom_limit() {
132        let mw = TurnToolLimitMiddleware::new(3);
133        let mut ctx = make_ctx(3, 1);
134        mw.on_post_llm(&mut ctx).await.unwrap();
135        assert!(ctx.tool_calls.is_empty());
136        assert!(ctx.follow_up_message.is_some());
137    }
138}
agent_base/engine/safety.rs

agent_base/engine/
safety.rs
