diff --git a/codex-rs/app-server/src/current_time.rs b/codex-rs/app-server/src/current_time.rs index 205068be2..ae873bfc6 100644 --- a/codex-rs/app-server/src/current_time.rs +++ b/codex-rs/app-server/src/current_time.rs @@ -23,6 +23,7 @@ use crate::outgoing_message::OutgoingMessageSender; use crate::thread_state::ThreadStateManager; const CURRENT_TIME_REQUEST_TIMEOUT: Duration = Duration::from_secs(5); +const CURRENT_TIME_POLL_INTERVAL: Duration = Duration::from_secs(1); pub(crate) fn app_server_time_provider( outgoing: Arc, @@ -51,10 +52,32 @@ impl TimeProvider for AppServerTimeProvider { }) } - fn sleep(&self, _thread_id: ThreadId, duration: Duration) -> SleepFuture<'_> { + fn sleep(&self, thread_id: ThreadId, duration: Duration) -> SleepFuture<'_> { + let outgoing = self.outgoing.clone(); + let thread_state_manager = self.thread_state_manager.clone(); Box::pin(async move { - tokio::time::sleep(duration).await; - Ok(()) + let outgoing = outgoing + .upgrade() + .context("app-server current-time provider is unavailable")?; + let started_at = + request_current_time(outgoing.clone(), thread_state_manager.clone(), thread_id) + .await?; + let wake_at = started_at + .checked_add_signed( + chrono::Duration::from_std(duration) + .context("external sleep duration is outside the supported range")?, + ) + .context("external sleep deadline is outside the supported range")?; + + loop { + tokio::time::sleep(CURRENT_TIME_POLL_INTERVAL).await; + if request_current_time(outgoing.clone(), thread_state_manager.clone(), thread_id) + .await? + >= wake_at + { + return Ok(()); + } + } }) } } diff --git a/codex-rs/app-server/tests/suite/v2/sleep.rs b/codex-rs/app-server/tests/suite/v2/sleep.rs index 872983f76..8d88816fd 100644 --- a/codex-rs/app-server/tests/suite/v2/sleep.rs +++ b/codex-rs/app-server/tests/suite/v2/sleep.rs @@ -1,11 +1,12 @@ use anyhow::Result; use app_test_support::TestAppServer; use app_test_support::to_response; +use codex_app_server_protocol::CurrentTimeReadResponse; use codex_app_server_protocol::ItemCompletedNotification; use codex_app_server_protocol::ItemStartedNotification; -use codex_app_server_protocol::JSONRPCMessage; use codex_app_server_protocol::JSONRPCResponse; use codex_app_server_protocol::RequestId; +use codex_app_server_protocol::ServerRequest; use codex_app_server_protocol::ThreadItem; use codex_app_server_protocol::ThreadStartParams; use codex_app_server_protocol::ThreadStartResponse; @@ -19,12 +20,16 @@ use std::time::Duration; use tempfile::TempDir; use tokio::time::timeout; +#[cfg(windows)] +const DEFAULT_READ_TIMEOUT: Duration = Duration::from_secs(25); +#[cfg(not(windows))] const DEFAULT_READ_TIMEOUT: Duration = Duration::from_secs(10); +const CURRENT_TIME_AT: i64 = 1_781_717_655; #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn sleep_emits_started_and_completed_items() -> Result<()> { +async fn external_sleep_polls_current_time_and_emits_items() -> Result<()> { const CALL_ID: &str = "sleep-1"; - const DURATION_MS: u64 = 1; + const DURATION_MS: u64 = 2_000; let server = responses::start_mock_server().await; responses::mount_sse_sequence( @@ -85,38 +90,19 @@ async fn sleep_emits_started_and_completed_items() -> Result<()> { .await??; let TurnStartResponse { turn, .. } = to_response(turn_start_response)?; - let (started, completed) = timeout(DEFAULT_READ_TIMEOUT, async { - let mut started = None; - let mut completed = None; - while started.is_none() || completed.is_none() { - let JSONRPCMessage::Notification(notification) = mcp.read_next_message().await? else { - continue; - }; - match notification.method.as_str() { - "item/started" => { - let payload: ItemStartedNotification = - serde_json::from_value(notification.params.expect("item/started params"))?; - if matches!(&payload.item, ThreadItem::Sleep { .. }) { - started = Some(payload); - } - } - "item/completed" => { - let payload: ItemCompletedNotification = serde_json::from_value( - notification.params.expect("item/completed params"), - )?; - if matches!(&payload.item, ThreadItem::Sleep { .. }) { - completed = Some(payload); - } - } - _ => {} - } - } - Ok::<_, anyhow::Error>(( - started.expect("sleep started"), - completed.expect("sleep completed"), - )) - }) - .await??; + // Read once for the initial reminder, then once to establish the sleep deadline. + respond_to_current_time_read(&mut mcp, &thread.id, CURRENT_TIME_AT).await?; + let started = wait_for_sleep_started(&mut mcp, CALL_ID).await?; + respond_to_current_time_read(&mut mcp, &thread.id, CURRENT_TIME_AT).await?; + + // The first poll remains below the deadline, so the provider must request time again. + respond_to_current_time_read(&mut mcp, &thread.id, CURRENT_TIME_AT + 1).await?; + respond_to_current_time_read(&mut mcp, &thread.id, CURRENT_TIME_AT + 2).await?; + + let completed = wait_for_sleep_completed(&mut mcp, CALL_ID).await?; + + // The next inference boundary reads the same external clock after the sleep completes. + respond_to_current_time_read(&mut mcp, &thread.id, CURRENT_TIME_AT + 2).await?; timeout( DEFAULT_READ_TIMEOUT, mcp.read_stream_until_notification_message("turn/completed"), @@ -150,6 +136,64 @@ async fn sleep_emits_started_and_completed_items() -> Result<()> { Ok(()) } +async fn wait_for_sleep_started( + mcp: &mut TestAppServer, + call_id: &str, +) -> Result { + loop { + let notification = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_notification_message("item/started"), + ) + .await??; + let started: ItemStartedNotification = + serde_json::from_value(notification.params.expect("item/started params"))?; + if matches!(&started.item, ThreadItem::Sleep { id, .. } if id == call_id) { + return Ok(started); + } + } +} + +async fn wait_for_sleep_completed( + mcp: &mut TestAppServer, + call_id: &str, +) -> Result { + loop { + let notification = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_notification_message("item/completed"), + ) + .await??; + let completed: ItemCompletedNotification = + serde_json::from_value(notification.params.expect("item/completed params"))?; + if matches!(&completed.item, ThreadItem::Sleep { id, .. } if id == call_id) { + return Ok(completed); + } + } +} + +async fn respond_to_current_time_read( + mcp: &mut TestAppServer, + thread_id: &str, + current_time_at: i64, +) -> Result<()> { + let request = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_request_message(), + ) + .await??; + let ServerRequest::CurrentTimeRead { request_id, params } = request else { + panic!("expected CurrentTimeRead request, got: {request:?}"); + }; + assert_eq!(params.thread_id, thread_id); + mcp.send_response( + request_id, + serde_json::to_value(CurrentTimeReadResponse { current_time_at })?, + ) + .await?; + Ok(()) +} + fn create_config_toml(codex_home: &Path, server_uri: &str) -> std::io::Result<()> { std::fs::write( codex_home.join("config.toml"), @@ -170,6 +214,7 @@ stream_max_retries = 0 [features.current_time_reminder] enabled = true sleep_tool = true +clock_source = "external" "# ), )