From 766d2377a8c0d31b080f06b680700a1342fd0fc0 Mon Sep 17 00:00:00 2001 From: jif Date: Thu, 24 Sep 2026 18:28:01 +0000 Subject: [PATCH] Add an opt-in flag to defer mailbox preemption (#47913) ## Why Agent messages arriving during a response can interrupt sampling at reasoning or commentary boundaries, cutting off remaining tool calls. ## What changed Add the disabled-by-default `defer_mailbox_preemption` feature flag. When enabled, sampling continues past these boundaries and queued agent messages reach the next model request alongside completed tool results. ## Testing Add streaming scenarios and snapshots for reasoning and commentary boundaries with deferral enabled and disabled. Verify that deferral preserves the planned tool call and that both queued messages arrive in order in the next request. GitOrigin-RevId: d6846dd73b3afb47e7293f3b1aa2abcd04e37853 --- codex-rs/core/config.schema.json | 6 + codex-rs/core/src/session/turn.rs | 11 +- codex-rs/core/tests/suite/scenarios.rs | 3 + .../scenarios_mailbox_preemption_tests.rs | 180 ++++++++++++++++++ ..._preemption_commentary_deferred_false.snap | 19 ++ ...x_preemption_commentary_deferred_true.snap | 21 ++ ...x_preemption_reasoning_deferred_false.snap | 18 ++ ...ox_preemption_reasoning_deferred_true.snap | 20 ++ codex-rs/features/src/lib.rs | 9 + 9 files changed, 285 insertions(+), 2 deletions(-) create mode 100644 codex-rs/core/tests/suite/scenarios_mailbox_preemption_tests.rs create mode 100644 codex-rs/core/tests/suite/snapshots/all__suite__scenarios__mailbox_preemption__mailbox_preemption_commentary_deferred_false.snap create mode 100644 codex-rs/core/tests/suite/snapshots/all__suite__scenarios__mailbox_preemption__mailbox_preemption_commentary_deferred_true.snap create mode 100644 codex-rs/core/tests/suite/snapshots/all__suite__scenarios__mailbox_preemption__mailbox_preemption_reasoning_deferred_false.snap create mode 100644 codex-rs/core/tests/suite/snapshots/all__suite__scenarios__mailbox_preemption__mailbox_preemption_reasoning_deferred_true.snap diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index 89ef425034..8aa2114281 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -789,6 +789,9 @@ "default_mode_request_user_input": { "type": "boolean" }, + "defer_mailbox_preemption": { + "type": "boolean" + }, "deferred_executor": { "type": "boolean" }, @@ -6744,6 +6747,9 @@ "default_mode_request_user_input": { "type": "boolean" }, + "defer_mailbox_preemption": { + "type": "boolean" + }, "deferred_executor": { "type": "boolean" }, diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index cdb61e7597..77af477646 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -2707,8 +2707,15 @@ async fn try_run_sampling_request( last_agent_message = Some(agent_message); } needs_follow_up |= output_result.needs_follow_up; - // todo: remove before stabilizing multi-agent v2 - if preempt_for_mailbox_mail && sess.input_queue.has_pending_mailbox_items().await { + // Hosts can keep the current response intact and deliver mail before the next + // model request instead of cutting off its remaining tool calls. + if preempt_for_mailbox_mail + && !turn_context + .config + .features + .enabled(Feature::DeferMailboxPreemption) + && sess.input_queue.has_pending_mailbox_items().await + { break Ok(SamplingRequestResult { needs_follow_up: true, last_agent_message, diff --git a/codex-rs/core/tests/suite/scenarios.rs b/codex-rs/core/tests/suite/scenarios.rs index ed3d675686..c113311288 100644 --- a/codex-rs/core/tests/suite/scenarios.rs +++ b/codex-rs/core/tests/suite/scenarios.rs @@ -82,6 +82,9 @@ const ONE_PIXEL_PNG_BASE64: &str = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJ #[path = "scenarios_agent_message_board.rs"] mod agent_message_board; +#[path = "scenarios_mailbox_preemption_tests.rs"] +mod mailbox_preemption; + #[path = "scenarios_guardian_extra_policy.rs"] mod guardian_extra_policy; diff --git a/codex-rs/core/tests/suite/scenarios_mailbox_preemption_tests.rs b/codex-rs/core/tests/suite/scenarios_mailbox_preemption_tests.rs new file mode 100644 index 0000000000..3fa281bff9 --- /dev/null +++ b/codex-rs/core/tests/suite/scenarios_mailbox_preemption_tests.rs @@ -0,0 +1,180 @@ +//! Mailbox scheduling preserves the current response when deferral is enabled. +//! Queued mail still reaches the next request, alongside completed tool results. + +use codex_core::TurnInputRequest; +use codex_features::Feature; +use codex_protocol::AgentPath; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::InterAgentCommunication; +use codex_protocol::protocol::Op; +use core_test_support::context_snapshot; +use core_test_support::context_snapshot::ContextSnapshotOptions; +use core_test_support::context_snapshot::SnapshotEntry; +use core_test_support::responses; +use core_test_support::streaming_sse::StreamingSseChunk; +use core_test_support::streaming_sse::start_streaming_sse_server; +use core_test_support::test_codex::test_codex; +use core_test_support::wait_for_event; +use pretty_assertions::assert_eq; +use serde_json::Value; +use serde_json::json; +use test_case::test_case; +use tokio::sync::oneshot; + +#[derive(Clone, Copy)] +enum Boundary { + Reasoning, + Commentary, +} + +#[test_case(Boundary::Reasoning, false; "reasoning_default")] +#[test_case(Boundary::Commentary, false; "commentary_default")] +#[test_case(Boundary::Reasoning, true; "reasoning_deferred")] +#[test_case(Boundary::Commentary, true; "commentary_deferred")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn mailbox_preemption_preserves_response_when_deferred( + boundary: Boundary, + defer_mailbox_preemption: bool, +) -> anyhow::Result<()> { + let (release, gate) = oneshot::channel(); + let (added, done, boundary_name) = match boundary { + Boundary::Reasoning => ( + responses::ev_reasoning_item_added("boundary", &["Preparing the next action"]), + responses::ev_reasoning_item("boundary", &["Preparing the next action"], &[]), + "reasoning", + ), + Boundary::Commentary => { + let mut done = responses::ev_assistant_message("boundary", "I will update the plan."); + done["item"]["phase"] = json!("commentary"); + ( + responses::ev_message_item_added("boundary", ""), + done, + "commentary", + ) + } + }; + let (streaming, _completions) = start_streaming_sse_server(vec![ + vec![ + StreamingSseChunk { + gate: None, + body: responses::sse(vec![responses::ev_response_created("first"), added]), + }, + StreamingSseChunk { + gate: Some(gate), + body: responses::sse(vec![ + done, + responses::ev_function_call( + "planned-action", + "update_plan", + r#"{"plan":[{"step":"Reply to the user","status":"in_progress"}]}"#, + ), + responses::ev_completed("first"), + ]), + }, + ], + vec![StreamingSseChunk { + gate: None, + body: responses::sse(vec![ + responses::ev_response_created("follow-up"), + responses::ev_assistant_message("answer", "The worker update is now included."), + responses::ev_completed("follow-up"), + ]), + }], + ]) + .await; + let config_server = responses::start_mock_server().await; + let base_url = format!("{}/v1", streaming.uri()); + let test = test_codex() + .with_model("gpt-5.4") + .with_config(move |config| { + config.model_provider.base_url = Some(base_url); + config.update_plan_enabled = true; + config + .features + .set_enabled(Feature::DeferMailboxPreemption, defer_mailbox_preemption) + .expect("set mailbox policy"); + // The gated server captures uncompressed request bodies for assertions. + config + .features + .disable(Feature::EnableRequestCompression) + .expect("disable request compression"); + }) + .build_with_auto_env(&config_server) + .await?; + test.codex + .start_or_steer_turn(TurnInputRequest::user_input(vec![super::text( + "Plan the next action while the worker finishes.", + )])) + .await?; + wait_for_event( + &test.codex, + |event| matches!(event, EventMsg::ItemStarted(item) if item.item.id() == "boundary"), + ) + .await; + for message in ["Worker found the result.", "Worker checked the result."] { + test.codex + .submit(Op::InterAgentCommunication { + communication: InterAgentCommunication::new( + AgentPath::root().join("worker").expect("worker path"), + AgentPath::root(), + Vec::new(), + message.to_string(), + /*trigger_turn*/ false, + ), + start_options: Default::default(), + }) + .await?; + } + // The barrier confirms that both messages were queued before finishing the item. + test.codex + .submit(Op::RealtimeConversationListVoices) + .await?; + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::RealtimeConversationListVoicesResponse(_)) + }) + .await; + release.send(()).expect("release response"); + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + + let requests = streaming + .requests() + .await + .iter() + .map(|body| serde_json::from_slice::(body)) + .collect::>>()?; + assert_eq!(requests.len(), 2); + let input = requests[1]["input"].as_array().expect("follow-up input"); + let output = input + .iter() + .find(|item| item["type"] == "function_call_output" && item["call_id"] == "planned-action"); + assert_eq!( + output.map(|item| &item["output"]), + defer_mailbox_preemption.then_some(&json!("Plan updated")), + ); + let messages = input + .iter() + .filter(|item| item["type"] == "agent_message") + .map(|item| item["content"].clone()) + .collect::>(); + assert_eq!( + messages, + vec![ + json!([{"type": "input_text", "text": "Worker found the result."}]), + json!([{"type": "input_text", "text": "Worker checked the result."}]), + ], + ); + let entries = requests.iter().map(SnapshotEntry::body).collect::>(); + insta::assert_snapshot!( + format!("mailbox_preemption_{boundary_name}_deferred_{defer_mailbox_preemption}"), + context_snapshot::format_context_snapshot( + "Worker messages arrive during a response. Deferred mailbox preemption lets the planned tool execute before the next request receives both messages.", + &entries, + &ContextSnapshotOptions::default().rewrite_known_segments(), + ) + ); + streaming.shutdown().await; + Ok(()) +} diff --git a/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__mailbox_preemption__mailbox_preemption_commentary_deferred_false.snap b/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__mailbox_preemption__mailbox_preemption_commentary_deferred_false.snap new file mode 100644 index 0000000000..b1782e8825 --- /dev/null +++ b/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__mailbox_preemption__mailbox_preemption_commentary_deferred_false.snap @@ -0,0 +1,19 @@ +--- +source: core/tests/suite/scenarios_mailbox_preemption_tests.rs +expression: "context_snapshot::format_context_snapshot(\"Worker messages arrive during a response. Deferred mailbox preemption lets the planned tool execute before the next request receives both messages.\",\n&entries, &ContextSnapshotOptions::default().rewrite_known_segments(),)" +--- +Scenario: Worker messages arrive during a response. Deferred mailbox preemption lets the planned tool execute before the next request receives both messages. + +## Window 1 +-- request 1 (request) -- +00:message/developer: + +01:message/user: + +02:message/user: + Plan the next action while the worker finishes. +-- request 2 (request) -- +03:message/assistant: + I will update the plan. +04:agent_message +05:agent_message diff --git a/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__mailbox_preemption__mailbox_preemption_commentary_deferred_true.snap b/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__mailbox_preemption__mailbox_preemption_commentary_deferred_true.snap new file mode 100644 index 0000000000..4e68050228 --- /dev/null +++ b/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__mailbox_preemption__mailbox_preemption_commentary_deferred_true.snap @@ -0,0 +1,21 @@ +--- +source: core/tests/suite/scenarios_mailbox_preemption_tests.rs +expression: "context_snapshot::format_context_snapshot(\"Worker messages arrive during a response. Deferred mailbox preemption lets the planned tool execute before the next request receives both messages.\",\n&entries, &ContextSnapshotOptions::default().rewrite_known_segments(),)" +--- +Scenario: Worker messages arrive during a response. Deferred mailbox preemption lets the planned tool execute before the next request receives both messages. + +## Window 1 +-- request 1 (request) -- +00:message/developer: + +01:message/user: + +02:message/user: + Plan the next action while the worker finishes. +-- request 2 (request) -- +03:message/assistant: + I will update the plan. +04:function_call/update_plan:{"plan":[{"status":"in_progress","step":"Reply to the user"}]} +05:function_call_output:Plan updated +06:agent_message +07:agent_message diff --git a/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__mailbox_preemption__mailbox_preemption_reasoning_deferred_false.snap b/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__mailbox_preemption__mailbox_preemption_reasoning_deferred_false.snap new file mode 100644 index 0000000000..e27e94fc74 --- /dev/null +++ b/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__mailbox_preemption__mailbox_preemption_reasoning_deferred_false.snap @@ -0,0 +1,18 @@ +--- +source: core/tests/suite/scenarios_mailbox_preemption_tests.rs +expression: "context_snapshot::format_context_snapshot(\"Worker messages arrive during a response. Deferred mailbox preemption lets the planned tool execute before the next request receives both messages.\",\n&entries, &ContextSnapshotOptions::default().rewrite_known_segments(),)" +--- +Scenario: Worker messages arrive during a response. Deferred mailbox preemption lets the planned tool execute before the next request receives both messages. + +## Window 1 +-- request 1 (request) -- +00:message/developer: + +01:message/user: + +02:message/user: + Plan the next action while the worker finishes. +-- request 2 (request) -- +03:reasoning:summary=Preparing the next action:encrypted=true +04:agent_message +05:agent_message diff --git a/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__mailbox_preemption__mailbox_preemption_reasoning_deferred_true.snap b/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__mailbox_preemption__mailbox_preemption_reasoning_deferred_true.snap new file mode 100644 index 0000000000..2b3a2dab74 --- /dev/null +++ b/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__mailbox_preemption__mailbox_preemption_reasoning_deferred_true.snap @@ -0,0 +1,20 @@ +--- +source: core/tests/suite/scenarios_mailbox_preemption_tests.rs +expression: "context_snapshot::format_context_snapshot(\"Worker messages arrive during a response. Deferred mailbox preemption lets the planned tool execute before the next request receives both messages.\",\n&entries, &ContextSnapshotOptions::default().rewrite_known_segments(),)" +--- +Scenario: Worker messages arrive during a response. Deferred mailbox preemption lets the planned tool execute before the next request receives both messages. + +## Window 1 +-- request 1 (request) -- +00:message/developer: + +01:message/user: + +02:message/user: + Plan the next action while the worker finishes. +-- request 2 (request) -- +03:reasoning:summary=Preparing the next action:encrypted=true +04:function_call/update_plan:{"plan":[{"status":"in_progress","step":"Reply to the user"}]} +05:function_call_output:Plan updated +06:agent_message +07:agent_message diff --git a/codex-rs/features/src/lib.rs b/codex-rs/features/src/lib.rs index c27189e043..171e6dc2ea 100644 --- a/codex-rs/features/src/lib.rs +++ b/codex-rs/features/src/lib.rs @@ -205,6 +205,9 @@ pub enum Feature { Collab, /// Enable task-path-based multi-agent routing. MultiAgentV2, + /// Keep sampling through reasoning and commentary boundaries when agent mail arrives. + /// Pending mail is delivered at the next normal input boundary instead. + DeferMailboxPreemption, /// Enable shared discussion tools for an agent tree. AgentMessageBoard, /// Removed compatibility flag retained as a no-op. @@ -1317,6 +1320,12 @@ pub const FEATURES: &[FeatureSpec] = &[ stage: Stage::Stable, default_enabled: false, }, + FeatureSpec { + id: Feature::DeferMailboxPreemption, + key: "defer_mailbox_preemption", + stage: Stage::UnderDevelopment, + default_enabled: false, + }, FeatureSpec { id: Feature::AgentMessageBoard, key: "agent_message_board",