From b4b055cfc8fccc0040d13aa4304cf9ba61c2e27e Mon Sep 17 00:00:00 2001 From: dermanyang-oai Date: Wed, 23 Sep 2026 22:40:56 +0000 Subject: [PATCH] Materialize rollout persistence for pending inter-agent messages (#47691) ## What changed Call `ensure_rollout_materialized` after recording pending inter-agent communication, using the supplied persistence context. This also triggers the first execution checkpoint when an agent message starts the first turn. ## Testing Extend the first-input persistence test to cover inter-agent communication, checking that the checkpoint is requested and sampling proceeds while durable metadata is blocked. Verify that prepared context precedes the agent message and that its author and recipient reach the model. GitOrigin-RevId: 55779edb1a38fcf4d1c1cff24c09f5a810460c3d --- codex-rs/core/src/hook_runtime.rs | 1 + .../tests/suite/pending_input_persistence.rs | 85 ++++++++++++++++--- 2 files changed, 75 insertions(+), 11 deletions(-) diff --git a/codex-rs/core/src/hook_runtime.rs b/codex-rs/core/src/hook_runtime.rs index 7808d695b0..e401dae04e 100644 --- a/codex-rs/core/src/hook_runtime.rs +++ b/codex-rs/core/src/hook_runtime.rs @@ -766,6 +766,7 @@ pub(crate) async fn record_pending_input( TurnInput::InterAgentCommunication(communication) => { sess.record_inter_agent_communication(turn_context, model_info, communication) .await; + sess.ensure_rollout_materialized(persist_context).await; } } record_additional_contexts(sess, turn_context, additional_contexts).await; diff --git a/codex-rs/core/tests/suite/pending_input_persistence.rs b/codex-rs/core/tests/suite/pending_input_persistence.rs index 09261abfa7..e94db81946 100644 --- a/codex-rs/core/tests/suite/pending_input_persistence.rs +++ b/codex-rs/core/tests/suite/pending_input_persistence.rs @@ -12,10 +12,12 @@ use codex_core::TurnInputSubmission; use codex_features::Feature; use codex_history::InitialHistory; use codex_history::ResumedHistory; +use codex_protocol::AgentPath; use codex_protocol::ThreadId; use codex_protocol::mcp::ClientMcpExtensions; use codex_protocol::models::ResponseItem; use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::InterAgentCommunication; use codex_protocol::protocol::Op; use codex_protocol::protocol::ThreadHistoryMode; use codex_protocol::user_input::UserInput; @@ -64,6 +66,12 @@ enum InputKind { ToolOutput, } +#[derive(Clone, Copy)] +enum FirstInputKind { + User, + InterAgentCommunication, +} + #[derive(Debug)] struct PendingCheckpoint { thread_id: ThreadId, @@ -299,8 +307,12 @@ async fn steered_input_checkpoint_controls_next_request( Ok(()) } +#[test_case(FirstInputKind::User; "user_input")] +#[test_case(FirstInputKind::InterAgentCommunication; "inter_agent_communication")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn preparation_and_first_sampling_do_not_wait_for_durable_metadata() -> anyhow::Result<()> { +async fn preparation_and_first_sampling_do_not_wait_for_durable_metadata( + input_kind: FirstInputKind, +) -> anyhow::Result<()> { let (complete, completion) = oneshot::channel(); let (server, _completions) = start_streaming_sse_server(vec![vec![StreamingSseChunk { gate: Some(completion), @@ -332,16 +344,44 @@ async fn preparation_and_first_sampling_do_not_wait_for_durable_metadata() -> an "type": "message", "role": "developer", "content": [{"type": "input_text", "text": "prepared context"}] }))?])).await??; assert!(checkpoint_requests.try_recv().is_err()); - test.codex - .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text { - text: "first prompt".to_string(), - text_elements: Vec::new(), - }])) - .await?; + match input_kind { + FirstInputKind::User => { + assert!(matches!( + test.codex + .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text { + text: "first prompt".to_string(), + text_elements: Vec::new(), + }])) + .await?, + TurnInputSubmission::Started { .. } + )); + } + FirstInputKind::InterAgentCommunication => { + test.codex + .submit(Op::InterAgentCommunication { + communication: InterAgentCommunication::new( + AgentPath::try_from("/root/worker").expect("valid agent path"), + AgentPath::root(), + /*other_recipients*/ Vec::new(), + "first prompt".to_string(), + /*trigger_turn*/ true, + ), + start_options: Default::default(), + }) + .await?; + } + } let checkpoint = timeout(Duration::from_secs(10), checkpoint_requests.recv()) .await? .expect("first execution checkpoint"); - assert_eq!(checkpoint.context, PersistContext::TurnStart); + assert_eq!( + checkpoint.context, + match input_kind { + FirstInputKind::User => PersistContext::TurnStart, + // Mailbox work enters through the pending-input drain, even on its first turn. + FirstInputKind::InterAgentCommunication => PersistContext::SteeredUserInput, + } + ); timeout( Duration::from_secs(10), server.wait_for_request_count(/*count*/ 1), @@ -361,14 +401,37 @@ async fn preparation_and_first_sampling_do_not_wait_for_durable_metadata() -> an .flatten() .filter_map(move |content| { let text = content["text"].as_str()?; - matches!(text, "prepared context" | "first prompt") - .then(|| (item["role"].as_str().expect("message role"), text)) + matches!(text, "prepared context" | "first prompt").then(|| { + ( + item["role"] + .as_str() + .or_else(|| item["author"].as_str()) + .expect("message role or agent author"), + text, + ) + }) }) }) .collect::>(); + let first_input_author = match input_kind { + FirstInputKind::User => "user", + FirstInputKind::InterAgentCommunication => { + let communication = request["input"] + .as_array() + .expect("request input") + .iter() + .find(|item| item["type"] == "agent_message") + .expect("accepted communication should reach the model"); + assert_eq!(communication["recipient"], "/root"); + "/root/worker" + } + }; assert_eq!( prepared_input, - vec![("developer", "prepared context"), ("user", "first prompt")] + vec![ + ("developer", "prepared context"), + (first_input_author, "first prompt") + ] ); drop(checkpoint.complete); complete.send(()).expect("complete inference");