From 960e878df4bf87b9d7fa0849f14561d7cdcca942 Mon Sep 17 00:00:00 2001 From: jif Date: Thu, 1 Oct 2026 19:58:59 +0000 Subject: [PATCH] Preserve queued agent mail across session eviction (#50087) ## Why Unread queue-only messages prevented idle agents from unloading, and sending a message to an evicted agent reloaded its session. Pending mail should not require an idle recipient to occupy a loaded thread slot. ## What changed - Retain queue-only mail in an in-memory runtime mailbox when a session unloads. Queue new messages for locally evicted agents without reloading them or starting a turn. - Add `AgentControl::take_mailbox` and `AgentControl::watch_mailbox` so sessions can consume retained mail and react to host mailbox notifications. - Preserve message order across eviction, reload, and subsequent follow-ups. Keep agents with outstanding durable sleep or turn-triggering mail loaded, and discard retained mail when an agent is removed or its tree shuts down. ## Testing Add integration coverage for retaining mail across eviction, sending while unloaded, and ordering messages around a follow-up after reload. Add a host controller test verifying that mailbox notifications wake `wait_agent` and that the message is delivered exactly once. GitOrigin-RevId: 4b55647f9ee9d1aaf3753f0688863bda81617327 --- codex-rs/core/src/agent/api.rs | 11 + codex-rs/core/src/agent/control.rs | 25 +++ codex-rs/core/src/agent/control/api.rs | 24 ++- codex-rs/core/src/agent/control/legacy.rs | 1 + codex-rs/core/src/agent/control/mailbox.rs | 114 ++++++++++ codex-rs/core/src/agent/control/residency.rs | 24 ++- codex-rs/core/src/agent/control/runtime.rs | 3 + codex-rs/core/src/agent_communication.rs | 3 + codex-rs/core/src/session/handlers.rs | 22 ++ codex-rs/core/src/session/input_queue.rs | 78 +++++-- codex-rs/core/src/session/session.rs | 2 +- .../core/tests/suite/agent_control_tests.rs | 89 ++++++++ codex-rs/core/tests/suite/agent_execution.rs | 3 + codex-rs/core/tests/suite/agent_mailbox.rs | 197 ++++++++++++++++++ 14 files changed, 573 insertions(+), 23 deletions(-) create mode 100644 codex-rs/core/src/agent/control/mailbox.rs create mode 100644 codex-rs/core/tests/suite/agent_mailbox.rs diff --git a/codex-rs/core/src/agent/api.rs b/codex-rs/core/src/agent/api.rs index e4d434667f..d9c894832c 100644 --- a/codex-rs/core/src/agent/api.rs +++ b/codex-rs/core/src/agent/api.rs @@ -63,6 +63,17 @@ pub trait AgentControl: Send + Sync { /// the root. Legacy user input can address loaded threads outside the agent registry. fn send(&self, request: SendRequest) -> BoxFuture<'_, Result>; + /// Take queued, non-turn-starting mail in order, without loading the recipient. + /// This in-memory operation performs no I/O; returning transfers ownership to the caller. + fn take_mailbox( + &self, + agent: ThreadId, + ) -> Vec; + + /// Observe whether unread mail is available. Subscribe before the first read so + /// arrivals cannot be missed; notifications do not consume mail or start a turn. + fn watch_mailbox(&self, agent: ThreadId) -> tokio::sync::watch::Receiver; + /// Load a recorded V2 child through its live immediate parent, without sending input. /// Implementations validate ownership and restore the child under the parent's current /// authority. Success makes the child available for attachment through its thread manager. diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index 8ab0fca3fb..a7af5100a8 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/codex-rs/core/src/agent/control.rs @@ -72,6 +72,7 @@ mod execution; mod inspection; mod interrupt; mod legacy; +mod mailbox; mod residency; mod resume; mod root_handoff; @@ -260,6 +261,30 @@ impl LocalAgentControl { ) -> CodexResult { let communication_for_log = crate::agent_communication::logging_enabled().then(|| communication.clone()); + { + // Keep unloaded delivery atomic with publication of a reloaded session. + let threads = state.threads.read().await; + if !communication.trigger_turn && !threads.contains_key(&agent_id) { + let _membership = self.runtime.admit_start()?; + self.runtime.ensure_agent_known(agent_id)?; + let submission_id = uuid::Uuid::now_v7().to_string(); + self.runtime.mailboxes.enqueue( + agent_id, + Some(submission_id.clone()), + vec![communication], + )?; + if let Some(communication) = communication_for_log { + crate::agent_communication::emit_agent_communication_send( + &submission_id, + &context, + &communication, + agent_id, + ); + } + return Ok(submission_id); + } + } + // Loaded recipients retain submission ordering with follow-ups and interrupts. let (parent_turn_id, root_turn_id) = if communication.trigger_turn { ( start_options.parent_turn_id.clone(), diff --git a/codex-rs/core/src/agent/control/api.rs b/codex-rs/core/src/agent/control/api.rs index 9ee1d50c57..ef0a699486 100644 --- a/codex-rs/core/src/agent/control/api.rs +++ b/codex-rs/core/src/agent/control/api.rs @@ -136,8 +136,17 @@ impl AgentControl for LocalAgentControl { "target agent is missing an agent_path".to_string(), ) })?; - self.ensure_v2_agent_loaded(resume_config, target, /*parent*/ None) - .await?; + // Cold-restored children still reload lazily on any message. Only + // locally evicted recipients can retain mail without reloading. + // Loaded recipients go straight to delivery, which rejects sends + // racing an in-progress eviction when it acquires the residency pin. + if mode == MessageDeliveryMode::TriggerTurn + || (self.runtime.upgrade()?.get_thread(target).await.is_err() + && self.runtime.registry.evicted_environments(target).is_none()) + { + self.ensure_v2_agent_loaded(resume_config, target, /*parent*/ None) + .await?; + } let communication = message.into_communication(author, receiver_path, mode); let kind = match mode { MessageDeliveryMode::QueueOnly => { @@ -165,6 +174,17 @@ impl AgentControl for LocalAgentControl { }) } + fn take_mailbox( + &self, + agent: ThreadId, + ) -> Vec { + self.runtime.mailboxes.take(agent) + } + + fn watch_mailbox(&self, agent: ThreadId) -> tokio::sync::watch::Receiver { + self.runtime.mailboxes.watch(agent) + } + fn ensure_child_loaded(&self, parent: ThreadId, child: ThreadId) -> BoxFuture<'_, Result<()>> { Box::pin(async move { let parent = self.runtime.upgrade()?.get_thread(parent).await?; diff --git a/codex-rs/core/src/agent/control/legacy.rs b/codex-rs/core/src/agent/control/legacy.rs index 1f4ee3834c..1c0b6e2fd7 100644 --- a/codex-rs/core/src/agent/control/legacy.rs +++ b/codex-rs/core/src/agent/control/legacy.rs @@ -41,6 +41,7 @@ impl LocalAgentControl { let _ = state.remove_thread(&agent_id).await; self.forget_v2_residency(agent_id); self.runtime.registry.release_spawned_thread(agent_id); + self.runtime.mailboxes.remove(agent_id); result } diff --git a/codex-rs/core/src/agent/control/mailbox.rs b/codex-rs/core/src/agent/control/mailbox.rs new file mode 100644 index 0000000000..6b5f4d88de --- /dev/null +++ b/codex-rs/core/src/agent/control/mailbox.rs @@ -0,0 +1,114 @@ +//! Retains queue-only mail while local sessions are unloaded. +//! Eviction transfers mail before releasing the recipient's residency guard. + +use crate::agent_communication::PENDING_MAILBOX_MESSAGES; +use codex_diagnostics::GaugeGuard; +use codex_protocol::ThreadId; +use codex_protocol::error::CodexErr; +use codex_protocol::error::Result; +use codex_protocol::protocol::InterAgentCommunication; +use std::collections::HashMap; +use std::sync::Mutex; +use std::sync::PoisonError; +use tokio::sync::watch; + +pub(super) struct Mailboxes(Mutex>>); + +impl Default for Mailboxes { + fn default() -> Self { + Self(Mutex::new(Some(HashMap::new()))) + } +} + +#[derive(Default)] +struct Mailbox { + pending: Vec, + state: watch::Sender, +} + +struct StoredMail { + id: Option, + communication: InterAgentCommunication, + _diagnostics_guard: GaugeGuard, +} + +impl Mailboxes { + pub(super) fn enqueue( + &self, + agent: ThreadId, + id: Option, + messages: Vec, + ) -> Result<()> { + let mut mailboxes = self.0.lock().unwrap_or_else(PoisonError::into_inner); + let mailbox = mailboxes + .as_mut() + .ok_or_else(|| CodexErr::InvalidRequest("agent runtime is shutting down".into()))? + .entry(agent) + .or_default(); + mailbox + .pending + .extend(messages.into_iter().map(|communication| StoredMail { + id: id.clone(), + communication, + _diagnostics_guard: PENDING_MAILBOX_MESSAGES.track(), + })); + mailbox.state.send_replace(!mailbox.pending.is_empty()); + Ok(()) + } + + pub(super) fn take(&self, agent: ThreadId) -> Vec { + let mut mailboxes = self.0.lock().unwrap_or_else(PoisonError::into_inner); + let Some(mailbox) = mailboxes + .as_mut() + .and_then(|mailboxes| mailboxes.get_mut(&agent)) + else { + return Vec::new(); + }; + let pending = std::mem::take(&mut mailbox.pending); + mailbox.state.send_if_modified(std::mem::take); + pending + .into_iter() + .map(|mail| { + // Mail transferred during eviction was already received by the old session. + if let Some(id) = mail.id { + crate::agent_communication::emit_agent_communication_receive(&id); + } + mail.communication + }) + .collect() + } + + pub(super) fn watch(&self, agent: ThreadId) -> watch::Receiver { + let mut mailboxes = self.0.lock().unwrap_or_else(PoisonError::into_inner); + match mailboxes.as_mut() { + Some(mailboxes) => { + // Private delegates and cancelled startups can stop without explicit removal. + // Retain unread mail and live subscriptions, but reclaim abandoned empty entries. + mailboxes.retain(|_, mailbox| { + !mailbox.pending.is_empty() || mailbox.state.receiver_count() != 0 + }); + mailboxes.entry(agent).or_default().state.subscribe() + } + None => watch::channel(/*init*/ false).1, + } + } + + pub(super) fn remove(&self, agent: ThreadId) { + if let Some(mailboxes) = self + .0 + .lock() + .unwrap_or_else(PoisonError::into_inner) + .as_mut() + { + mailboxes.remove(&agent); + } + } + + pub(super) fn close(&self) { + if let Some(mailboxes) = self.0.lock().unwrap_or_else(PoisonError::into_inner).take() { + for mailbox in mailboxes.values() { + mailbox.state.send_replace(false); + } + } + } +} diff --git a/codex-rs/core/src/agent/control/residency.rs b/codex-rs/core/src/agent/control/residency.rs index edaa97215b..b53f48b2ca 100644 --- a/codex-rs/core/src/agent/control/residency.rs +++ b/codex-rs/core/src/agent/control/residency.rs @@ -193,6 +193,23 @@ impl V2Residency { teardown.complete(); return false; } + // The submission loop has stopped and the residency guard excludes senders. + // Preserve unread queue-only mail before dropping the session that held it. + let mail = candidate_thread + .session + .input_queue + .drain_mailbox() + .await + .into_iter() + .map(|mail| mail.communication) + .collect(); + // A concurrent tree shutdown deliberately discards its unread mail. + let _ = candidate_thread + .session + .services + .local_agent_runtime + .mailboxes + .enqueue(candidate_thread_id, /*id*/ None, mail); let environments = candidate_thread.environment_selections().await; let mut threads = manager.threads.write().await; if threads @@ -276,7 +293,12 @@ async fn is_unloadable(thread: &CodexThread) -> bool { thread.agent_status().await, AgentStatus::Completed(_) | AgentStatus::Errored(_) | AgentStatus::Interrupted ) && thread.session.active_turn.lock().await.is_none() - && !thread.session.input_queue.has_pending_mailbox_items().await + && !thread.session.has_outstanding_durable_sleep() + && !thread + .session + .input_queue + .has_trigger_turn_mailbox_items() + .await } #[cfg(test)] diff --git a/codex-rs/core/src/agent/control/runtime.rs b/codex-rs/core/src/agent/control/runtime.rs index 5a621e7307..5fdf62560f 100644 --- a/codex-rs/core/src/agent/control/runtime.rs +++ b/codex-rs/core/src/agent/control/runtime.rs @@ -109,6 +109,7 @@ pub(crate) struct LocalAgentRuntime { Arc>>, pub(super) registry: Arc, pub(super) residency: Arc, + pub(super) mailboxes: Arc, /// Shared by every session in this tree, including private delegates. pub(crate) shutdown: CancellationToken, shutdown_state: Arc, @@ -125,6 +126,7 @@ impl LocalAgentRuntime { thread_id_generator, registry: Arc::default(), residency: Arc::default(), + mailboxes: Arc::default(), shutdown: CancellationToken::new(), shutdown_state: Arc::default(), agent_execution_limiter: Arc::default(), @@ -204,6 +206,7 @@ impl LocalAgentRuntime { pub(crate) fn request_shutdown(&self) -> Arc { self.shutdown_state.members.close(); self.shutdown.cancel(); + self.mailboxes.close(); Arc::clone(&self.shutdown_state) } diff --git a/codex-rs/core/src/agent_communication.rs b/codex-rs/core/src/agent_communication.rs index 91022fdc3e..9925d61f9f 100644 --- a/codex-rs/core/src/agent_communication.rs +++ b/codex-rs/core/src/agent_communication.rs @@ -1,6 +1,9 @@ use codex_protocol::ThreadId; use codex_protocol::protocol::InterAgentCommunication; +pub(crate) static PENDING_MAILBOX_MESSAGES: codex_diagnostics::Gauge = + codex_diagnostics::Gauge::new("core.mailbox.pending"); + const AGENT_COMMUNICATION_TARGET: &str = "codex_otel.agent_communication"; #[derive(Clone, Copy, Debug, PartialEq, Eq)] diff --git a/codex-rs/core/src/session/handlers.rs b/codex-rs/core/src/session/handlers.rs index 17af075cdf..32f1a7a73d 100644 --- a/codex-rs/core/src/session/handlers.rs +++ b/codex-rs/core/src/session/handlers.rs @@ -427,6 +427,10 @@ pub(super) async fn submission_loop( ) { // Session shutdown and tree shutdown both use the existing teardown handler. let mut shutdown_received = false; + let mut mailbox = sess.input_queue.mailbox_updates(); + if let Some(updates) = &mut mailbox { + updates.mark_changed(); + } loop { let sub = tokio::select! { biased; @@ -438,6 +442,24 @@ pub(super) async fn submission_loop( Ok(sub) => sub, Err(_) => break, }, + update = async { + match &mut mailbox { + Some(updates) => updates.changed().await.map(|()| *updates.borrow_and_update()), + None => std::future::pending().await, + } + } => { + match update { + Ok(true) => { + sess.input_queue.notify_mailbox(); + if sess.has_outstanding_durable_sleep() { + sess.maybe_start_turn_for_pending_work().await; + } + } + Ok(false) => {} + Err(_) => mailbox = None, + } + continue; + } }; if matches!(sub.op, Op::ResolveElicitation { .. }) { debug!(submission_id = %sub.id, operation = sub.op.kind(), "Submission"); diff --git a/codex-rs/core/src/session/input_queue.rs b/codex-rs/core/src/session/input_queue.rs index c73b75f8eb..de46d6a1d6 100644 --- a/codex-rs/core/src/session/input_queue.rs +++ b/codex-rs/core/src/session/input_queue.rs @@ -1,9 +1,11 @@ +use crate::agent::api::AgentControl; +use crate::agent_communication::PENDING_MAILBOX_MESSAGES; use crate::state::ActiveTurn; use crate::state::MailboxDeliveryPhase; use crate::state::TurnState; -use codex_diagnostics::Gauge; use codex_diagnostics::GaugeGuard; use codex_history::ResponseItemEnvelope; +use codex_protocol::ThreadId; use codex_protocol::models::ResponseItem; use codex_protocol::protocol::InterAgentCommunication; use codex_protocol::turn_input::TurnStartOptions; @@ -17,8 +19,6 @@ use tokio::sync::watch; use tokio_util::sync::CancellationToken; use tokio_util::task::AbortOnDropHandle; -static PENDING_MAILBOX_MESSAGES: Gauge = Gauge::new("core.mailbox.pending"); - /// Host capture metadata belonging to one input, including steers within another turn. #[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)] pub struct UserInputMetadata { @@ -94,11 +94,12 @@ pub(crate) struct TurnInputQueue { /// Session-scoped pending input storage and active-turn mailbox delivery coordination. pub(crate) struct InputQueue { activity_tx: watch::Sender, + controller: Option<(ThreadId, Arc, watch::Receiver)>, mailbox_pending_mails: Mutex>, } -struct PendingMailboxCommunication { - communication: InterAgentCommunication, +pub(crate) struct PendingMailboxCommunication { + pub(crate) communication: InterAgentCommunication, start_options: TurnStartOptions, _diagnostics_guard: GaugeGuard, } @@ -108,10 +109,44 @@ impl InputQueue { let (activity_tx, _) = watch::channel(InputQueueActivity::Mailbox); Self { activity_tx, + controller: None, mailbox_pending_mails: Mutex::new(VecDeque::new()), } } + pub(crate) fn with_controller(thread_id: ThreadId, control: Arc) -> Self { + let updates = control.watch_mailbox(thread_id); + Self { + controller: Some((thread_id, control, updates)), + ..Self::new() + } + } + + pub(crate) fn mailbox_updates(&self) -> Option> { + self.controller + .as_ref() + .map(|(_, _, updates)| updates.clone()) + } + + pub(crate) fn notify_mailbox(&self) { + self.activity_tx.send_replace(InputQueueActivity::Mailbox); + } + + fn read_mailbox(&self, pending: &mut VecDeque) { + if let Some((thread_id, control, _)) = &self.controller { + pending.extend( + control + .take_mailbox(*thread_id) + .into_iter() + .map(|communication| PendingMailboxCommunication { + communication, + start_options: TurnStartOptions::default(), + _diagnostics_guard: PENDING_MAILBOX_MESSAGES.track(), + }), + ); + } + } + pub(crate) async fn subscribe_activity( &self, turn_state: Option<&Mutex>, @@ -172,19 +207,23 @@ impl InputQueue { communication: InterAgentCommunication, start_options: TurnStartOptions, ) { - self.mailbox_pending_mails - .lock() - .await - .push_back(PendingMailboxCommunication { - communication, - start_options, - _diagnostics_guard: PENDING_MAILBOX_MESSAGES.track(), - }); + let mut pending = self.mailbox_pending_mails.lock().await; + // Mail retained while unloaded precedes new submissions to the loaded session. + self.read_mailbox(&mut pending); + pending.push_back(PendingMailboxCommunication { + communication, + start_options, + _diagnostics_guard: PENDING_MAILBOX_MESSAGES.track(), + }); self.activity_tx.send_replace(InputQueueActivity::Mailbox); } pub(crate) async fn has_pending_mailbox_items(&self) -> bool { !self.mailbox_pending_mails.lock().await.is_empty() + || self + .controller + .as_ref() + .is_some_and(|(_, _, updates)| *updates.borrow()) } pub(crate) async fn has_trigger_turn_mailbox_items(&self) -> bool { @@ -195,13 +234,14 @@ impl InputQueue { .any(|mail| mail.communication.trigger_turn) } + pub(crate) async fn drain_mailbox(&self) -> Vec { + let mut pending = self.mailbox_pending_mails.lock().await; + self.read_mailbox(&mut pending); + pending.drain(..).collect() + } + pub(crate) async fn drain_mailbox_input_items(&self) -> (Vec, TurnStartOptions) { - let pending_mails = self - .mailbox_pending_mails - .lock() - .await - .drain(..) - .collect::>(); + let pending_mails = self.drain_mailbox().await; // A later follow-up supersedes the earlier choice, including an omitted choice. let mut start_options = pending_mails .iter() diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index d4e6b96bc8..64f93fd42b 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -1842,7 +1842,7 @@ impl Session { .then(|| Mutex::new(Default::default())), active_turn: Mutex::new(None), async_hook_results, - input_queue: InputQueue::new(), + input_queue: InputQueue::with_controller(thread_id, Arc::clone(&services.agent_control)), services, git_enrichment_policy, fork_persistence, diff --git a/codex-rs/core/tests/suite/agent_control_tests.rs b/codex-rs/core/tests/suite/agent_control_tests.rs index 8e12fad0de..f47635a98f 100644 --- a/codex-rs/core/tests/suite/agent_control_tests.rs +++ b/codex-rs/core/tests/suite/agent_control_tests.rs @@ -49,6 +49,8 @@ use wiremock::MockServer; struct TestAgentControl { thread_id: ThreadId, service_tier: Mutex>, + mail: Mutex>, + mailbox_state: tokio::sync::watch::Sender, agents: Mutex>, } @@ -57,6 +59,8 @@ impl TestAgentControl { Self { thread_id, service_tier: Mutex::default(), + mail: Mutex::default(), + mailbox_state: tokio::sync::watch::channel(/*init*/ false).0, agents: Mutex::new(HashMap::from([("/root".into(), thread_id)])), } } @@ -96,6 +100,19 @@ impl AgentControl for TestAgentControl { Box::pin(async { Err(CodexErr::InvalidRequest("host send rejection".to_string())) }) } + fn take_mailbox( + &self, + _agent: ThreadId, + ) -> Vec { + let messages = std::mem::take(&mut *self.mail.lock().expect("mail lock")); + self.mailbox_state.send_if_modified(std::mem::take); + messages + } + + fn watch_mailbox(&self, _agent: ThreadId) -> tokio::sync::watch::Receiver { + self.mailbox_state.subscribe() + } + fn ensure_child_loaded( &self, parent: ThreadId, @@ -498,3 +515,75 @@ async fn collaboration_tools_dispatch_to_the_host_controller() -> anyhow::Result } use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::InternalSessionSource; + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn host_mailbox_notifies_waiting_agent_and_delivers_once() -> anyhow::Result<()> { + let server = responses::start_mock_server().await; + let (test, controller) = test_with_host_control(&server).await?; + responses::mount_sse_once( + &server, + responses::sse(vec![ + responses::ev_response_created("waiting"), + responses::ev_function_call_with_namespace( + "wait-mail", + "collaboration", + "wait_agent", + r#"{"timeout_ms":10000}"#, + ), + responses::ev_completed("waiting"), + ]), + ) + .await; + test.codex + .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text { + text: "Wait for a colleague.".into(), + text_elements: Vec::new(), + }])) + .await?; + wait_for_event(&test.codex, |event| matches!(event, + EventMsg::ItemStarted(event) if matches!(&event.item, + codex_protocol::items::TurnItem::CollabAgentToolCall(call) if call.id == "wait-mail") + )).await; + let next = responses::mount_sse_once( + &server, + responses::sse(vec![ + responses::ev_response_created("done"), + responses::ev_assistant_message("done", "Received the note."), + responses::ev_completed("done"), + ]), + ) + .await; + controller.mail.lock().expect("mail lock").push( + codex_protocol::protocol::InterAgentCommunication::new( + AgentPath::root().join("worker").expect("valid worker path"), + AgentPath::root(), + Vec::new(), + "Note from the host mailbox".into(), + /*trigger_turn*/ false, + ), + ); + controller.mailbox_state.send_replace(true); + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + let request = next.single_request(); + let wait_output = request + .function_call_output_text("wait-mail") + .expect("wait_agent output"); + assert_eq!( + serde_json::from_str::(&wait_output)?, + serde_json::json!({"message": "Wait completed.", "timed_out": false}) + ); + assert_eq!( + request + .inputs_of_type("agent_message") + .iter() + .filter(|item| item.to_string().contains("Note from the host mailbox")) + .count(), + 1 + ); + assert!(controller.mail.lock().expect("mail lock").is_empty()); + test.codex.shutdown_and_wait().await?; + Ok(()) +} diff --git a/codex-rs/core/tests/suite/agent_execution.rs b/codex-rs/core/tests/suite/agent_execution.rs index 3d1c99dfcb..2f74b79568 100644 --- a/codex-rs/core/tests/suite/agent_execution.rs +++ b/codex-rs/core/tests/suite/agent_execution.rs @@ -814,3 +814,6 @@ async fn tree_shutdown_reports_persistence_writer_failure() -> Result<()> { assert!(result.is_err()); Ok(()) } + +#[path = "agent_mailbox.rs"] +mod agent_mailbox; diff --git a/codex-rs/core/tests/suite/agent_mailbox.rs b/codex-rs/core/tests/suite/agent_mailbox.rs new file mode 100644 index 0000000000..357631a517 --- /dev/null +++ b/codex-rs/core/tests/suite/agent_mailbox.rs @@ -0,0 +1,197 @@ +//! Pending mail survives eviction without loading idle recipients. + +use super::*; +use crate::suite::settings_commits::COMMITTED_MODEL; +use crate::suite::settings_commits::PauseAfterCommit; +use codex_core::config::Config; +use codex_extension_api::ExtensionRegistryBuilder; +use codex_protocol::protocol::AgentStatus; +use codex_protocol::protocol::Op; +use core_test_support::ThreadIdle; +use pretty_assertions::assert_eq; +use std::sync::Arc; +use std::sync::Mutex; +use std::sync::mpsc; +use tokio::sync::oneshot; + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn idle_mail_survives_eviction_and_preserves_followup_order() -> Result<()> { + const MAIL: &str = "Keep this note for later"; + const LATE_MAIL: &str = "Another note while unloaded"; + const FOLLOWUP: &str = "Read the pending note"; + const AFTER_FOLLOWUP: &str = "This note must follow the new task"; + let server = start_mock_server().await; + mount_root_collaboration_call( + &server, + FIRST_PROMPT, + "first-call", + "spawn_agent", + json!({ "message": FIRST_TASK, "task_name": "first", "fork_turns": "none" }), + ) + .await; + mount_completed_worker(&server, FIRST_TASK, "first-call").await; + let (entered_tx, entered_rx) = oneshot::channel(); + let (release_tx, release_rx) = mpsc::channel(); + let mut extensions = ExtensionRegistryBuilder::::new(); + extensions.config_contributor(Arc::new(PauseAfterCommit { + gate: Mutex::new(Some((entered_tx, release_rx))), + })); + extensions.thread_lifecycle_contributor(Arc::new(ThreadIdle)); + let test = test_codex() + .with_model("gpt-5.6-sol") + .with_extensions(Arc::new(extensions.build())) + .with_config(|config| { + config.features.enable(Feature::Collab).unwrap(); + config.features.enable(Feature::MultiAgentV2).unwrap(); + config.multi_agent_v2.max_concurrent_threads_per_session = 2; + config.multi_agent_v2.hide_spawn_agent_metadata = true; + }) + .build_with_auto_env(&server) + .await?; + let mut created = test.thread_manager.subscribe_thread_created(); + test.submit_turn(FIRST_PROMPT).await?; + ThreadIdle::wait(&test.codex).await; + let first_id = created.recv().await?; + let first = test.thread_manager.get_thread(first_id).await?; + wait_for_event(&first, |event| matches!(event, EventMsg::TurnComplete(_))).await; + ThreadIdle::wait(&first).await; + + mount_root_collaboration_call( + &server, + "send a note", + "queued-mail", + "send_message", + json!({ "target": "first", "message": MAIL }), + ) + .await; + test.submit_turn("send a note").await?; + ThreadIdle::wait(&test.codex).await; + // Wait for delivery through the submission queue before attempting eviction. + submit_thread_settings(&first, ThreadSettingsOverrides::default()).await?; + assert_eq!( + first.agent_status().await, + AgentStatus::Completed(Some("worker completed".into())) + ); + drop(first); + + mount_root_collaboration_call( + &server, + "spawn while child is idle", + "replacement-spawn", + "spawn_agent", + json!({ "message": SECOND_TASK, "task_name": "replacement", "fork_turns": "none" }), + ) + .await; + mount_completed_worker(&server, SECOND_TASK, "replacement-spawn").await; + test.submit_turn("spawn while child is idle").await?; + ThreadIdle::wait(&test.codex).await; + let replacement_id = created.recv().await?; + let replacement = test.thread_manager.get_thread(replacement_id).await?; + wait_for_event(&replacement, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + ThreadIdle::wait(&replacement).await; + assert!(test.thread_manager.get_thread(first_id).await.is_err()); + assert_eq!(test.thread_manager.list_thread_ids().await.len(), 2); + mount_root_collaboration_call( + &server, + "send while unloaded", + "late-mail", + "send_message", + json!({ "target": "first", "message": LATE_MAIL }), + ) + .await; + test.submit_turn("send while unloaded").await?; + ThreadIdle::wait(&test.codex).await; + assert!(test.thread_manager.get_thread(first_id).await.is_err()); + + // Reload without starting a turn, then hold dispatch while both tools submit. + // The older queued notes must precede the follow-up, and the newer note must follow it. + test.thread_manager + .ensure_multi_agent_v2_child_loaded(first_id) + .await?; + let resumed = test.thread_manager.get_thread(first_id).await?; + resumed + .submit(Op::ThreadSettings { + thread_settings: ThreadSettingsOverrides { + model: Some(COMMITTED_MODEL.to_string()), + ..Default::default() + }, + reply: None, + }) + .await?; + entered_rx.await?; + mount_root_collaboration_call( + &server, + "drain the note", + "followup", + "followup_task", + json!({ "target": "first", "message": FOLLOWUP }), + ) + .await; + mount_sse_once_match( + &server, + |request: &wiremock::Request| { + body_contains(request, FOLLOWUP) && !has_function_call_output(request, "followup") + }, + sse(vec![ + ev_response_created("worker-wait"), + ev_function_call_with_namespace( + "wait-note", + MULTI_AGENT_V2_NAMESPACE, + "wait_agent", + r#"{"timeout_ms":1}"#, + ), + ev_completed("worker-wait"), + ]), + ) + .await; + let delivered = mount_sse_once_match( + &server, + |request: &wiremock::Request| has_function_call_output(request, "wait-note"), + sse(vec![ + ev_assistant_message("done", "Read all notes"), + ev_completed("done"), + ]), + ) + .await; + test.submit_turn("drain the note").await?; + ThreadIdle::wait(&test.codex).await; + mount_root_collaboration_call( + &server, + "send after followup", + "after-followup", + "send_message", + json!({ "target": "first", "message": AFTER_FOLLOWUP }), + ) + .await; + test.submit_turn("send after followup").await?; + ThreadIdle::wait(&test.codex).await; + release_tx.send(())?; + wait_for_event(&resumed, |event| matches!(event, EventMsg::TurnComplete(_))).await; + ThreadIdle::wait(&resumed).await; + let request = delivered + .requests() + .into_iter() + .find(|request| request.function_call_output_text("wait-note").is_some()) + .expect("worker continuation after waiting for mail"); + let messages = request.inputs_of_type("agent_message"); + let expected = [MAIL, LATE_MAIL, FOLLOWUP, AFTER_FOLLOWUP]; + let received = messages + .iter() + .filter_map(|item| { + expected + .iter() + .find(|text| item.to_string().contains(**text)) + .copied() + }) + .collect::>(); + assert_eq!(received, expected); + assert_eq!(test.thread_manager.list_thread_ids().await.len(), 2); + + test.thread_manager + .shutdown_all_threads_bounded(Duration::from_secs(5)) + .await; + Ok(()) +}