From 44d76c6a6dd04fa2efc302b906ac8774267a1272 Mon Sep 17 00:00:00 2001 From: Boyang Niu Date: Thu, 23 Jul 2026 00:49:18 +0000 Subject: [PATCH] Wake sleeping threads for queued agent mail (#34852) ## Why An idle thread with an outstanding durable sleep must resume when agent work arrives, even when that message would normally remain queued without triggering a turn. ## What changed - Treat any pending mailbox message as wake-up work while a durable sleep is attached to the thread. - Keep requiring `trigger_turn` for idle threads that are not durably asleep. - Cover waking a sleeping root thread from queue-only agent mail and persisting the message in thread history. GitOrigin-RevId: e361ac1b104436d4d63aebfe5e2422691a1cc139 --- codex-rs/core/src/session/handlers.rs | 2 +- codex-rs/core/src/tasks/mod.rs | 20 +++++-- codex-rs/core/tests/suite/pending_input.rs | 64 +++++++++++++++++++++- 3 files changed, 80 insertions(+), 6 deletions(-) diff --git a/codex-rs/core/src/session/handlers.rs b/codex-rs/core/src/session/handlers.rs index 23d24c3f5ce5..1594b7358413 100644 --- a/codex-rs/core/src/session/handlers.rs +++ b/codex-rs/core/src/session/handlers.rs @@ -280,7 +280,7 @@ pub async fn inter_agent_communication( .enqueue_mailbox_communication(communication) .await; crate::agent_communication::emit_agent_communication_receive(&sub_id); - if trigger_turn { + if trigger_turn || sess.has_outstanding_durable_sleep() { sess.maybe_start_turn_for_pending_work_with_sub_id(sub_id) .await; } diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index 1893568e416e..d931e8bd6575 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -451,9 +451,18 @@ impl Session { turn.task = Some(running_task); } + /// Returns whether an extension has marked this thread as durably asleep. + pub(crate) fn has_outstanding_durable_sleep(&self) -> bool { + self.services + .thread_extension_data + .get::() + .is_some() + } + /// Starts a regular turn when the session is idle and pending work is waiting. /// - /// Pending work currently includes mailbox mail marked with `trigger_turn`. + /// Pending work includes mailbox mail marked with `trigger_turn`, or any mailbox mail while + /// an outstanding durable sleep is attached to the thread. /// /// This helper generates a fresh sub-id for the synthetic turn before delegating to the /// explicit-sub-id variant. @@ -469,13 +478,16 @@ impl Session { /// Starts a regular turn with the provided sub-id when pending work should wake an idle /// session. /// - /// The turn is created only when there is mailbox mail marked with `trigger_turn`, and only - /// if the session is currently idle. + /// The turn is created only when the session is idle and mailbox mail either requests a turn + /// or can wake an outstanding durable sleep. pub(crate) async fn maybe_start_turn_for_pending_work_with_sub_id( self: &Arc, sub_id: String, ) { - if !self.input_queue.has_trigger_turn_mailbox_items().await { + if !self.input_queue.has_pending_mailbox_items().await + || (!self.input_queue.has_trigger_turn_mailbox_items().await + && !self.has_outstanding_durable_sleep()) + { return; } diff --git a/codex-rs/core/tests/suite/pending_input.rs b/codex-rs/core/tests/suite/pending_input.rs index d7c972d5cb68..a5a3df20225e 100644 --- a/codex-rs/core/tests/suite/pending_input.rs +++ b/codex-rs/core/tests/suite/pending_input.rs @@ -193,7 +193,7 @@ async fn steer_user_input(codex: &CodexThread, text: &str) { .expect("steer user input"); } -async fn submit_queue_only_agent_mail(codex: &CodexThread, text: &str) { +async fn enqueue_queue_only_agent_mail(codex: &CodexThread, text: &str) { codex .submit(Op::InterAgentCommunication { communication: InterAgentCommunication::new( @@ -206,6 +206,10 @@ async fn submit_queue_only_agent_mail(codex: &CodexThread, text: &str) { }) .await .expect("submit queue-only agent mail"); +} + +async fn submit_queue_only_agent_mail(codex: &CodexThread, text: &str) { + enqueue_queue_only_agent_mail(codex, text).await; codex .submit(Op::RealtimeConversationListVoices) .await @@ -294,6 +298,64 @@ async fn wait_for_sleep_item_completed(codex: &CodexThread, call_id: &str, durat ); } +struct SleepingRootExtension; + +impl codex_extension_api::ThreadLifecycleContributor + for SleepingRootExtension +{ + fn on_thread_start<'a>( + &'a self, + input: codex_extension_api::ThreadStartInput<'a, codex_core::config::Config>, + ) -> codex_extension_api::ExtensionFuture<'a, ()> { + Box::pin(async move { + input.thread_store.insert(SleepItem { + id: "clock-wait-1".to_string(), + duration_ms: 60_000, + }); + }) + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn queue_only_agent_mail_wakes_sleeping_root_and_persists_message() { + const CHILD_MESSAGE: &str = "worker completed"; + + let (server, _completions) = + start_streaming_sse_server(vec![response_completed_chunks("resp-1")]).await; + let mut extensions = + codex_extension_api::ExtensionRegistryBuilder::::new(); + extensions.thread_lifecycle_contributor(Arc::new(SleepingRootExtension)); + let codex = test_codex() + .with_model("gpt-5.4") + .with_extensions(Arc::new(extensions.build())) + .build_with_streaming_server(&server) + .await + .expect("build Codex test session") + .codex; + + enqueue_queue_only_agent_mail(&codex, CHILD_MESSAGE).await; + wait_for_turn_complete(&codex).await; + + assert_eq!(server.requests().await.len(), 1); + let history = codex + .load_history(/*include_archived*/ true) + .await + .expect("load persisted thread history"); + assert!(history.items.iter().any(|item| { + matches!( + item, + RolloutItem::ResponseItem(codex_protocol::models::ResponseItem::AgentMessage { + content, + .. + }) if content.iter().any(|content| matches!( + content, + codex_protocol::models::AgentMessageInputContent::InputText { text } + if text == CHILD_MESSAGE + )) + ) + })); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn steer_interrupts_wait_agent_and_is_sent_in_follow_up_request() { const WAIT_CALL_ID: &str = "wait-call";