PicoBot/src/tools/sleep.rs

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

use crate::agent::steering::{TurnWakeupState, WakeupSource};

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 input_interrupt_policy(&self) -> crate::tools::InputInterruptPolicy {
        crate::tools::InputInterruptPolicy::WakeOnly
    }

    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> {
        self.execute_with_context(&crate::tools::ToolExecutionContext::default(), args)
            .await
            .map(|output| output.result)
    }

    async fn execute_with_context(
        &self,
        context: &crate::tools::ToolExecutionContext,
        args: serde_json::Value,
    ) -> anyhow::Result<crate::tools::ToolOutput> {
        let seconds = match parse_seconds(&args) {
            Ok(seconds) => seconds,
            Err(error) => {
                return Ok(ToolResult {
                    success: false,
                    output: String::new(),
                    error: Some(error),
                }
                .into());
            }
        };

        let started = std::time::Instant::now();
        let mut wakeup_rx = context
            .turn_wakeup
            .as_ref()
            .map(|handle| handle.receiver.clone());

        // Root interactive Turn: if inputs are already pending, do not wait
        // at all.  The watch revision is monotonic, so an input arriving
        // between this check and the select below still fires `changed()`.
        if let Some(rx) = wakeup_rx.as_mut() {
            let state = rx.borrow_and_update();
            if state.pending_total() > 0 {
                return Ok(ToolResult {
                    success: true,
                    output: wake_message(&state, started.elapsed(), 0),
                    error: None,
                }
                .into());
            }
        }

        let outcome = match wakeup_rx.as_mut() {
            Some(rx) => {
                tokio::select! {
                    biased;
                    _ = context.cancellation.cancelled() => {
                        anyhow::bail!("sleep cancelled");
                    }
                    _ = tokio::time::sleep(Duration::from_secs(seconds)) => {
                        WakeOutcome::Elapsed
                    }
                    changed = rx.changed() => {
                        let _ = changed;
                        let state = rx.borrow_and_update();
                        WakeOutcome::InputArrived(state.clone())
                    }
                }
            }
            // Child runs and continuation Turns have no session input lane:
            // their sleep answers only the timer, run cancellation, timeout
            // and shutdown.
            None => {
                tokio::select! {
                    biased;
                    _ = context.cancellation.cancelled() => {
                        anyhow::bail!("sleep cancelled");
                    }
                    _ = tokio::time::sleep(Duration::from_secs(seconds)) => {
                        WakeOutcome::Elapsed
                    }
                }
            }
        };

        let output = match outcome {
            WakeOutcome::Elapsed => format!("Slept for {seconds} second(s)."),
            WakeOutcome::InputArrived(state) => wake_message(&state, started.elapsed(), seconds),
        };
        Ok(ToolResult {
            success: true,
            output,
            error: None,
        }
        .into())
    }
}

enum WakeOutcome {
    Elapsed,
    InputArrived(TurnWakeupState),
}

