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",