Respect delivery boundaries for injected agent messages (#58988)

## Why

Injected agent messages need to reach active sampling before stale tool calls execute, while messages queued around a final answer must wait for the next task without requesting another response.

## What changed

- Include injected `agent_message` items in mailbox preemption checks when the matching turn still accepts mail and `DeferMailboxPreemption` is disabled.
- Treat injected agent messages as queue-only mail so they do not reopen delivery after a final answer or force a follow-up response.

## Testing

Add a streaming scenario and context snapshot covering idle injection, active sampling preemption, deferred delivery around a final answer, and preservation of message payloads and ordering before the next task.

GitOrigin-RevId: 70f49cbe5074725c11e1fcb33eea05e0e92e17f7
This commit is contained in:
Won Park
2026-11-02 21:35:56 +00:00
committed by copyberry
parent 9da07658e2
commit ac21e58790
7 changed files with 327 additions and 18 deletions
+3 -3
View File
@@ -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()
+53 -13
View File
@@ -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<Option<ActiveTurn>>,
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<TurnState>,
input: Vec<TurnInput>,
) {
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(),
+4 -1
View File
@@ -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",
+1 -1
View File
@@ -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,
)
+3
View File
@@ -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;
@@ -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::<Value>(body))
.collect::<serde_json::Result<Vec<_>>>()?;
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::<Vec<_>>();
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::<Vec<_>>();
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(())
}
@@ -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:
<MODEL_INSTRUCTIONS>
01:message/developer:
<PERMISSIONS_INSTRUCTIONS>
02:message/user:
<ENVIRONMENT_CONTEXT>
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.