From 617523998027d7804afa0719681e66b8d3fdeef8 Mon Sep 17 00:00:00 2001 From: jif-oai Date: Thu, 2 Apr 2026 11:20:21 +0200 Subject: [PATCH] fix: pending input race --- codex-rs/core/src/agent/mailbox.rs | 40 ++++++++++++++++++++++-------- codex-rs/core/src/agent/mod.rs | 1 + codex-rs/core/src/codex.rs | 12 ++++++--- codex-rs/core/src/codex_tests.rs | 36 +++++++++++++++++++++++++++ codex-rs/core/src/tasks/mod.rs | 19 ++++++++++++++ 5 files changed, 94 insertions(+), 14 deletions(-) diff --git a/codex-rs/core/src/agent/mailbox.rs b/codex-rs/core/src/agent/mailbox.rs index c328236475..a447794ac3 100644 --- a/codex-rs/core/src/agent/mailbox.rs +++ b/codex-rs/core/src/agent/mailbox.rs @@ -19,6 +19,11 @@ pub(crate) struct MailboxReceiver { pending_mails: VecDeque, } +pub(crate) struct MailboxDrain { + pub(crate) mails: Vec, + pub(crate) saw_trigger_turn: bool, +} + impl Mailbox { pub(crate) fn new() -> (Self, MailboxReceiver) { let (tx, rx) = mpsc::unbounded_channel(); @@ -65,9 +70,14 @@ impl MailboxReceiver { self.pending_mails.iter().any(|mail| mail.trigger_turn) } - pub(crate) fn drain(&mut self) -> Vec { + pub(crate) fn drain(&mut self) -> MailboxDrain { self.sync_pending_mails(); - self.pending_mails.drain(..).collect() + let mails = self.pending_mails.drain(..).collect::>(); + let saw_trigger_turn = mails.iter().any(|mail| mail.trigger_turn); + MailboxDrain { + mails, + saw_trigger_turn, + } } } @@ -134,28 +144,38 @@ mod tests { mailbox.send(mail_one.clone()); mailbox.send(mail_two.clone()); - assert_eq!(receiver.drain(), vec![mail_one, mail_two]); + let drain = receiver.drain(); + assert_eq!(drain.mails, vec![mail_one, mail_two]); + assert!(!drain.saw_trigger_turn); assert!(!receiver.has_pending()); } #[tokio::test] - async fn mailbox_tracks_pending_trigger_turn_mail() { + async fn mailbox_drain_tracks_pending_trigger_turn_mail() { let (mailbox, mut receiver) = Mailbox::new(); - mailbox.send(make_mail( + let queue_only_mail = make_mail( AgentPath::root(), AgentPath::try_from("/root/worker").expect("agent path"), "queued", /*trigger_turn*/ false, - )); - assert!(!receiver.has_pending_trigger_turn()); + ); + mailbox.send(queue_only_mail.clone()); - mailbox.send(make_mail( + let drain = receiver.drain(); + assert_eq!(drain.mails, vec![queue_only_mail]); + assert!(!drain.saw_trigger_turn); + + let trigger_turn_mail = make_mail( AgentPath::root(), AgentPath::try_from("/root/worker").expect("agent path"), "wake", /*trigger_turn*/ true, - )); - assert!(receiver.has_pending_trigger_turn()); + ); + mailbox.send(trigger_turn_mail.clone()); + + let drain = receiver.drain(); + assert_eq!(drain.mails, vec![trigger_turn_mail]); + assert!(drain.saw_trigger_turn); } } diff --git a/codex-rs/core/src/agent/mod.rs b/codex-rs/core/src/agent/mod.rs index a60fc3004a..4f42ca4c4f 100644 --- a/codex-rs/core/src/agent/mod.rs +++ b/codex-rs/core/src/agent/mod.rs @@ -8,6 +8,7 @@ pub(crate) mod status; pub(crate) use codex_protocol::protocol::AgentStatus; pub(crate) use control::AgentControl; pub(crate) use mailbox::Mailbox; +pub(crate) use mailbox::MailboxDrain; pub(crate) use mailbox::MailboxReceiver; pub(crate) use registry::exceeds_thread_spawn_depth_limit; pub(crate) use registry::next_thread_spawn_depth; diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 47c6cf8600..92daa81681 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -9,6 +9,7 @@ use std::sync::atomic::AtomicU64; use crate::agent::AgentControl; use crate::agent::AgentStatus; use crate::agent::Mailbox; +use crate::agent::MailboxDrain; use crate::agent::MailboxReceiver; use crate::agent::agent_status_from_event; use crate::agent::status::is_final; @@ -4052,6 +4053,10 @@ impl Session { self.mailbox.send(communication); } + pub(crate) async fn drain_mailbox(&self) -> MailboxDrain { + self.mailbox_rx.lock().await.drain() + } + pub(crate) async fn has_trigger_turn_mailbox_items(&self) -> bool { self.mailbox_rx.lock().await.has_pending_trigger_turn() } @@ -4080,9 +4085,9 @@ impl Session { } }; let mailbox_items = { - let mut mailbox_rx = self.mailbox_rx.lock().await; - mailbox_rx - .drain() + self.drain_mailbox() + .await + .mails .into_iter() .map(|mail| mail.to_response_input_item()) .collect::>() @@ -4099,7 +4104,6 @@ impl Session { } /// Queue response items to be injected into the next active turn created for this session. - #[cfg(test)] pub(crate) async fn queue_response_items_for_next_turn(&self, items: Vec) { if items.is_empty() { return; diff --git a/codex-rs/core/src/codex_tests.rs b/codex-rs/core/src/codex_tests.rs index 24c24eb935..5f8ba2c96d 100644 --- a/codex-rs/core/src/codex_tests.rs +++ b/codex-rs/core/src/codex_tests.rs @@ -54,6 +54,7 @@ use codex_execpolicy::NetworkRuleProtocol; use codex_execpolicy::Policy; use codex_network_proxy::NetworkProxyConfig; use codex_otel::TelemetryAuthMode; +use codex_protocol::AgentPath; use codex_protocol::config_types::CollaborationMode; use codex_protocol::config_types::ModeKind; use codex_protocol::config_types::Settings; @@ -69,6 +70,7 @@ use codex_protocol::protocol::ConversationAudioParams; use codex_protocol::protocol::CreditsSnapshot; use codex_protocol::protocol::GranularApprovalConfig; use codex_protocol::protocol::InitialHistory; +use codex_protocol::protocol::InterAgentCommunication; use codex_protocol::protocol::NetworkApprovalProtocol; use codex_protocol::protocol::RateLimitSnapshot; use codex_protocol::protocol::RateLimitWindow; @@ -4781,6 +4783,40 @@ async fn queued_response_items_for_next_turn_move_into_next_active_turn() { assert_eq!(sess.get_pending_input().await, vec![queued_item]); } +#[tokio::test] +async fn on_task_finished_records_queue_only_mailbox_mail_without_restart() { + let (sess, tc, _rx) = make_session_and_context_with_rx().await; + sess.spawn_task( + Arc::clone(&tc), + Vec::new(), + NeverEndingTask { + kind: TaskKind::Regular, + listen_to_cancellation_token: false, + }, + ) + .await; + + let communication = InterAgentCommunication::new( + AgentPath::try_from("/root/worker").expect("worker path should parse"), + AgentPath::root(), + Vec::new(), + "queued completion".to_string(), + /*trigger_turn*/ false, + ); + sess.enqueue_mailbox_communication(communication.clone()); + + sess.on_task_finished(Arc::clone(&tc), /*last_agent_message*/ None) + .await; + sleep(Duration::from_millis(50)).await; + + assert!(!sess.has_pending_input().await); + assert!(sess.active_turn.lock().await.is_none()); + assert_eq!( + sess.clone_history().await.raw_items(), + &[ResponseItem::from(communication.to_response_input_item())], + ); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn abort_review_task_emits_exited_then_aborted_and_records_history() { let (sess, tc, rx) = make_session_and_context_with_rx().await; diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index 74b2e471cd..65fac76664 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -356,6 +356,20 @@ impl Session { turn_tool_calls = ts.tool_calls; token_usage_at_turn_start = Some(ts.token_usage_at_turn_start.clone()); } + + let mailbox_drain = self.drain_mailbox().await; + let mut next_turn_mailbox_input = Vec::new(); + let mut queue_mail_for_next_turn = false; + for mail in mailbox_drain.mails { + let mailbox_input = mail.to_response_input_item(); + queue_mail_for_next_turn |= mail.trigger_turn; + if queue_mail_for_next_turn { + next_turn_mailbox_input.push(mailbox_input); + } else { + pending_input.push(mailbox_input); + } + } + if !pending_input.is_empty() { for pending_input_item in pending_input { match inspect_pending_input(self, &turn_context, pending_input_item).await { @@ -370,6 +384,11 @@ impl Session { } } } + if mailbox_drain.saw_trigger_turn { + self.queue_response_items_for_next_turn(next_turn_mailbox_input) + .await; + } + // Emit token usage metrics. if let Some(token_usage_at_turn_start) = token_usage_at_turn_start { // TODO(jif): drop this