Add an opt-in flag to defer mailbox preemption (#47913)

## Why

Agent messages arriving during a response can interrupt sampling at reasoning or commentary boundaries, cutting off remaining tool calls.

## What changed

Add the disabled-by-default `defer_mailbox_preemption` feature flag. When enabled, sampling continues past these boundaries and queued agent messages reach the next model request alongside completed tool results.

## Testing

Add streaming scenarios and snapshots for reasoning and commentary boundaries with deferral enabled and disabled. Verify that deferral preserves the planned tool call and that both queued messages arrive in order in the next request.

GitOrigin-RevId: d6846dd73b3afb47e7293f3b1aa2abcd04e37853
This commit is contained in:
jif
2026-09-24 18:38:58 +00:00
committed by copyberry
parent a8a4f14087
commit 766d2377a8
9 changed files with 285 additions and 2 deletions
+6
View File
@@ -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"
},
+9 -2
View File
@@ -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,
+3
View File
@@ -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;
@@ -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::<Value>(body))
.collect::<serde_json::Result<Vec<_>>>()?;
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::<Vec<_>>();
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::<Vec<_>>();
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(())
}
@@ -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:
<PERMISSIONS_INSTRUCTIONS>
01:message/user:
<ENVIRONMENT_CONTEXT>
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
@@ -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:
<PERMISSIONS_INSTRUCTIONS>
01:message/user:
<ENVIRONMENT_CONTEXT>
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
@@ -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:
<PERMISSIONS_INSTRUCTIONS>
01:message/user:
<ENVIRONMENT_CONTEXT>
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
@@ -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:
<PERMISSIONS_INSTRUCTIONS>
01:message/user:
<ENVIRONMENT_CONTEXT>
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
+9
View File
@@ -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",