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                "本轮工具调用已达上限。请根据已有结果总结并向用户报告。不要再调用工具。".to_string(),
54            );
55        }
56
57        Ok(())
58    }
59}
60
61#[cfg(test)]
62mod tests {
63    use super::*;
64    use crate::types::SessionId;
65
66    fn make_ctx(turn_tool_calls: usize, pending_calls: usize) -> PostLlmCtx {
67        let tool_calls = (0..pending_calls)
68            .map(|i| (format!("call_{}", i), "test_tool".to_string(), "{}".to_string()))
69            .collect();
70        PostLlmCtx {
71            session_id: SessionId::new(1),
72            full_text: String::new(),
73            is_tool_call: pending_calls > 0,
74            tool_calls,
75            available_tools: vec!["test_tool".to_string()],
76            turn_count: 1,
77            total_tool_calls: 0,
78            nudge_count: 0,
79            turn_tool_calls,
80            skip_push: false,
81            follow_up_message: None,
82        }
83    }
84
85    #[tokio::test]
86    async fn allows_calls_under_limit() {
87        let mw = TurnToolLimitMiddleware::new(8);
88        let mut ctx = make_ctx(3, 2);
89        mw.on_post_llm(&mut ctx).await.unwrap();
90        assert_eq!(ctx.tool_calls.len(), 2);
91        assert!(ctx.follow_up_message.is_none());
92    }
93
94    #[tokio::test]
95    async fn blocks_calls_at_limit() {
96        let mw = TurnToolLimitMiddleware::new(8);
97        let mut ctx = make_ctx(8, 3);
98        mw.on_post_llm(&mut ctx).await.unwrap();
99        assert!(ctx.tool_calls.is_empty());
100        assert!(!ctx.is_tool_call);
101        assert!(ctx.follow_up_message.is_some());
102    }
103
104    #[tokio::test]
105    async fn blocks_calls_over_limit() {
106        let mw = TurnToolLimitMiddleware::new(8);
107        let mut ctx = make_ctx(12, 2);
108        mw.on_post_llm(&mut ctx).await.unwrap();
109        assert!(ctx.tool_calls.is_empty());
110        assert!(ctx.follow_up_message.is_some());
111    }
112
113    #[tokio::test]
114    async fn ignores_text_only_responses() {
115        let mw = TurnToolLimitMiddleware::new(8);
116        let mut ctx = make_ctx(10, 0); // no tool calls
117        ctx.is_tool_call = false;
118        ctx.full_text = "just text".to_string();
119        mw.on_post_llm(&mut ctx).await.unwrap();
120        assert!(ctx.follow_up_message.is_none());
121    }
122
123    #[tokio::test]
124    async fn custom_limit() {
125        let mw = TurnToolLimitMiddleware::new(3);
126        let mut ctx = make_ctx(3, 1);
127        mw.on_post_llm(&mut ctx).await.unwrap();
128        assert!(ctx.tool_calls.is_empty());
129        assert!(ctx.follow_up_message.is_some());
130    }
131}
agent_base/engine/safety.rs

agent_base/engine/
safety.rs
