Allow agent messages to wake idle root sessions behind a feature flag (#54961)

## Why

A worker can finish after its parent has returned a final answer. Agent messages without `trigger_turn` previously needed an outstanding durable sleep to wake an idle session.

## What changed

Add the disabled-by-default `wake_on_agent_messages` feature to start a new root turn when agent mail arrives through direct delivery or the host mailbox. Reuse the previous turn context for messages without `trigger_turn`. Keep these automatic wakes disabled for interrupted sessions, Plan mode, and child agents.

## Testing

Add coverage for both delivery paths waking a completed root exactly once, child completion waking a parent after its final answer, and the interruption, Plan mode, and child-agent boundaries.

GitOrigin-RevId: f1a4fc7c22be8ca46b2173a85445126cbdd6c4b6
This commit is contained in:
alishobeiri-oai
2026-10-18 01:58:02 +00:00
committed by copyberry
parent 7b44ac1661
commit 8c076942cb
9 changed files with 459 additions and 15 deletions
+6
View File
@@ -1297,6 +1297,9 @@
"view_image": {
"type": "boolean"
},
"wake_on_agent_messages": {
"type": "boolean"
},
"web_search": {
"type": "boolean"
},
@@ -7689,6 +7692,9 @@
"view_image": {
"type": "boolean"
},
"wake_on_agent_messages": {
"type": "boolean"
},
"web_search": {
"type": "boolean"
},
+8 -2
View File
@@ -6,6 +6,7 @@ use crate::realtime_conversation::handle_speech as handle_realtime_conversation_
use crate::realtime_conversation::handle_start as handle_realtime_conversation_start;
use crate::realtime_conversation::handle_text as handle_realtime_conversation_text;
use async_channel::Receiver;
use codex_features::Feature;
use codex_otel::set_parent_from_w3c_trace_context;
use tracing::Instrument;
use tracing::debug_span;
@@ -90,7 +91,10 @@ pub async fn inter_agent_communication(
.enqueue_mailbox_communication(communication, start_options)
.await;
crate::agent_communication::emit_agent_communication_receive(&sub_id);
if trigger_turn || sess.has_outstanding_durable_sleep() {
if trigger_turn
|| sess.features.enabled(Feature::WakeOnAgentMessages)
|| sess.has_outstanding_durable_sleep()
{
sess.maybe_start_turn_for_pending_work_with_sub_id(sub_id)
.await;
}
@@ -468,7 +472,9 @@ pub(super) async fn submission_loop(
match update {
Ok(true) => {
sess.input_queue.notify_mailbox();
if sess.has_outstanding_durable_sleep() {
if sess.features.enabled(Feature::WakeOnAgentMessages)
|| sess.has_outstanding_durable_sleep()
{
sess.maybe_start_turn_for_pending_work().await;
}
}
+21 -3
View File
@@ -25,6 +25,7 @@ use crate::context::GuardianSenderMessages;
use crate::state::ActiveTurn;
use crate::state::TurnState;
use crate::tasks::RegularTask;
use codex_features::Feature;
use codex_history::CodexHarnessMetadata;
use codex_history::ResponseItemEnvelope;
use codex_history::SenderUserMessages;
@@ -708,12 +709,15 @@ impl Session {
if active_turn.is_some() {
return None;
}
let needs_new_turn = self.input_queue.has_trigger_turn_mailbox_items().await;
if !needs_new_turn && !self.has_outstanding_durable_sleep() {
let has_trigger_turn = self.input_queue.has_trigger_turn_mailbox_items().await;
if !has_trigger_turn
&& !self.has_outstanding_durable_sleep()
&& !self.should_wake_on_agent_messages().await
{
return None;
}
let active_turn = active_turn.get_or_insert_with(ActiveTurn::default);
let previous_options = if needs_new_turn {
let previous_options = if has_trigger_turn {
Default::default()
} else {
self.state
@@ -727,6 +731,20 @@ impl Session {
Some((Arc::clone(&active_turn.turn_state), previous_options))
}
async fn should_wake_on_agent_messages(&self) -> bool {
if !self.features.enabled(Feature::WakeOnAgentMessages)
|| self.is_interrupted()
|| !self.input_queue.has_pending_mailbox_items().await
{
return false;
}
let state = self.state.lock().await;
let config = &state.session_configuration;
!config.session_source.is_non_root_agent()
&& TurnStartKind::Automatic.permits_mode(config.step_settings.collaboration_mode.mode)
}
pub(crate) async fn route_realtime_text_input(
self: &Arc<Self>,
text: String,
@@ -7,6 +7,7 @@ use crate::session::turn_context::TurnContext;
use crate::state::TaskKind;
use crate::tasks::SessionTask;
use crate::tasks::SessionTaskResult;
use codex_features::Feature;
use codex_protocol::AgentPath;
use codex_protocol::config_types::ApprovalsReviewer;
use codex_protocol::config_types::CollaborationMode;
@@ -852,6 +853,51 @@ async fn start_only_rejects_pending_trigger_turn_without_injecting() {
assert_eq!(session.collaboration_mode().await.mode, ModeKind::Plan);
}
/// Queue-only agent mail respects stop, Plan mode, and inherited child configuration.
#[test_case(SessionSource::Exec, true, ModeKind::Default; "interrupted_root")]
#[test_case(SessionSource::Exec, false, ModeKind::Plan; "plan_root")]
#[test_case(SessionSource::SubAgent(SubAgentSource::Review), false, ModeKind::Default; "inherited_child_flag")]
#[tokio::test]
async fn agent_message_wake_reservation_preserves_existing_boundaries(
source: SessionSource,
interrupted: bool,
mode: ModeKind,
) {
let (mut session, _turn_context) = make_session_and_context().await;
session
.features
.enable(Feature::WakeOnAgentMessages)
.expect("enable agent-message wakes");
{
let mut state = session.state.lock().await;
state.session_configuration.session_source = source;
Arc::make_mut(&mut state.session_configuration.step_settings)
.collaboration_mode
.mode = mode;
}
if interrupted {
session.mark_interrupted();
}
let communication = InterAgentCommunication::new(
AgentPath::try_from("/root/worker").expect("valid worker path"),
AgentPath::root(),
Vec::new(),
"late worker update".to_string(),
/*trigger_turn*/ false,
);
session
.input_queue
.enqueue_mailbox_communication(communication.clone(), Default::default())
.await;
assert!(session.reserve_pending_work_turn().await.is_none());
assert!(session.active_turn.lock().await.is_none());
assert_eq!(
session.input_queue.drain_mailbox_input_items().await.0,
vec![TurnInput::InterAgentCommunication(communication)]
);
}
#[tokio::test]
async fn steer_only_requires_active_turn() {
let (session, _turn_context, _rx) = make_session_and_context_with_rx().await;
+4 -5
View File
@@ -449,8 +449,8 @@ impl Session {
/// Starts a regular turn when the session is idle and pending work is waiting.
///
/// Pending work includes mailbox mail marked with `trigger_turn`, or any mailbox mail while
/// an outstanding durable sleep is attached to the thread.
/// Mail can request a new task with `trigger_turn`, wake durable sleep, or wake a root with
/// `wake_on_agent_messages` enabled.
///
/// This helper generates a fresh sub-id for the synthetic turn before delegating to the
/// explicit-sub-id variant.
@@ -466,8 +466,7 @@ impl Session {
/// Starts a regular turn with the provided sub-id when pending work should wake an idle
/// session.
///
/// The turn is created only when the session is idle and mailbox mail either requests a turn
/// or can wake an outstanding durable sleep.
/// Starts an idle turn for explicit follow-ups, durable sleep, or enabled message wakes.
pub(crate) async fn maybe_start_turn_for_pending_work_with_sub_id(
self: &Arc<Self>,
sub_id: String,
@@ -499,7 +498,7 @@ impl Session {
if !input.iter().any(
|item| matches!(item, TurnInput::InterAgentCommunication(mail) if mail.trigger_turn),
) {
// Queue-only mail wakes durable sleep without selecting a new task's settings.
// Queue-only wakes retain the previous turn context without selecting a new task.
start_options = previous_options;
start_options.cyber_access_program = self
.reference_context_item()
@@ -31,7 +31,9 @@ use codex_protocol::ThreadId;
use codex_protocol::error::CodexErr;
use codex_protocol::error::CodexErrorDetails;
use codex_protocol::error::Result as CodexResult;
use codex_protocol::protocol::InterAgentCommunication;
use codex_protocol::protocol::MultiAgentVersion;
use codex_protocol::protocol::Op;
use codex_protocol::protocol::SessionSource;
use codex_protocol::protocol::SubAgentSource;
use codex_protocol::protocol::TokenUsage;
@@ -44,6 +46,8 @@ use core_test_support::test_codex::test_codex;
use core_test_support::wait_for_event;
use futures::future::BoxFuture;
use pretty_assertions::assert_eq;
use serde_json::json;
use test_case::test_case;
use wiremock::MockServer;
struct TestAgentControl {
@@ -214,16 +218,21 @@ impl AgentControl for TestAgentControl {
async fn test_with_host_control(
server: &MockServer,
extra_features: &[Feature],
) -> anyhow::Result<(TestCodex, Arc<TestAgentControl>)> {
let thread_id = ThreadId::new();
let controller = Arc::new(TestAgentControl::new(thread_id));
let extra_features = extra_features.to_vec();
let test = test_codex()
.with_config(|config| {
.with_config(move |config| {
for feature in [
Feature::Collab,
Feature::MultiAgentV2,
Feature::AgentMessageBoard,
] {
]
.into_iter()
.chain(extra_features.iter().copied())
{
config
.features
.enable(feature)
@@ -251,7 +260,7 @@ async fn test_with_host_control(
async fn host_threads_preserve_lineage_settings_and_resume_routing() -> anyhow::Result<()> {
skip_if_no_network!(Ok(()));
let server = responses::start_mock_server().await;
let (test, controller) = test_with_host_control(&server).await?;
let (test, controller) = test_with_host_control(&server, &[]).await?;
let root_id = test.session_configured.thread_id;
let mut child_config = test.config.clone();
child_config.service_tier = None;
@@ -490,7 +499,7 @@ async fn host_factory_follows_thread_lifecycle() -> anyhow::Result<()> {
async fn collaboration_tools_dispatch_to_the_host_controller() -> anyhow::Result<()> {
skip_if_no_network!(Ok(()));
let server = responses::start_mock_server().await;
let (test, _controller) = test_with_host_control(&server).await?;
let (test, _controller) = test_with_host_control(&server, &[]).await?;
let mut events = vec![responses::ev_response_created("tools")];
for (call, name, arguments) in [
(
@@ -541,7 +550,7 @@ use codex_protocol::protocol::InternalSessionSource;
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn host_mailbox_notifies_waiting_agent_and_delivers_once() -> anyhow::Result<()> {
let server = responses::start_mock_server().await;
let (test, controller) = test_with_host_control(&server).await?;
let (test, controller) = test_with_host_control(&server, &[]).await?;
responses::mount_sse_once(
&server,
responses::sse(vec![
@@ -609,3 +618,80 @@ async fn host_mailbox_notifies_waiting_agent_and_delivers_once() -> anyhow::Resu
test.codex.shutdown_and_wait().await?;
Ok(())
}
/// Both delivery paths wake an opted-in completed root exactly once.
#[test_case(false; "direct_op")]
#[test_case(true; "host_mailbox")]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn agent_mail_wakes_completed_root_when_enabled(
via_host_mailbox: bool,
) -> anyhow::Result<()> {
let server = responses::start_mock_server().await;
let (test, controller) =
test_with_host_control(&server, &[Feature::WakeOnAgentMessages]).await?;
let mut final_answer = responses::ev_assistant_message("first-answer", "initial work finished");
final_answer["item"]["phase"] = json!("final_answer");
let response = responses::mount_sse_sequence(
&server,
vec![
responses::sse(vec![
responses::ev_response_created("initial"),
final_answer,
responses::ev_completed("initial"),
]),
responses::sse(vec![
responses::ev_response_created("agent-mail"),
responses::ev_assistant_message("mail-answer", "received the agent update"),
responses::ev_completed("agent-mail"),
]),
],
)
.await;
test.submit_text_turn("Finish the initial work.").await?;
assert_eq!(response.requests().len(), 1);
let message = InterAgentCommunication::new(
AgentPath::root().join("worker").expect("valid worker path"),
AgentPath::root(),
Vec::new(),
"Late agent update".into(),
/*trigger_turn*/ false,
);
if via_host_mailbox {
controller.mail.lock().expect("mail lock").push(message);
controller.mailbox_state.send_replace(true);
} else {
test.codex
.submit(Op::InterAgentCommunication {
communication: message,
start_options: Default::default(),
})
.await?;
}
wait_for_event(&test.codex, |event| {
matches!(event, EventMsg::TurnComplete(_))
})
.await;
test.codex.shutdown_and_wait().await?;
let requests = response.requests();
assert_eq!(requests.len(), 2);
assert_eq!(
responses::received_responses_requests(&server).await.len(),
2
);
assert_ne!(
requests[0].body_json()["client_metadata"]["turn_id"],
requests[1].body_json()["client_metadata"]["turn_id"]
);
let messages = requests[1].inputs_of_type("agent_message");
assert_eq!(messages.len(), 1);
assert_eq!(messages[0]["author"], "/root/worker");
assert_eq!(messages[0]["recipient"], "/root");
assert_eq!(
messages[0]["content"],
json!([{"type": "input_text", "text": "Late agent update"}])
);
assert!(controller.mail.lock().expect("mail lock").is_empty());
Ok(())
}
+178
View File
@@ -1167,6 +1167,184 @@ async fn astra_omits_disabled_executor_and_plugin_skills_from_model_context(
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn child_completion_wakes_parent_after_final_answer() -> Result<()> {
use core_test_support::ThreadIdle;
use core_test_support::responses;
use wiremock::matchers::header;
skip_if_no_network!(Ok(()));
let server = start_mock_server().await;
let mut extensions = ExtensionRegistryBuilder::<Config>::new();
extensions.thread_lifecycle_contributor(Arc::new(ThreadIdle));
let mut test = test_codex()
.with_model("gpt-6-astra")
.with_extensions(Arc::new(extensions.build()))
.with_config(|config| {
configure_scenario_catalog(config);
config.model_reasoning_effort = Some(ReasoningEffort::Medium);
config.include_environment_context = false;
for feature in [
Feature::Collab,
Feature::MultiAgentV2,
Feature::MultiAgentV2DynamicTools,
Feature::WakeOnAgentMessages,
] {
config.features.enable(feature).unwrap();
}
})
.build_with_auto_env(&server)
.await?;
// A pending client tool holds the real child result until the parent is idle.
test.codex.shutdown_and_wait().await?;
let root = test
.thread_manager
.start_thread(StartThreadOptions {
dynamic_tools: vec![DynamicToolSpec::Function(DynamicToolFunctionSpec {
name: "check_release".to_string(),
description: "Check whether the release is ready.".to_string(),
input_schema: json!({"type": "object", "properties": {}, "additionalProperties": false}),
defer_loading: false,
})],
environments: Some(vec![test.executor_environment().request()]),
..StartThreadOptions::new(test.config.clone())
})
.await?;
test.codex = root.thread;
test.session_configured = root.session_configured;
let parent_id = test.session_configured.thread_id.to_string();
let mut created = test.thread_manager.subscribe_thread_created();
responses::mount_sse_once_match(
&server,
header("thread-id", parent_id.clone()),
sse(vec![
ev_function_call_with_namespace(
"spawn-worker",
"collaboration",
"spawn_agent",
r#"{"task_name":"worker","fork_turns":"none","message":"Check whether the release is ready."}"#,
),
ev_completed("parent-spawn"),
]),
)
.await;
let expected_parent = parent_id.clone();
let mut final_answer =
ev_assistant_message("parent-final", "I have delegated the release check.");
final_answer["item"]["phase"] = json!("final_answer");
let parent_finished = responses::mount_sse_once_match(
&server,
move |request: &wiremock::Request| {
request
.headers
.get("thread-id")
.is_some_and(|id| id == expected_parent.as_str())
&& String::from_utf8_lossy(&request.body).contains("spawn-worker")
},
sse(vec![final_answer, ev_completed("parent-finished")]),
)
.await;
let expected_parent = parent_id.clone();
responses::mount_sse_once_match(
&server,
move |request: &wiremock::Request| {
request
.headers
.get("thread-id")
.is_some_and(|id| id != expected_parent.as_str())
},
sse(vec![
responses::ev_function_call("worker-check", "check_release", "{}"),
ev_completed("child-check"),
]),
)
.await;
test.submit_text_turn("Ask a worker to check the release, then finish your turn.")
.await?;
ThreadIdle::wait(&test.codex).await;
let child_id = created.try_recv()?;
let child = test.thread_manager.get_thread(child_id).await?;
let call = wait_for_event_match(&child, |event| match event {
EventMsg::DynamicToolCallRequest(call) => Some(call.clone()),
_ => None,
})
.await;
assert_eq!(call.call_id, "worker-check");
responses::mount_sse_once_match(
&server,
header("thread-id", child_id.to_string()),
sse(vec![
ev_assistant_message("child-final", "All release checks passed."),
ev_completed("child-finished"),
]),
)
.await;
let wake = responses::mount_sse_once_match(
&server,
header("thread-id", parent_id.clone()),
sse(vec![
ev_assistant_message(
"parent-update",
"The worker confirmed that the release is ready.",
),
ev_completed("parent-woken"),
]),
)
.await;
child
.submit(Op::DynamicToolResponse {
id: call.call_id,
response: DynamicToolResponse {
content_items: vec![DynamicToolCallOutputContentItem::InputText {
text: "All release checks passed.".to_string(),
}],
success: true,
},
})
.await?;
wait_for_event(&child, |event| matches!(event, EventMsg::TurnComplete(_))).await;
wait_for_event(&test.codex, |event| {
matches!(event, EventMsg::TurnComplete(_))
})
.await;
ThreadIdle::wait(&test.codex).await;
assert_ne!(
parent_finished.single_request().body_json()["client_metadata"]["turn_id"],
wake.single_request().body_json()["client_metadata"]["turn_id"],
);
let mut requests = responses::received_responses_requests(&server).await;
assert_eq!(requests.len(), 5);
// Preserve each complete stream while removing parent/child scheduling order.
requests
.sort_by_key(|request| request.header("thread-id").as_deref() != Some(parent_id.as_str()));
let history = context_snapshot::format_request_history_snapshot(
"The parent finishes at medium reasoning without persistent mode. A real child then completes its release check and wakes a new parent turn. Complete model streams are grouped as parent, then worker.",
&requests,
&ContextSnapshotOptions::default()
.rewrite_known_segments()
.include_request_settings(),
);
// The shared formatter labels agent messages without rendering their payload.
let mut mailbox =
responses::strip_metadata_from_json(responses::strip_response_item_ids_from_json(
serde_json::Value::Array(wake.single_request().inputs_of_type("agent_message")),
));
mailbox.sort_all_objects();
insta::assert_snapshot!(format!(
"{history}\n\nWoken parent mailbox:\n{}",
serde_json::to_string_pretty(&mailbox)?,
));
test.thread_manager
.shutdown_all_threads_bounded(Duration::from_secs(5))
.await;
Ok(())
}
#[test_case::test_case(false; "native_tools")]
#[test_case::test_case(true; "host_tools")]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
@@ -0,0 +1,97 @@
---
source: core/tests/suite/scenarios.rs
expression: "format!(\"{history}\\n\\nWoken parent mailbox:\\n{}\",\nserde_json::to_string_pretty(&mailbox)?,)"
---
Scenario: The parent finishes at medium reasoning without persistent mode. A real child then completes its release check and wakes a new parent turn. Complete model streams are grouped as parent, then worker.
## Window 1
Settings:
include: ["reasoning.encrypted_content"]
model: "gpt-6-astra"
parallel_tool_calls: false
prompt_cache_key: "<PROMPT_CACHE_KEY 1>"
reasoning: {"context":"all_turns","effort":"medium"}
store: false
stream: true
text: {"verbosity":"low"}
tool_choice: "auto"
-- request 1 (turn) --
00:additional_tools/developer (3; hash=37af7b347b5c1eb7):
- namespace/functions; hash=a3c60999d0288751
- custom/exec
- function/wait
- function/request_user_input
- function/request_user_input_async
- namespace/clock: Tools for reading and waiting on time.; hash=f402e8c5e9b5e317
- function/sleep
- namespace/collaboration: Tools for spawning and managing sub-agents.; hash=2c8d6295978dbe63
- function/followup_task
- function/interrupt_agent
- function/list_agents
- function/send_message
- function/spawn_agent
- function/wait_agent
01:message/developer:
<MODEL_INSTRUCTIONS>
02:message/developer[2]:
[01] <PERMISSIONS_INSTRUCTIONS>
[02] <COLLABORATION_MODE>
03:message/developer:
<MULTI_AGENT_ROLE>
04:message/developer:
<MULTI_AGENT_MODE>
05:message/user:
Ask a worker to check the release, then finish your turn.
-- request 2 (turn) --
06:function_call/collaboration.spawn_agent:{"fork_turns":"none","message":"Check whether the release is ready.","task_name":"worker"}
07:function_call_output:{"task_name":"/root/worker"}
-- request 3 (turn) --
08:message/assistant:
I have delegated the release check.
09:agent_message
## Window 2 (after request 3: input diverged at item 00)
Settings: same as window 1
-- request 4 (turn) --
00:additional_tools/developer (3; hash=ce2de39d992cd930):
- namespace/functions; hash=fa0731964313a476
- custom/exec
- function/wait
- function/request_user_input
- namespace/clock: Tools for reading and waiting on time.; hash=f402e8c5e9b5e317
- function/sleep
- namespace/collaboration: Tools for spawning and managing sub-agents.; hash=2c8d6295978dbe63
- function/followup_task
- function/interrupt_agent
- function/list_agents
- function/send_message
- function/spawn_agent
- function/wait_agent
01:message/developer:
<MODEL_INSTRUCTIONS>
02:message/developer[2]:
[01] <PERMISSIONS_INSTRUCTIONS>
[02] <COLLABORATION_MODE>
03:message/developer:
<MULTI_AGENT_ROLE>
04:message/developer:
<MULTI_AGENT_MODE>
05:agent_message
-- request 5 (turn) --
06:function_call/check_release:{}
07:function_call_output:All release checks passed.
Woken parent mailbox:
[
{
"author": "/root/worker",
"content": [
{
"text": "Message Type: FINAL_ANSWER\nTask name: /root\nSender: /root/worker\nPayload:\nAll release checks passed.",
"type": "input_text"
}
],
"recipient": "/root",
"type": "agent_message"
}
]
+8
View File
@@ -237,6 +237,8 @@ pub enum Feature {
/// Keep sampling through reasoning and commentary boundaries when agent mail arrives.
/// Pending mail is delivered at the next normal input boundary instead.
DeferMailboxPreemption,
/// Allow agent messages to wake idle root sessions after a completed turn.
WakeOnAgentMessages,
/// Preempt responses and yield foreground code-mode observations on new user input.
InstantInterrupt,
/// Enable shared discussion tools for an agent tree.
@@ -1522,6 +1524,12 @@ pub const FEATURES: &[FeatureSpec] = &[
stage: Stage::UnderDevelopment,
default_enabled: false,
},
FeatureSpec {
id: Feature::WakeOnAgentMessages,
key: "wake_on_agent_messages",
stage: Stage::UnderDevelopment,
default_enabled: false,
},
FeatureSpec {
id: Feature::AgentMessageBoard,
key: "agent_message_board",