agent_base/engine/
tool_enforcement.rs

1use async_trait::async_trait;
2
3use crate::engine::middleware::{Middleware, PostLlmCtx};
4use crate::types::AgentResult;
5
6pub struct ToolEnforcementConfig {
7    pub max_nudges: usize,
8    pub nudge_message: String,
9    pub first_turn_only: bool,
10    pub min_tools_threshold: usize,
11}
12
13impl Default for ToolEnforcementConfig {
14    fn default() -> Self {
15        Self {
16            max_nudges: 3,
17            nudge_message: "CRITICAL: You have tools available but did not call any. \
18                             Call the appropriate tool NOW. \
19                             关键提示：你有可用的工具但没有调用。立即使用工具执行。"
20                .to_string(),
21            first_turn_only: true,
22            min_tools_threshold: 1,
23        }
24    }
25}
26
27pub struct ToolEnforcementMiddleware {
28    config: ToolEnforcementConfig,
29}
30
31impl ToolEnforcementMiddleware {
32    pub fn new(config: ToolEnforcementConfig) -> Self {
33        Self { config }
34    }
35}
36
37#[async_trait]
38impl Middleware for ToolEnforcementMiddleware {
39    async fn on_post_llm(&self, ctx: &mut PostLlmCtx) -> AgentResult<()> {
40        if ctx.available_tools.len() < self.config.min_tools_threshold {
41            return Ok(());
42        }
43        if self.config.first_turn_only && ctx.total_tool_calls > 0 {
44            return Ok(());
45        }
46        if ctx.is_tool_call {
47            return Ok(());
48        }
49        if ctx.full_text.is_empty() {
50            return Ok(());
51        }
52
53        if ctx.nudge_count >= self.config.max_nudges {
54            return Ok(());
55        }
56
57        // Increment nudge_count in the context; the caller will write it back to the session
58        ctx.nudge_count += 1;
59
60        tracing::info!(
61            nudge_count = ctx.nudge_count,
62            full_text_len = ctx.full_text.len(),
63            "ToolEnforcement: suppressing text response, injecting nudge"
64        );
65
66        ctx.skip_push = true;
67        ctx.follow_up_message = Some(self.config.nudge_message.clone());
68
69        Ok(())
70    }
71}
72
73#[cfg(test)]
74mod tests {
75    use super::*;
76    use crate::types::{FinishReason, SessionId};
77
78    fn ctx(available_tools: Vec<String>) -> PostLlmCtx {
79        PostLlmCtx {
80            session_id: SessionId::new(1),
81            full_text: "I will do it.".to_string(),
82            is_tool_call: false,
83            tool_calls: vec![],
84            available_tools,
85            turn_count: 1,
86            total_tool_calls: 0,
87            nudge_count: 0,
88            turn_tool_calls: 0,
89            skip_push: false,
90            follow_up_message: None,
91            finish_reason: FinishReason::Stop,
92        }
93    }
94
95    #[test]
96    fn config_defaults() {
97        let cfg = ToolEnforcementConfig::default();
98        assert_eq!(cfg.max_nudges, 3);
99        assert!(cfg.first_turn_only);
100        assert_eq!(cfg.min_tools_threshold, 1);
101        assert!(cfg.nudge_message.contains("CRITICAL"));
102    }
103
104    #[tokio::test]
105    async fn nudges_when_no_tool_call() {
106        let config = ToolEnforcementConfig::default();
107        let expected_nudge = config.nudge_message.clone();
108        let mw = ToolEnforcementMiddleware::new(config);
109
110        let mut c = ctx(vec!["read".to_string()]);
111        mw.on_post_llm(&mut c).await.unwrap();
112
113        assert!(c.skip_push);
114        assert_eq!(c.nudge_count, 1);
115        assert_eq!(c.follow_up_message, Some(expected_nudge));
116    }
117
118    #[tokio::test]
119    async fn skips_when_below_tools_threshold() {
120        let mw = ToolEnforcementMiddleware::new(ToolEnforcementConfig::default());
121        let mut c = ctx(vec![]);
122        mw.on_post_llm(&mut c).await.unwrap();
123
124        assert!(!c.skip_push);
125        assert_eq!(c.nudge_count, 0);
126        assert!(c.follow_up_message.is_none());
127    }
128
129    #[tokio::test]
130    async fn skips_when_is_tool_call() {
131        let mw = ToolEnforcementMiddleware::new(ToolEnforcementConfig::default());
132        let mut c = ctx(vec!["read".to_string()]);
133        c.is_tool_call = true;
134        mw.on_post_llm(&mut c).await.unwrap();
135
136        assert!(!c.skip_push);
137        assert_eq!(c.nudge_count, 0);
138    }
139
140    #[tokio::test]
141    async fn skips_when_full_text_empty() {
142        let mw = ToolEnforcementMiddleware::new(ToolEnforcementConfig::default());
143        let mut c = ctx(vec!["read".to_string()]);
144        c.full_text = String::new();
145        mw.on_post_llm(&mut c).await.unwrap();
146
147        assert!(!c.skip_push);
148        assert_eq!(c.nudge_count, 0);
149    }
150
151    #[tokio::test]
152    async fn skips_when_max_nudges_reached() {
153        let config = ToolEnforcementConfig {
154            max_nudges: 1,
155            ..ToolEnforcementConfig::default()
156        };
157        let mw = ToolEnforcementMiddleware::new(config);
158
159        let mut c = ctx(vec!["read".to_string()]);
160        c.nudge_count = 1;
161        mw.on_post_llm(&mut c).await.unwrap();
162
163        assert!(!c.skip_push);
164        assert_eq!(c.nudge_count, 1);
165        assert!(c.follow_up_message.is_none());
166    }
167
168    #[tokio::test]
169    async fn first_turn_only_skips_after_tool_calls() {
170        let mw = ToolEnforcementMiddleware::new(ToolEnforcementConfig::default());
171        let mut c = ctx(vec!["read".to_string()]);
172        c.total_tool_calls = 1;
173        mw.on_post_llm(&mut c).await.unwrap();
174
175        assert!(!c.skip_push);
176        assert_eq!(c.nudge_count, 0);
177    }
178}
agent_base/engine/tool_enforcement.rs

agent_base/engine/
tool_enforcement.rs
