mirror of
https://github.com/openai/codex.git
synced 2026-10-20 00:18:49 +00:00
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
This commit is contained in:
@@ -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<DeliveryReceipt>>;
|
||||
|
||||
/// 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<codex_protocol::protocol::InterAgentCommunication>;
|
||||
|
||||
/// 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<bool>;
|
||||
|
||||
/// 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.
|
||||
|
||||
@@ -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<String> {
|
||||
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(),
|
||||
|
||||
@@ -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<codex_protocol::protocol::InterAgentCommunication> {
|
||||
self.runtime.mailboxes.take(agent)
|
||||
}
|
||||
|
||||
fn watch_mailbox(&self, agent: ThreadId) -> tokio::sync::watch::Receiver<bool> {
|
||||
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?;
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
|
||||
@@ -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<Option<HashMap<ThreadId, Mailbox>>>);
|
||||
|
||||
impl Default for Mailboxes {
|
||||
fn default() -> Self {
|
||||
Self(Mutex::new(Some(HashMap::new())))
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct Mailbox {
|
||||
pending: Vec<StoredMail>,
|
||||
state: watch::Sender<bool>,
|
||||
}
|
||||
|
||||
struct StoredMail {
|
||||
id: Option<String>,
|
||||
communication: InterAgentCommunication,
|
||||
_diagnostics_guard: GaugeGuard,
|
||||
}
|
||||
|
||||
impl Mailboxes {
|
||||
pub(super) fn enqueue(
|
||||
&self,
|
||||
agent: ThreadId,
|
||||
id: Option<String>,
|
||||
messages: Vec<InterAgentCommunication>,
|
||||
) -> 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<InterAgentCommunication> {
|
||||
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<bool> {
|
||||
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);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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)]
|
||||
|
||||
@@ -109,6 +109,7 @@ pub(crate) struct LocalAgentRuntime {
|
||||
Arc<OnceLock<Arc<dyn ThreadInstructionsProvider>>>,
|
||||
pub(super) registry: Arc<AgentRegistry>,
|
||||
pub(super) residency: Arc<V2Residency>,
|
||||
pub(super) mailboxes: Arc<super::mailbox::Mailboxes>,
|
||||
/// Shared by every session in this tree, including private delegates.
|
||||
pub(crate) shutdown: CancellationToken,
|
||||
shutdown_state: Arc<AgentTreeShutdownState>,
|
||||
@@ -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<AgentTreeShutdownState> {
|
||||
self.shutdown_state.members.close();
|
||||
self.shutdown.cancel();
|
||||
self.mailboxes.close();
|
||||
Arc::clone(&self.shutdown_state)
|
||||
}
|
||||
|
||||
|
||||
@@ -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)]
|
||||
|
||||
@@ -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");
|
||||
|
||||
@@ -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<InputQueueActivity>,
|
||||
controller: Option<(ThreadId, Arc<dyn AgentControl>, watch::Receiver<bool>)>,
|
||||
mailbox_pending_mails: Mutex<VecDeque<PendingMailboxCommunication>>,
|
||||
}
|
||||
|
||||
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<dyn AgentControl>) -> Self {
|
||||
let updates = control.watch_mailbox(thread_id);
|
||||
Self {
|
||||
controller: Some((thread_id, control, updates)),
|
||||
..Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn mailbox_updates(&self) -> Option<watch::Receiver<bool>> {
|
||||
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<PendingMailboxCommunication>) {
|
||||
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<TurnState>>,
|
||||
@@ -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<PendingMailboxCommunication> {
|
||||
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<TurnInput>, TurnStartOptions) {
|
||||
let pending_mails = self
|
||||
.mailbox_pending_mails
|
||||
.lock()
|
||||
.await
|
||||
.drain(..)
|
||||
.collect::<Vec<_>>();
|
||||
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()
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -49,6 +49,8 @@ use wiremock::MockServer;
|
||||
struct TestAgentControl {
|
||||
thread_id: ThreadId,
|
||||
service_tier: Mutex<Option<String>>,
|
||||
mail: Mutex<Vec<codex_protocol::protocol::InterAgentCommunication>>,
|
||||
mailbox_state: tokio::sync::watch::Sender<bool>,
|
||||
agents: Mutex<HashMap<String, ThreadId>>,
|
||||
}
|
||||
|
||||
@@ -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<codex_protocol::protocol::InterAgentCommunication> {
|
||||
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<bool> {
|
||||
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::<serde_json::Value>(&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(())
|
||||
}
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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::<Config>::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::<Vec<_>>();
|
||||
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(())
|
||||
}
|
||||
Reference in New Issue
Block a user