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::{FinishReason, 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            finish_reason: FinishReason::Stop,
90        }
91    }
92
93    #[tokio::test]
94    async fn allows_calls_under_limit() {
95        let mw = TurnToolLimitMiddleware::new(8);
96        let mut ctx = make_ctx(3, 2);
97        mw.on_post_llm(&mut ctx).await.unwrap();
98        assert_eq!(ctx.tool_calls.len(), 2);
99        assert!(ctx.follow_up_message.is_none());
100    }
101
102    #[tokio::test]
103    async fn blocks_calls_at_limit() {
104        let mw = TurnToolLimitMiddleware::new(8);
105        let mut ctx = make_ctx(8, 3);
106        mw.on_post_llm(&mut ctx).await.unwrap();
107        assert!(ctx.tool_calls.is_empty());
108        assert!(!ctx.is_tool_call);
109        assert!(ctx.follow_up_message.is_some());
110    }
111
112    #[tokio::test]
113    async fn blocks_calls_over_limit() {
114        let mw = TurnToolLimitMiddleware::new(8);
115        let mut ctx = make_ctx(12, 2);
116        mw.on_post_llm(&mut ctx).await.unwrap();
117        assert!(ctx.tool_calls.is_empty());
118        assert!(ctx.follow_up_message.is_some());
119    }
120
121    #[tokio::test]
122    async fn ignores_text_only_responses() {
123        let mw = TurnToolLimitMiddleware::new(8);
124        let mut ctx = make_ctx(10, 0); // no tool calls
125        ctx.is_tool_call = false;
126        ctx.full_text = "just text".to_string();
127        mw.on_post_llm(&mut ctx).await.unwrap();
128        assert!(ctx.follow_up_message.is_none());
129    }
130
131    #[tokio::test]
132    async fn custom_limit() {
133        let mw = TurnToolLimitMiddleware::new(3);
134        let mut ctx = make_ctx(3, 1);
135        mw.on_post_llm(&mut ctx).await.unwrap();
136        assert!(ctx.tool_calls.is_empty());
137        assert!(ctx.follow_up_message.is_some());
138    }
139}
agent_base/engine/safety.rs

agent_base/engine/
safety.rs
