PicoBot/src/tools/sleep.rs

use super::traits::{Tool, ToolResult};
use async_trait::async_trait;
use serde_json::json;
use std::time::Duration;

const MAX_SLEEP_SECONDS: u64 = 86_400;

pub struct SleepTool;

impl SleepTool {
    pub fn new() -> Self {
        Self
    }
}

impl Default for SleepTool {
    fn default() -> Self {
        Self::new()
    }
}

fn parse_seconds(args: &serde_json::Value) -> Result<u64, String> {
    let seconds = args
        .get("seconds")
        .and_then(serde_json::Value::as_u64)
        .ok_or_else(|| "seconds must be a non-negative integer".to_string())?;
    if seconds > MAX_SLEEP_SECONDS {
        return Err(format!(
            "seconds must not exceed {MAX_SLEEP_SECONDS} (24 hours)"
        ));
    }
    Ok(seconds)
}

#[async_trait]
impl Tool for SleepTool {
    fn name(&self) -> &str {
        "sleep"
    }

    fn description(&self) -> &str {
        "Pause the current agent execution for a specified number of whole seconds."
    }

    fn parameters_schema(&self) -> serde_json::Value {
        json!({
            "type": "object",
            "properties": {
                "seconds": {
                    "type": "integer",
                    "minimum": 0,
                    "maximum": MAX_SLEEP_SECONDS,
                    "description": "Number of whole seconds to wait, up to 24 hours."
                }
            },
            "required": ["seconds"]
        })
    }

    async fn execute(&self, args: serde_json::Value) -> anyhow::Result<ToolResult> {
        let seconds = match parse_seconds(&args) {
            Ok(seconds) => seconds,
            Err(error) => {
                return Ok(ToolResult {
                    success: false,
                    output: String::new(),
                    error: Some(error),
                });
            }
        };

        tokio::time::sleep(Duration::from_secs(seconds)).await;

        Ok(ToolResult {
            success: true,
            output: format!("Slept for {seconds} second(s)."),
            error: None,
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::agent::TurnEvent;
    use crate::providers::ToolCall;
    use crate::session::{ToolStatus, TurnBlock, TurnController, TurnStatus};
    use crate::tools::Tool;
    use serde_json::json;
    use std::time::Duration;

    #[test]
    fn exposes_sleep_metadata_and_schema() {
        let tool = SleepTool::new();
        let schema = tool.parameters_schema();

        assert_eq!(tool.name(), "sleep");
        assert!(tool.description().contains("current agent execution"));
        assert!(tool.description().contains("whole seconds"));
        assert_eq!(schema["type"], "object");
        assert_eq!(schema["required"], json!(["seconds"]));
        assert_eq!(schema["properties"]["seconds"]["type"], "integer");
        assert_eq!(schema["properties"]["seconds"]["minimum"], 0);
        assert_eq!(
            schema["properties"]["seconds"]["maximum"],
            MAX_SLEEP_SECONDS
        );
        assert!(schema.get("additionalProperties").is_none());
        assert!(!tool.read_only());
        assert!(!tool.concurrency_safe());
        assert!(!tool.exclusive());
    }

    #[tokio::test]
    async fn zero_seconds_returns_exact_success() {
        let result = SleepTool::new()
            .execute(json!({"seconds": 0}))
            .await
            .unwrap();

        assert!(result.success);
        assert_eq!(result.output, "Slept for 0 second(s).");
        assert_eq!(result.error, None);
    }

    #[tokio::test]
    async fn rejects_invalid_seconds() {
        let invalid_args = [
            json!({}),
            json!({"seconds": -1}),
            json!({"seconds": 0.5}),
            json!({"seconds": "1"}),
            json!({"seconds": 18_446_744_073_709_552_000.0_f64}),
            json!({"seconds": MAX_SLEEP_SECONDS + 1}),
        ];

        for args in invalid_args {
            let result = SleepTool::new().execute(args).await.unwrap();
            assert!(!result.success);
            assert!(result.output.is_empty());
            assert!(result.error.is_some());
        }
    }

    #[test]
    fn accepts_24_hour_boundary() {
        assert_eq!(
            parse_seconds(&json!({"seconds": MAX_SLEEP_SECONDS})),
            Ok(MAX_SLEEP_SECONDS)
        );
    }

    #[tokio::test(start_paused = true)]
    async fn waits_for_requested_seconds() {
        let handle = tokio::spawn(async { SleepTool::new().execute(json!({"seconds": 2})).await });
        tokio::task::yield_now().await;
        tokio::time::advance(Duration::from_secs(1)).await;
        tokio::task::yield_now().await;
        assert!(!handle.is_finished());
        tokio::time::advance(Duration::from_secs(1)).await;
        tokio::task::yield_now().await;
        assert!(handle.await.unwrap().unwrap().success);
    }

    #[tokio::test(start_paused = true)]
    async fn waits_up_to_24_hour_boundary() {
        let handle = tokio::spawn(async {
            SleepTool::new()
                .execute(json!({"seconds": MAX_SLEEP_SECONDS}))
                .await
        });
        tokio::task::yield_now().await;
        tokio::time::advance(Duration::from_secs(MAX_SLEEP_SECONDS - 1)).await;
        tokio::task::yield_now().await;
        assert!(!handle.is_finished());
        tokio::time::advance(Duration::from_secs(1)).await;
        tokio::task::yield_now().await;
        assert!(handle.await.unwrap().unwrap().success);
    }

    #[tokio::test(start_paused = true)]
    async fn cancellation_drops_an_active_sleep() {
        let handle = tokio::spawn(async {
            SleepTool::new()
                .execute(json!({"seconds": MAX_SLEEP_SECONDS}))
                .await
        });
        tokio::task::yield_now().await;
        assert!(!handle.is_finished());
        handle.abort();
        assert!(handle.await.unwrap_err().is_cancelled());
    }

    #[tokio::test(start_paused = true)]
    async fn user_cancellation_stops_sleep_and_terminalizes_its_tool_block() {
        let (controller, emitter, receiver) =
            TurnController::start("cli:test:sleep", "assistant-message");
        emitter
            .emit(TurnEvent::ToolStarted {
                iteration: 0,
                call: ToolCall {
                    id: "sleep-call".into(),
                    name: "sleep".into(),
                    arguments: json!({"seconds": MAX_SLEEP_SECONDS}),
                },
            })
            .unwrap();
        let (cancel_tx, cancel_rx) = tokio::sync::oneshot::channel::<()>();

        let handle = tokio::spawn(async move {
            let tool = SleepTool::new();
            tokio::select! {
                result = tool.execute(json!({"seconds": MAX_SLEEP_SECONDS})) => {
                    result.unwrap();
                    false
                }
                _ = cancel_rx => {
                    controller.cancel(Some("stopped by user".into()));
                    true
                }
            }
        });
        tokio::task::yield_now().await;
        drop(cancel_tx);

        assert!(handle.await.unwrap());
        let snapshot = receiver.borrow().clone();
        assert_eq!(snapshot.status, TurnStatus::Cancelled);
        assert!(matches!(
            &snapshot.blocks[0],
            TurnBlock::Tool {
                id,
                status: ToolStatus::Cancelled,
                ..
            } if id == "sleep-call"
        ));
    }
}