mirror of
https://github.com/pchuan98/codex.git
synced 2026-07-01 00:31:56 +08:00
[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:
committed by
GitHub
Unverified
parent
3b22498f69
commit
62c7f506d9
@@ -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(());
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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"
|
||||
"#
|
||||
),
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user