Skip to main content

recursive/tools/
schedule_wakeup.rs

1//! schedule_wakeup tool — lets the agent request a timed re-invocation.
2//!
3//! The tool writes a `WakeupRequest` into a shared slot; `AgentRuntime::run_loop`
4//! reads it after each turn completes and decides whether to sleep-then-loop.
5
6use std::sync::{Arc, Mutex};
7use std::time::Duration;
8
9use async_trait::async_trait;
10use serde_json::{json, Value};
11
12use super::Tool;
13use crate::error::Result;
14use crate::llm::ToolSpec;
15
16/// Shared slot where the tool writes a wakeup request.
17pub type WakeupSlot = Arc<Mutex<Option<WakeupRequest>>>;
18
19/// A wakeup request placed by the `schedule_wakeup` tool during a turn.
20#[derive(Debug, Clone)]
21pub struct WakeupRequest {
22    pub delay: Duration,
23    pub reason: String,
24    pub prompt: String,
25}
26
27/// Tool that lets the agent schedule its own next invocation.
28pub struct ScheduleWakeup {
29    slot: WakeupSlot,
30}
31
32impl ScheduleWakeup {
33    pub fn new(slot: WakeupSlot) -> Self {
34        Self { slot }
35    }
36}
37
38#[async_trait]
39impl Tool for ScheduleWakeup {
40    fn spec(&self) -> ToolSpec {
41        ToolSpec {
42            name: "schedule_wakeup".into(),
43            description: "Schedule the next loop iteration. The runner will \
44                          sleep for delay_secs then re-invoke the agent with \
45                          the given prompt. Call this to keep the loop alive."
46                .into(),
47            parameters: json!({
48                "type": "object",
49                "properties": {
50                    "delay_secs": {
51                        "type": "integer",
52                        "description": "Seconds to sleep before next turn (1-3600)",
53                        "minimum": 1,
54                        "maximum": 3600
55                    },
56                    "reason": {
57                        "type": "string",
58                        "description": "Why this wakeup is needed (shown in logs)"
59                    },
60                    "prompt": {
61                        "type": "string",
62                        "description": "Goal/context for the next turn"
63                    }
64                },
65                "required": ["delay_secs", "reason", "prompt"]
66            }),
67        }
68    }
69
70    async fn execute(&self, args: Value) -> Result<String> {
71        let delay_secs = args["delay_secs"].as_u64().unwrap_or(60).clamp(1, 3600);
72        let reason = args["reason"].as_str().unwrap_or("continue").to_string();
73        let prompt = args["prompt"].as_str().unwrap_or("").to_string();
74
75        let request = WakeupRequest {
76            delay: Duration::from_secs(delay_secs),
77            reason: reason.clone(),
78            prompt,
79        };
80
81        if let Ok(mut slot) = self.slot.lock() {
82            *slot = Some(request);
83        }
84
85        Ok(format!("Wakeup scheduled: {reason} in {delay_secs}s"))
86    }
87}
88
89#[cfg(test)]
90mod tests {
91    use super::*;
92
93    #[tokio::test]
94    async fn stores_wakeup_request() {
95        let slot: WakeupSlot = Arc::new(Mutex::new(None));
96        let tool = ScheduleWakeup::new(slot.clone());
97        let args = json!({"delay_secs": 30, "reason": "check status", "prompt": "check if done"});
98        let result = tool.execute(args).await.unwrap();
99        assert!(result.contains("30s"));
100        let req = slot.lock().unwrap().take().unwrap();
101        assert_eq!(req.delay, Duration::from_secs(30));
102        assert_eq!(req.prompt, "check if done");
103        assert_eq!(req.reason, "check status");
104    }
105
106    #[tokio::test]
107    async fn clamps_delay_max() {
108        let slot: WakeupSlot = Arc::new(Mutex::new(None));
109        let tool = ScheduleWakeup::new(slot.clone());
110        let args = json!({"delay_secs": 9999, "reason": "x", "prompt": "y"});
111        tool.execute(args).await.unwrap();
112        let req = slot.lock().unwrap().take().unwrap();
113        assert_eq!(req.delay, Duration::from_secs(3600));
114    }
115
116    #[tokio::test]
117    async fn clamps_delay_min() {
118        let slot: WakeupSlot = Arc::new(Mutex::new(None));
119        let tool = ScheduleWakeup::new(slot.clone());
120        let args = json!({"delay_secs": 0, "reason": "x", "prompt": "y"});
121        tool.execute(args).await.unwrap();
122        let req = slot.lock().unwrap().take().unwrap();
123        assert_eq!(req.delay, Duration::from_secs(1));
124    }
125
126    #[tokio::test]
127    async fn overwrites_previous_request() {
128        let slot: WakeupSlot = Arc::new(Mutex::new(None));
129        let tool = ScheduleWakeup::new(slot.clone());
130        tool.execute(json!({"delay_secs": 10, "reason": "first", "prompt": "a"}))
131            .await
132            .unwrap();
133        tool.execute(json!({"delay_secs": 20, "reason": "second", "prompt": "b"}))
134            .await
135            .unwrap();
136        let req = slot.lock().unwrap().take().unwrap();
137        assert_eq!(req.delay, Duration::from_secs(20));
138        assert_eq!(req.prompt, "b");
139    }
140}