[codex] poll external clock during sleep (#30113)

## Summary

- make the external app-server time provider establish sleep deadlines
using `currentTime/read`
- poll the external clock once per second and complete `clock.sleep`
when the deadline is reached
- keep the system-clock timer and existing steer/agent-message
interruption behavior unchanged

## Why

This lets training control `clock.sleep` through its existing external
simulated clock without adding separate sleep/wake protocol methods.

## Testing

- `just fmt`
- `just test -p codex-app-server
external_sleep_polls_current_time_and_emits_items`
This commit is contained in:
rka-oai
2026-06-25 13:46:42 -07:00
committed by GitHub
Unverified
parent 3b22498f69
commit 62c7f506d9
2 changed files with 106 additions and 38 deletions
+26 -3
View File
@@ -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<OutgoingMessageSender>,
@@ -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(());
}
}
})
}
}
+80 -35
View File
@@ -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<ItemStartedNotification> {
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<ItemCompletedNotification> {
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"
"#
),
)