From c28770a42f9ed7bff549d7283d14172b8b061eaa Mon Sep 17 00:00:00 2001 From: jif Date: Wed, 15 Jul 2026 15:42:28 +0000 Subject: [PATCH] Respect final-answer boundaries for queued agent mail (#33367) ## Why Queue-only child-agent updates that arrive around a final answer should not restart sampling, but they must remain available to the next turn. Explicitly injected response items still need to reopen the current turn. ## What changed - Defer queue-only inter-agent mail after a final answer without treating it as same-turn pending input. - Re-enable current-turn mailbox delivery when the model requests a follow-up, a stop hook continues the turn, or a response item is explicitly injected. - Check for turn-triggering mailbox work after clearing the completed active turn so pending work can start safely. ## Testing Added coverage that queued child mail waits for the next user turn and that an injected response item after a final answer triggers a follow-up request. GitOrigin-RevId: aa81d00707943284b7ab069b60970e547908c895 --- codex-rs/core/src/session/inject.rs | 2 +- codex-rs/core/src/session/input_queue.rs | 27 +++-- codex-rs/core/src/session/turn.rs | 14 +++ codex-rs/core/src/tasks/mod.rs | 13 ++- codex-rs/core/tests/suite/pending_input.rs | 125 +++++++++++++++++++++ 5 files changed, 169 insertions(+), 12 deletions(-) diff --git a/codex-rs/core/src/session/inject.rs b/codex-rs/core/src/session/inject.rs index b0b9ef79266b..3e086964c391 100644 --- a/codex-rs/core/src/session/inject.rs +++ b/codex-rs/core/src/session/inject.rs @@ -24,7 +24,7 @@ impl Session { match active.as_mut() { Some(active_turn) => { self.input_queue - .extend_pending_input_for_turn_state( + .extend_pending_input_and_accept_mailbox_delivery_for_turn_state( active_turn.turn_state.as_ref(), input.into_iter().map(TurnInput::ResponseItem).collect(), ) diff --git a/codex-rs/core/src/session/input_queue.rs b/codex-rs/core/src/session/input_queue.rs index e46a71892be5..3992f44d4af9 100644 --- a/codex-rs/core/src/session/input_queue.rs +++ b/codex-rs/core/src/session/input_queue.rs @@ -133,7 +133,14 @@ impl InputQueue { return; }; let mut turn_state = turn_state.lock().await; - if !turn_state.pending_input.items.is_empty() { + // Explicit same-turn work still needs a follow-up. Queue-only child 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 + ) + }) { return; } turn_state.set_mailbox_delivery_phase(MailboxDeliveryPhase::NextTurn); @@ -203,10 +210,14 @@ impl InputQueue { match active.as_mut() { Some(active_turn) => { let mut turn_state = active_turn.turn_state.lock().await; - ( - turn_state.pending_input.items.split_off(0), - turn_state.accepts_mailbox_delivery_for_current_turn(), - ) + let accepts_mailbox_delivery = + turn_state.accepts_mailbox_delivery_for_current_turn(); + let pending_input = if accepts_mailbox_delivery { + turn_state.pending_input.items.split_off(0) + } else { + Vec::new() + }; + (pending_input, accepts_mailbox_delivery) } None => (Vec::new(), true), } @@ -242,12 +253,12 @@ impl InputQueue { None => (false, true), } }; - if has_turn_pending_input { - return true; - } if !accepts_mailbox_delivery { return false; } + if has_turn_pending_input { + return true; + } self.has_pending_mailbox_items().await } } diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 54e4cf070d39..6fa82b55e39b 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -302,6 +302,14 @@ pub(crate) async fn run_turn( needs_follow_up: model_needs_follow_up, last_agent_message: sampling_request_last_agent_message, } = sampling_request_output; + if model_needs_follow_up { + sess.input_queue + .accept_mailbox_delivery_for_current_turn( + &sess.active_turn, + &turn_context.sub_id, + ) + .await; + } can_drain_pending_input = true; let (has_pending_input, token_status, estimated_token_count) = async { let has_pending_input = @@ -391,6 +399,12 @@ pub(crate) async fn run_turn( hook_prompt_message, ) .await; + sess.input_queue + .accept_mailbox_delivery_for_current_turn( + &sess.active_turn, + &turn_context.sub_id, + ) + .await; stop_hook_active = true; continue; } else { diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index 3faf10436441..a89015e9d790 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -456,9 +456,13 @@ impl Session { /// /// This helper generates a fresh sub-id for the synthetic turn before delegating to the /// explicit-sub-id variant. - pub(crate) async fn maybe_start_turn_for_pending_work(self: &Arc) { - self.maybe_start_turn_for_pending_work_with_sub_id(uuid::Uuid::new_v4().to_string()) - .await; + pub(crate) fn maybe_start_turn_for_pending_work(self: &Arc) -> BoxFuture<'static, ()> { + let session = Arc::clone(self); + Box::pin(async move { + session + .maybe_start_turn_for_pending_work_with_sub_id(uuid::Uuid::new_v4().to_string()) + .await; + }) } /// Starts a regular turn with the provided sub-id when pending work should wake an idle @@ -806,6 +810,9 @@ impl Session { if let Err(err) = self.flush_rollout().await { warn!("failed to flush rollout after emitting terminal turn event: {err}"); } + if cleared_active_turn { + self.maybe_start_turn_for_pending_work().await; + } } async fn take_active_turn(&self) -> Option { diff --git a/codex-rs/core/tests/suite/pending_input.rs b/codex-rs/core/tests/suite/pending_input.rs index 2f15ff495317..d7c972d5cb68 100644 --- a/codex-rs/core/tests/suite/pending_input.rs +++ b/codex-rs/core/tests/suite/pending_input.rs @@ -712,6 +712,131 @@ async fn queued_inter_agent_mail_triggers_follow_up_after_commentary_message_ite server.shutdown().await; } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn queued_inter_agent_mail_does_not_restart_after_final_answer() { + let first_chunks = vec![ + chunk(ev_response_created("resp-1")), + chunk(ev_message_item_added("msg-1", "")), + chunk(ev_output_text_delta("first answer")), + chunk(json!({ + "type": "response.output_item.done", + "item": { + "type": "message", + "role": "assistant", + "id": "msg-1", + "content": [{"type": "output_text", "text": "first answer"}], + "phase": "final_answer", + } + })), + chunk(ev_completed("resp-1")), + ]; + + let (server, _completions) = start_streaming_sse_server(vec![ + first_chunks, + response_completed_chunks("unexpected-resp-2"), + ]) + .await; + let codex = build_codex(&server).await; + + submit_queue_only_agent_mail(&codex, "queued child update").await; + submit_user_input(&codex, "first prompt").await; + wait_for_turn_complete(&codex).await; + + let mut requests = server.requests().await; + assert_eq!(requests.len(), 1); + let request: Value = from_slice(&requests[0]).expect("parse request"); + assert!( + request["input"] + .as_array() + .expect("request input") + .iter() + .all(|item| item.get("type").and_then(Value::as_str) != Some("agent_message")) + ); + + submit_user_input(&codex, "second prompt").await; + wait_for_turn_complete(&codex).await; + + requests = server.requests().await; + assert_eq!(requests.len(), 2); + let request: Value = from_slice(&requests[1]).expect("parse request"); + let input = request["input"].as_array().expect("request input"); + let agent_message = input + .iter() + .find(|item| item.get("type").and_then(Value::as_str) == Some("agent_message")) + .expect("queued child update should be included in the next turn"); + assert_eq!( + agent_message["content"], + json!([{"type": "input_text", "text": "queued child update"}]) + ); + let user_input = message_input_texts(&request, "user") + .into_iter() + .filter(|text| text == "second prompt") + .collect::>(); + assert_eq!(user_input, vec!["second prompt"]); + + server.shutdown().await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn injected_response_item_reopens_turn_after_final_answer() { + const INITIAL_PROMPT: &str = "first prompt"; + const INJECTED_CONTEXT: &str = "late injected context"; + let (gate_completed_tx, gate_completed_rx) = oneshot::channel(); + + let first_chunks = vec![ + chunk(ev_response_created("resp-1")), + chunk(ev_message_item_added("msg-1", "")), + chunk(ev_output_text_delta("first answer")), + chunk(json!({ + "type": "response.output_item.done", + "item": { + "type": "message", + "role": "assistant", + "id": "msg-1", + "content": [{"type": "output_text", "text": "first answer"}], + "phase": "final_answer", + } + })), + // Keep the response open past an observable event so the answer boundary is established + // before the late context is injected. + chunk(ev_reasoning_item_added("reason-after-final", &["done"])), + gated_chunk( + gate_completed_rx, + vec![ + ev_reasoning_item("reason-after-final", &["done"], &[]), + ev_completed("resp-1"), + ], + ), + ]; + let (server, _completions) = + start_streaming_sse_server(vec![first_chunks, response_completed_chunks("resp-2")]).await; + let codex = build_codex(&server).await; + + submit_user_input(&codex, INITIAL_PROMPT).await; + wait_for_reasoning_item_started(&codex).await; + + assert!( + codex + .inject_if_running(vec![responses::user_message_item(INJECTED_CONTEXT)]) + .await + .is_ok() + ); + let _ = gate_completed_tx.send(()); + + wait_for_turn_complete(&codex).await; + + let requests = server.requests().await; + assert_eq!(requests.len(), 2); + let second: Value = from_slice(&requests[1]).expect("parse second request"); + let relevant_user_input = message_input_texts(&second, "user") + .into_iter() + .filter(|text| text == INITIAL_PROMPT || text == INJECTED_CONTEXT) + .collect::>(); + assert_eq!(relevant_user_input, vec![INITIAL_PROMPT, INJECTED_CONTEXT]); + + server.shutdown().await; +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn user_input_does_not_preempt_after_reasoning_item() { let (gate_reasoning_done_tx, gate_reasoning_done_rx) = oneshot::channel();