/// Build the model-visible wake message.  Steer wakes describe the source,
/// run identity and a safe preview; queue wakes only state the type/count and
/// explicitly promise the content stays out of the current Turn.
fn wake_message(state: &TurnWakeupState, waited: std::time::Duration, planned: u64) -> String {
    let waited_secs = waited.as_secs();
    let mut message = format!("Sleep 提前结束：已等待 {waited_secs} 秒");
    if planned > 0 {
        message.push_str(&format!("（原计划 {planned} 秒）"));
    }
    message.push('。');
    match &state.latest_source {
        Some(WakeupSource::UserSteer) => {
            message.push_str(" 收到一条新的用户输入，将在当前 Turn 的下一个安全边界注入。");
        }
        Some(WakeupSource::UserQueue) => {
            message.push_str(&format!(
                " 收到 {} 条排队输入。内容不会进入当前 Turn，将在当前工作结束后的下一 Turn处理。",
                state.pending_user_queue.max(1)
            ));
        }
        Some(WakeupSource::AgentSignal { run_id, agent_id }) => {
            message.push_str(&format!(
                " 收到一条 steer AgentSignal（run_id={run_id}, agent={agent_id}）"
            ));
            if let Some(preview) = state.latest_safe_preview.as_deref() {
                message.push_str(&format!("：{preview}"));
            }
            message.push_str("。该信号将在当前 Turn 的下一个安全边界注入。");
        }
        Some(WakeupSource::AgentCompletion { run_id, agent_id }) => {
            message.push_str(&format!(
                " 收到一条 steer AgentCompletion（run_id={run_id}, agent={agent_id}），将在当前 Turn 的下一个安全边界注入。"
            ));
        }
        Some(WakeupSource::AgentQueue) | None => {
            message.push_str(&format!(
                " 收到 {} 条排队输入。内容不会进入当前 Turn，将在当前工作结束后的下一 Turn处理。",
                state.pending_agent_queue.max(1)
            ));
        }
    }
    message
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::agent::TurnEvent;
    use crate::agent::steering::TurnWakeupPublisher;
    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());
        assert_eq!(
            tool.input_interrupt_policy(),
            crate::tools::InputInterruptPolicy::WakeOnly
        );
    }

    #[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"
        ));
    }

    #[tokio::test(start_paused = true)]
    async fn cancellation_token_ends_sleep_before_timer() {
        let context = crate::tools::ToolExecutionContext::default();
        let token = context.cancellation.clone();
        let handle = tokio::spawn(async move {
            SleepTool::new()
                .execute_with_context(&context, json!({"seconds": MAX_SLEEP_SECONDS}))
                .await
        });
        tokio::task::yield_now().await;
        assert!(!handle.is_finished());
        token.cancel();
        tokio::task::yield_now().await;
        let error = handle.await.unwrap().unwrap_err();
        assert!(error.to_string().contains("cancelled"));
    }

    #[tokio::test(start_paused = true)]
    async fn pre_cancelled_context_never_enters_sleep() {
        let context = crate::tools::ToolExecutionContext::default();
        context.cancellation.cancel();
        let error = SleepTool::new()
            .execute_with_context(&context, json!({"seconds": 60}))
            .await
            .unwrap_err();
        assert!(error.to_string().contains("cancelled"));
    }

    #[tokio::test(start_paused = true)]
    async fn pending_input_before_listen_returns_immediately() {
        let publisher = TurnWakeupPublisher::new();
        let handle = publisher.subscribe();
        publisher.publish(TurnWakeupState {
            pending_user_steer: 1,
            latest_source: Some(WakeupSource::UserSteer),
            ..Default::default()
        });
        let context = crate::tools::ToolExecutionContext::default().with_turn_wakeup(handle);
        let result = SleepTool::new()
            .execute_with_context(&context, json!({"seconds": 3600}))
            .await
            .unwrap();
        assert!(result.result.success);
        assert!(result.result.output.contains("提前结束"));
        assert!(result.result.output.contains("用户输入"));
    }

    #[tokio::test(start_paused = true)]
    async fn steer_publish_wakes_sleep_with_source_and_preview() {
        let publisher = TurnWakeupPublisher::new();
        let handle = publisher.subscribe();
        let context = crate::tools::ToolExecutionContext::default().with_turn_wakeup(handle);
        let tool = SleepTool::new();
        let wait = tokio::spawn(async move {
            tool.execute_with_context(&context, json!({"seconds": 3600}))
                .await
                .unwrap()
                .result
                .output
        });
        tokio::task::yield_now().await;
        assert!(!wait.is_finished());
        publisher.publish(TurnWakeupState {
            pending_agent_steer: 1,
            latest_source: Some(WakeupSource::AgentSignal {
                run_id: "run-123".to_string(),
                agent_id: "monitor".to_string(),
            }),
            latest_safe_preview: Some("服务错误率超过 5%".to_string()),
            ..Default::default()
        });
        tokio::task::yield_now().await;
        let output = wait.await.unwrap();
        assert!(output.contains("提前结束"));
        assert!(output.contains("run-123"));
        assert!(output.contains("服务错误率超过 5%"));
        assert!(output.contains("安全边界注入"));
    }

    #[tokio::test(start_paused = true)]
    async fn queue_publish_wakes_sleep_without_content() {
        let publisher = TurnWakeupPublisher::new();
        let handle = publisher.subscribe();
        let context = crate::tools::ToolExecutionContext::default().with_turn_wakeup(handle);
        let tool = SleepTool::new();
        let wait = tokio::spawn(async move {
            tool.execute_with_context(&context, json!({"seconds": 3600}))
                .await
                .unwrap()
                .result
                .output
        });
        tokio::task::yield_now().await;
        publisher.publish(TurnWakeupState {
            pending_agent_queue: 1,
            latest_source: Some(WakeupSource::AgentQueue),
            ..Default::default()
        });
        tokio::task::yield_now().await;
        let output = wait.await.unwrap();
        assert!(output.contains("排队输入"));
        assert!(output.contains("不会进入当前 Turn"));
        assert!(!output.contains("run-"));
    }

    #[tokio::test(start_paused = true)]
    async fn child_sleep_without_handle_is_not_woken_by_publishes() {
        let publisher = TurnWakeupPublisher::new();
        let _handle = publisher.subscribe();
        let context = crate::tools::ToolExecutionContext::default();
        let tool = SleepTool::new();
        let wait = tokio::spawn(async move {
            tool.execute_with_context(&context, json!({"seconds": 30}))
                .await
                .unwrap()
                .result
                .output
        });
        tokio::task::yield_now().await;
        publisher.publish(TurnWakeupState {
            pending_agent_steer: 1,
            latest_source: Some(WakeupSource::AgentSignal {
                run_id: "run-9".to_string(),
                agent_id: "a".to_string(),
            }),
            ..Default::default()
        });
        tokio::task::yield_now().await;
        assert!(!wait.is_finished());
        tokio::time::advance(Duration::from_secs(30)).await;
        tokio::task::yield_now().await;
        assert!(wait.await.unwrap().contains("Slept for 30"));
    }

    #[tokio::test(start_paused = true)]
    async fn pre_listen_publish_does_not_lose_the_wake() {
        // Publish BEFORE the sleep subscribes its own receiver: watch keeps
        // the latest value, so the borrow_and_update pre-check sees it.
        let publisher = TurnWakeupPublisher::new();
        let handle = publisher.subscribe();
        publisher.publish(TurnWakeupState {
            pending_agent_queue: 2,
            latest_source: Some(WakeupSource::AgentQueue),
            ..Default::default()
        });
        let context = crate::tools::ToolExecutionContext::default().with_turn_wakeup(handle);
        let result = SleepTool::new()
            .execute_with_context(&context, json!({"seconds": 3600}))
            .await
            .unwrap();
        assert!(result.result.success);
        assert!(result.result.output.contains("排队输入"));
    }
}