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 { 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 { 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 { 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("排队输入")); } }