diff --git a/codex-rs/core/src/session/inject.rs b/codex-rs/core/src/session/inject.rs index e8e61dc643..ae0ad1e3c4 100644 --- a/codex-rs/core/src/session/inject.rs +++ b/codex-rs/core/src/session/inject.rs @@ -45,7 +45,7 @@ impl Session { match active.as_mut().filter(|turn| turn.finishing.is_none()) { Some(active_turn) => { self.input_queue - .extend_pending_input_and_accept_mailbox_delivery_for_turn_state( + .extend_pending_input_and_notify_for_turn_state( active_turn.turn_state.as_ref(), input .into_iter() @@ -77,7 +77,7 @@ impl Session { return Err(input); } self.input_queue - .extend_pending_input_and_accept_mailbox_delivery_for_turn_state( + .extend_pending_input_and_notify_for_turn_state( active_turn.turn_state.as_ref(), input .into_iter() @@ -106,7 +106,7 @@ impl Session { let mut active = self.active_turn.lock().await; if let Some(active_turn) = active.as_mut().filter(|turn| turn.finishing.is_none()) { self.input_queue - .extend_pending_input_and_accept_mailbox_delivery_for_turn_state( + .extend_pending_input_and_notify_for_turn_state( active_turn.turn_state.as_ref(), items .into_iter() diff --git a/codex-rs/core/src/session/input_queue.rs b/codex-rs/core/src/session/input_queue.rs index a22fc2684f..de016d7ac9 100644 --- a/codex-rs/core/src/session/input_queue.rs +++ b/codex-rs/core/src/session/input_queue.rs @@ -90,6 +90,17 @@ pub(crate) enum InputQueueActivity { } impl TurnInput { + fn is_queue_only_mail(&self) -> bool { + match self { + Self::InterAgentCommunication(communication) => !communication.trigger_turn, + Self::ResponseItem(ResponseItemEnvelope { + item: ResponseItem::AgentMessage { .. }, + .. + }) => true, + Self::UserInput { .. } | Self::FunctionCallOutput(_) | Self::ResponseItem(_) => false, + } + } + fn wait_activity(&self) -> InputQueueActivity { match self { Self::InterAgentCommunication(_) @@ -247,6 +258,32 @@ impl InputQueue { .is_some_and(|(_, _, updates)| *updates.borrow()) } + /// Native mailbox backlog qualifies regardless of turn ID or delivery phase. + /// Injected mail qualifies only while its matching turn still accepts mail. + pub(crate) async fn should_preempt_sampling_for_mail( + &self, + active_turn: &Mutex>, + sub_id: &str, + ) -> bool { + if self.has_pending_mailbox_items().await { + return true; + } + let Some(turn_state) = self.turn_state_for_sub_id(active_turn, sub_id).await else { + return false; + }; + let turn_state = turn_state.lock().await; + turn_state.accepts_mailbox_delivery_for_current_turn() + && turn_state.pending_input.items.iter().any(|input| { + matches!( + input, + TurnInput::ResponseItem(ResponseItemEnvelope { + item: ResponseItem::AgentMessage { .. }, + .. + }) + ) + }) + } + pub(crate) async fn has_trigger_turn_mailbox_items(&self) -> bool { self.mailbox_pending_mails .lock() @@ -353,14 +390,14 @@ impl InputQueue { return; }; let mut turn_state = turn_state.lock().await; - // Explicit same-turn work still needs a follow-up. Queue-only child mail does not: keep + // Explicit same-turn work still needs a follow-up. Queue-only agent mail does not: keep // it pending so task completion records it for the next turn without sampling again. - if turn_state.pending_input.items.iter().any(|input| { - !matches!( - input, - TurnInput::InterAgentCommunication(communication) if !communication.trigger_turn - ) - }) { + if turn_state + .pending_input + .items + .iter() + .any(|input| !input.is_queue_only_mail()) + { return; } turn_state.set_mailbox_delivery_phase(MailboxDeliveryPhase::NextTurn); @@ -389,18 +426,21 @@ impl InputQueue { .accept_mailbox_delivery_for_current_turn(); } - pub(super) async fn extend_pending_input_and_accept_mailbox_delivery_for_turn_state( + pub(super) async fn extend_pending_input_and_notify_for_turn_state( &self, turn_state: &Mutex, input: Vec, ) { let activity = input.iter().map(TurnInput::wait_activity).max(); + let queue_only_mail = !input.is_empty() && input.iter().all(TurnInput::is_queue_only_mail); { let mut turn_state = turn_state.lock().await; turn_state.pending_input.notified_activity = turn_state.pending_input.notified_activity.max(activity); turn_state.pending_input.items.extend(input); - turn_state.accept_mailbox_delivery_for_current_turn(); + if !queue_only_mail { + turn_state.accept_mailbox_delivery_for_current_turn(); + } } self.activity_tx .send_replace(activity.unwrap_or(InputQueueActivity::Steer)); @@ -628,7 +668,7 @@ mod tests { assert_eq!(pending_activity, None); input_queue - .extend_pending_input_and_accept_mailbox_delivery_for_turn_state( + .extend_pending_input_and_notify_for_turn_state( &turn_state, vec![TurnInput::UserInput { metadata: Default::default(), @@ -651,7 +691,7 @@ mod tests { /*trigger_turn*/ true, ); input_queue - .extend_pending_input_and_accept_mailbox_delivery_for_turn_state( + .extend_pending_input_and_notify_for_turn_state( &turn_state, vec![TurnInput::InterAgentCommunication(mail.into())], ) @@ -683,7 +723,7 @@ mod tests { None ); input_queue - .extend_pending_input_and_accept_mailbox_delivery_for_turn_state( + .extend_pending_input_and_notify_for_turn_state( &turn_state, vec![passive_output.clone()], ) @@ -720,7 +760,7 @@ mod tests { Some(InputQueueActivity::Mailbox) ); input_queue - .extend_pending_input_and_accept_mailbox_delivery_for_turn_state( + .extend_pending_input_and_notify_for_turn_state( &turn_state, vec![TurnInput::UserInput { metadata: Default::default(), diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index c2f2a9b08f..74cc233167 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -3155,7 +3155,10 @@ async fn try_run_sampling_request( .config .features .enabled(Feature::DeferMailboxPreemption) - && sess.input_queue.has_pending_mailbox_items().await + && sess + .input_queue + .should_preempt_sampling_for_mail(&sess.active_turn, &turn_context.sub_id) + .await { tracing::event!( name: "codex.mailbox_preemption", diff --git a/codex-rs/core/src/session/turn_input.rs b/codex-rs/core/src/session/turn_input.rs index 7d9f0c8f0f..94f1664e12 100644 --- a/codex-rs/core/src/session/turn_input.rs +++ b/codex-rs/core/src/session/turn_input.rs @@ -970,7 +970,7 @@ impl Session { }; pending_input.push(input); self.input_queue - .extend_pending_input_and_accept_mailbox_delivery_for_turn_state( + .extend_pending_input_and_notify_for_turn_state( active_turn.turn_state.as_ref(), pending_input, ) diff --git a/codex-rs/core/tests/suite/scenarios.rs b/codex-rs/core/tests/suite/scenarios.rs index dc62ed14e4..d5daec3b38 100644 --- a/codex-rs/core/tests/suite/scenarios.rs +++ b/codex-rs/core/tests/suite/scenarios.rs @@ -111,6 +111,9 @@ use tokio::sync::oneshot; const ONE_PIXEL_PNG_BASE64: &str = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGP4z8DwHwAFAAH/iZk9HQAAAABJRU5ErkJggg=="; +#[path = "scenarios_injected_agent_mail.rs"] +mod injected_agent_mail; + #[path = "scenarios_dynamic_tools.rs"] mod dynamic_tools; diff --git a/codex-rs/core/tests/suite/scenarios_injected_agent_mail.rs b/codex-rs/core/tests/suite/scenarios_injected_agent_mail.rs new file mode 100644 index 0000000000..1590892853 --- /dev/null +++ b/codex-rs/core/tests/suite/scenarios_injected_agent_mail.rs @@ -0,0 +1,235 @@ +//! Injected agent messages share native delivery boundaries while retaining their wire payloads. + +use codex_core::InputWaitOutcome; +use codex_core::TurnInputRequest; +use codex_core::TurnInputSubmission; +use codex_features::Feature; +use codex_protocol::models::ResponseItem; +use codex_protocol::protocol::EventMsg; +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 std::time::Duration; +use tokio::sync::oneshot; + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn injected_mail_respects_delivery_boundaries() -> anyhow::Result<()> { + let (release_boundary, boundary_gate) = oneshot::channel(); + let (release_final, final_gate) = oneshot::channel(); + let (release_completion, completion_gate) = oneshot::channel(); + let chunk = |gate, events| StreamingSseChunk { + gate, + body: responses::sse(events), + }; + let mut final_answer = + responses::ev_assistant_message("answer", "The worker update is included."); + final_answer["item"]["phase"] = json!("final_answer"); + let (server, _completions) = start_streaming_sse_server(vec![ + vec![ + chunk(None, vec![ + responses::ev_response_created("first"), + responses::ev_reasoning_item_added("boundary", &["Preparing the plan"]), + ]), + chunk(Some(boundary_gate), vec![ + responses::ev_reasoning_item("boundary", &["Preparing the plan"], &[]), + responses::ev_function_call( + "stale-plan", "update_plan", + r#"{"plan":[{"step":"Act before reading the worker update","status":"in_progress"}]}"#, + ), + responses::ev_completed("first"), + ]), + ], + vec![ + chunk(None, vec![ + responses::ev_response_created("after-mail"), + responses::ev_message_item_added("answer", ""), + ]), + chunk(Some(final_gate), vec![ + final_answer, + // This event proves that final-answer handling has closed delivery. + responses::ev_reasoning_item_added("after-final", &["Finishing"]), + ]), + chunk(Some(completion_gate), vec![ + responses::ev_reasoning_item("after-final", &["Finishing"], &[]), + responses::ev_completed("after-mail"), + ]), + ], + vec![chunk(None, vec![ + responses::ev_response_created("next-task"), + responses::ev_assistant_message("next-answer", "The remaining updates are included."), + responses::ev_completed("next-task"), + ])], + ]).await; + let config_server = responses::start_mock_server().await; + let base_url = format!("{}/v1", server.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 + .disable(Feature::DeferMailboxPreemption) + .expect("enable preemption"); + config + .features + .disable(Feature::EnableRequestCompression) + .expect("capture request bodies"); + }) + .build_with_auto_env(&config_server) + .await?; + let mail = |id: &str, text: &str| -> ResponseItem { + serde_json::from_value(json!({ + "type": "agent_message", "id": id, + "author": "/root/worker", "recipient": "/root", + "content": [{"type": "input_text", "text": text}], + })) + .expect("agent message") + }; + let idle_mail = mail("agent_idle-mail", "The worker started checking."); + let active_mail: ResponseItem = serde_json::from_value(json!({ + "type": "agent_message", "id": "agent_active-mail", + "author": "/root/worker", "recipient": "/root", + "content": [ + {"type": "input_text", "text": "Use the updated worker result."}, + {"type": "encrypted_content", "encrypted_content": "opaque-worker-result"}, + ], + }))?; + let before_final_mail = mail( + "agent_before-final-mail", + "Another check finished while you answered.", + ); + let after_final_mail = mail( + "agent_after-final-mail", + "The worker finished the remaining checks.", + ); + + test.codex + .inject_response_items(vec![idle_mail.clone()]) + .await?; + assert!( + server.requests().await.is_empty(), + "idle injection must not start a turn" + ); + let TurnInputSubmission::Started { turn_id, .. } = test + .codex + .start_or_steer_turn(TurnInputRequest::user_input(vec![super::text( + "Plan the next action.", + )])) + .await? + else { + panic!("initial task should start a turn") + }; + wait_for_event( + &test.codex, + |event| matches!(event, EventMsg::ItemStarted(item) if item.item.id() == "boundary"), + ) + .await; + + let (outcome, ()) = tokio::try_join!(biased; + async { anyhow::Ok(test.codex.wait_for_input_activity(&turn_id, Duration::from_secs(10)).await) }, + async { test.codex.inject_response_items(vec![active_mail.clone()]).await?; anyhow::Ok(()) }, + )?; + assert_eq!(outcome, InputWaitOutcome::MailboxActivity); + release_boundary + .send(()) + .expect("release sampling boundary"); + wait_for_event(&test.codex, |event| { + assert!( + !matches!(event, EventMsg::PlanUpdate(_)), + "stale plan must not execute" + ); + matches!(event, EventMsg::ItemStarted(item) if item.item.id() == "answer") + }) + .await; + + // Mail already pending when the final answer finishes must not request another response. + test.codex + .inject_response_items(vec![before_final_mail.clone()]) + .await?; + release_final.send(()).expect("release final answer"); + wait_for_event( + &test.codex, + |event| matches!(event, EventMsg::ItemStarted(item) if item.item.id() == "after-final"), + ) + .await; + // A later arrival must not reopen delivery, including at a trailing reasoning boundary. + test.codex + .inject_response_items(vec![after_final_mail.clone()]) + .await?; + release_completion.send(()).expect("release completion"); + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + assert_eq!( + server.requests().await.len(), + 2, + "queued mail must wait for another task" + ); + + test.codex + .start_or_steer_turn(TurnInputRequest::user_input(vec![super::text( + "Continue with the remaining worker updates.", + )])) + .await?; + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + test.codex.shutdown_and_wait().await?; + + let requests = server + .requests() + .await + .iter() + .map(|body| serde_json::from_slice::(body)) + .collect::>>()?; + assert_eq!(requests.len(), 3); + for (request, expected) in requests.iter().zip([ + vec![idle_mail.clone()], + vec![idle_mail.clone(), active_mail.clone()], + vec![idle_mail, active_mail, before_final_mail, after_final_mail], + ]) { + let input = request["input"].as_array().expect("request input"); + let actual = input + .iter() + .filter(|item| item["type"] == "agent_message") + .cloned() + .collect::>(); + assert_eq!( + responses::strip_metadata_from_json(json!(actual)), + responses::strip_metadata_from_json(serde_json::to_value(expected)?) + ); + assert!(!input.iter().any(|item| item["call_id"] == "stale-plan")); + } + let next_input = requests[2]["input"].as_array().expect("next task input"); + let last_mail = next_input + .iter() + .rposition(|item| item["type"] == "agent_message") + .unwrap(); + let next_task = next_input + .iter() + .rposition(|item| item["role"] == "user") + .unwrap(); + assert!( + last_mail < next_task, + "deferred mail precedes the next task" + ); + let entries = requests.iter().map(SnapshotEntry::body).collect::>(); + insta::assert_snapshot!(context_snapshot::format_context_snapshot( + "Injected agent messages enter idle history, interrupt active sampling, and remain for the next task once the final answer closes delivery.", + &entries, + &ContextSnapshotOptions::default().rewrite_known_segments(), + )); + server.shutdown().await; + Ok(()) +} diff --git a/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__injected_agent_mail__injected_mail_respects_delivery_boundaries.snap b/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__injected_agent_mail__injected_mail_respects_delivery_boundaries.snap new file mode 100644 index 0000000000..4104389d2e --- /dev/null +++ b/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__injected_agent_mail__injected_mail_respects_delivery_boundaries.snap @@ -0,0 +1,28 @@ +--- +source: core/tests/suite/scenarios_injected_agent_mail.rs +expression: "context_snapshot::format_context_snapshot(\"Injected agent messages enter idle history, interrupt active sampling, and remain for the next task once the final answer closes delivery.\",\n&entries, &ContextSnapshotOptions::default().rewrite_known_segments(),)" +--- +Scenario: Injected agent messages enter idle history, interrupt active sampling, and remain for the next task once the final answer closes delivery. + +## Window 1 +-- request 1 (request) -- +00:message/developer: + +01:message/developer: + +02:message/user: + +03:agent_message +04:message/user: + Plan the next action. +-- request 2 (request) -- +05:reasoning:summary=Preparing the plan:encrypted=true +06:agent_message +-- request 3 (request) -- +07:message/assistant: + The worker update is included. +08:reasoning:summary=Finishing:encrypted=true +09:agent_message +10:agent_message +11:message/user: + Continue with the remaining worker updates.