Skip to main content

bamboo_tools/tools/
sleep.rs

1use async_trait::async_trait;
2use bamboo_agent_core::{Tool, ToolError, ToolResult};
3use serde::Deserialize;
4use serde_json::json;
5use tokio::time::{sleep, Duration};
6
7const MAX_SLEEP_SECONDS: f64 = 300.0;
8
9#[derive(Debug, Deserialize)]
10struct SleepArgs {
11    seconds: f64,
12    #[serde(default)]
13    reason: Option<String>,
14}
15
16/// Pause execution for a short duration.
17pub struct SleepTool;
18
19impl SleepTool {
20    pub fn new() -> Self {
21        Self
22    }
23}
24
25impl Default for SleepTool {
26    fn default() -> Self {
27        Self::new()
28    }
29}
30
31#[async_trait]
32impl Tool for SleepTool {
33    fn name(&self) -> &str {
34        "Sleep"
35    }
36
37    fn description(&self) -> &str {
38        "Pause execution for a specified number of seconds (max 300s)"
39    }
40
41    fn mutability(&self) -> crate::ToolMutability {
42        crate::ToolMutability::ReadOnly
43    }
44
45    fn concurrency_safe(&self) -> bool {
46        true
47    }
48
49    fn parameters_schema(&self) -> serde_json::Value {
50        json!({
51            "type": "object",
52            "properties": {
53                "seconds": {
54                    "type": "number",
55                    "description": "Seconds to sleep, can be fractional"
56                },
57                "reason": {
58                    "type": "string",
59                    "description": "Optional reason for logging"
60                }
61            },
62            "required": ["seconds"],
63            "additionalProperties": false
64        })
65    }
66
67    async fn execute(&self, args: serde_json::Value) -> Result<ToolResult, ToolError> {
68        let parsed: SleepArgs = serde_json::from_value(args)
69            .map_err(|e| ToolError::InvalidArguments(format!("Invalid sleep args: {e}")))?;
70
71        if parsed.seconds < 0.0 {
72            return Err(ToolError::InvalidArguments(
73                "seconds cannot be negative".to_string(),
74            ));
75        }
76        if parsed.seconds > MAX_SLEEP_SECONDS {
77            return Err(ToolError::InvalidArguments(format!(
78                "seconds cannot exceed {MAX_SLEEP_SECONDS}"
79            )));
80        }
81
82        if let Some(reason) = parsed.reason.as_deref() {
83            tracing::info!("Sleeping for {} seconds: {}", parsed.seconds, reason);
84        } else {
85            tracing::info!("Sleeping for {} seconds", parsed.seconds);
86        }
87
88        sleep(Duration::from_secs_f64(parsed.seconds)).await;
89
90        Ok(ToolResult {
91            success: true,
92            result: format!(
93                "Slept for {} seconds{}",
94                parsed.seconds,
95                parsed
96                    .reason
97                    .as_deref()
98                    .map(|r| format!(" ({r})"))
99                    .unwrap_or_default()
100            ),
101            display_preference: None,
102        })
103    }
104}
105
106#[cfg(test)]
107mod tests {
108    use super::*;
109    use std::time::Instant;
110
111    #[tokio::test]
112    async fn sleep_tool_waits_and_returns_success() {
113        let tool = SleepTool::new();
114        let start = Instant::now();
115        let result = tool.execute(json!({"seconds": 0.01})).await.unwrap();
116        assert!(result.success);
117        assert!(start.elapsed().as_millis() >= 10);
118    }
119
120    #[tokio::test(start_paused = true)]
121    async fn sleep_tool_accepts_valid_seconds() {
122        let tool = SleepTool::new();
123
124        // Minimum value
125        let result = tool.execute(json!({"seconds": 0.0})).await.unwrap();
126        assert!(result.success);
127
128        // Small positive value
129        let result = tool.execute(json!({"seconds": 0.001})).await.unwrap();
130        assert!(result.success);
131
132        // Maximum allowed value (300.0)
133        let result = tool.execute(json!({"seconds": 300.0})).await.unwrap();
134        assert!(result.success);
135    }
136
137    #[tokio::test]
138    async fn sleep_tool_rejects_negative_seconds() {
139        let tool = SleepTool::new();
140        let result = tool.execute(json!({"seconds": -1.0})).await;
141        assert!(result.is_err());
142        let error = result.unwrap_err();
143        assert!(matches!(error, ToolError::InvalidArguments(_)));
144    }
145
146    #[tokio::test]
147    async fn sleep_tool_rejects_seconds_exceeding_max() {
148        let tool = SleepTool::new();
149        let result = tool.execute(json!({"seconds": 300.1})).await;
150        assert!(result.is_err());
151        let error = result.unwrap_err();
152        if let ToolError::InvalidArguments(msg) = error {
153            assert!(msg.contains("cannot exceed"));
154            assert!(msg.contains("300"));
155        } else {
156            panic!("Expected InvalidArguments error");
157        }
158    }
159
160    #[tokio::test]
161    async fn sleep_tool_includes_reason_in_result() {
162        let tool = SleepTool::new();
163        let result = tool
164            .execute(json!({
165                "seconds": 0.001,
166                "reason": "testing sleep"
167            }))
168            .await
169            .unwrap();
170
171        assert!(result.success);
172        assert!(result.result.contains("testing sleep"));
173        assert!(result.result.contains("(testing sleep)"));
174    }
175
176    #[tokio::test]
177    async fn sleep_tool_works_without_reason() {
178        let tool = SleepTool::new();
179        let result = tool.execute(json!({"seconds": 0.001})).await.unwrap();
180
181        assert!(result.success);
182        assert!(result.result.contains("Slept for 0.001 seconds"));
183        assert!(!result.result.contains("("));
184    }
185
186    #[tokio::test]
187    async fn sleep_tool_rejects_missing_seconds() {
188        let tool = SleepTool::new();
189        let result = tool.execute(json!({})).await;
190        assert!(result.is_err());
191    }
192
193    #[tokio::test]
194    async fn sleep_tool_rejects_invalid_seconds_type() {
195        let tool = SleepTool::new();
196        let result = tool.execute(json!({"seconds": "not a number"})).await;
197        assert!(result.is_err());
198    }
199
200    #[tokio::test]
201    async fn sleep_tool_accepts_fractional_seconds() {
202        let tool = SleepTool::new();
203        let start = Instant::now();
204        let result = tool.execute(json!({"seconds": 0.05})).await.unwrap();
205
206        assert!(result.success);
207        assert!(result.result.contains("0.05"));
208        assert!(start.elapsed().as_millis() >= 50);
209    }
210
211    #[test]
212    fn sleep_tool_has_correct_name() {
213        let tool = SleepTool::new();
214        assert_eq!(tool.name(), "Sleep");
215    }
216
217    #[test]
218    fn sleep_tool_has_description() {
219        let tool = SleepTool::new();
220        assert!(!tool.description().is_empty());
221        assert!(tool.description().contains("300"));
222    }
223
224    #[test]
225    fn sleep_tool_parameters_schema_has_required_fields() {
226        let tool = SleepTool::new();
227        let schema = tool.parameters_schema();
228
229        assert_eq!(schema["type"], "object");
230        assert!(schema["properties"]["seconds"].is_object());
231        assert!(schema["properties"]["reason"].is_object());
232        assert!(schema["required"]
233            .as_array()
234            .unwrap()
235            .contains(&json!("seconds")));
236    }
237
238    #[tokio::test]
239    async fn sleep_tool_default_impl() {
240        let tool = SleepTool::default();
241        let result = tool.execute(json!({"seconds": 0.001})).await.unwrap();
242        assert!(result.success);
243    }
244
245    #[tokio::test]
246    async fn sleep_tool_handles_zero_seconds() {
247        let tool = SleepTool::new();
248        let result = tool.execute(json!({"seconds": 0.0})).await.unwrap();
249
250        assert!(result.success);
251        assert!(result.result.contains("0 seconds"));
252    }
253
254    #[tokio::test]
255    async fn sleep_tool_reason_with_special_characters() {
256        let tool = SleepTool::new();
257        let result = tool
258            .execute(json!({
259                "seconds": 0.001,
260                "reason": "等待数据 🎯 (waiting for data)"
261            }))
262            .await
263            .unwrap();
264
265        assert!(result.success);
266        assert!(result.result.contains("等待数据 🎯"));
267    }
268
269    #[tokio::test]
270    async fn sleep_tool_reason_empty_string() {
271        let tool = SleepTool::new();
272        let result = tool
273            .execute(json!({
274                "seconds": 0.001,
275                "reason": ""
276            }))
277            .await
278            .unwrap();
279
280        assert!(result.success);
281        // Empty string is still treated as a reason, so it will show " ()"
282        assert!(result.result.contains(" ()"));
283    }
284}