Skip to content

Commit 8c07694

Browse files
alishobeiri-oaicopyberry
authored andcommitted
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
1 parent 7b44ac1 commit 8c07694

9 files changed

Lines changed: 459 additions & 15 deletions

File tree

‎codex-rs/core/config.schema.json‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1297,6 +1297,9 @@
12971297
"view_image": {
12981298
"type": "boolean"
12991299
},
1300+
"wake_on_agent_messages": {
1301+
"type": "boolean"
1302+
},
13001303
"web_search": {
13011304
"type": "boolean"
13021305
},
@@ -7689,6 +7692,9 @@
76897692
"view_image": {
76907693
"type": "boolean"
76917694
},
7695+
"wake_on_agent_messages": {
7696+
"type": "boolean"
7697+
},
76927698
"web_search": {
76937699
"type": "boolean"
76947700
},

‎codex-rs/core/src/session/handlers.rs‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ use crate::realtime_conversation::handle_speech as handle_realtime_conversation_
66
use crate::realtime_conversation::handle_start as handle_realtime_conversation_start;
77
use crate::realtime_conversation::handle_text as handle_realtime_conversation_text;
88
use async_channel::Receiver;
9+
use codex_features::Feature;
910
use codex_otel::set_parent_from_w3c_trace_context;
1011
use tracing::Instrument;
1112
use tracing::debug_span;
@@ -90,7 +91,10 @@ pub async fn inter_agent_communication(
9091
.enqueue_mailbox_communication(communication, start_options)
9192
.await;
9293
crate::agent_communication::emit_agent_communication_receive(&sub_id);
93-
if trigger_turn || sess.has_outstanding_durable_sleep() {
94+
if trigger_turn
95+
|| sess.features.enabled(Feature::WakeOnAgentMessages)
96+
|| sess.has_outstanding_durable_sleep()
97+
{
9498
sess.maybe_start_turn_for_pending_work_with_sub_id(sub_id)
9599
.await;
96100
}
@@ -468,7 +472,9 @@ pub(super) async fn submission_loop(
468472
match update {
469473
Ok(true) => {
470474
sess.input_queue.notify_mailbox();
471-
if sess.has_outstanding_durable_sleep() {
475+
if sess.features.enabled(Feature::WakeOnAgentMessages)
476+
|| sess.has_outstanding_durable_sleep()
477+
{
472478
sess.maybe_start_turn_for_pending_work().await;
473479
}
474480
}

‎codex-rs/core/src/session/turn_input.rs‎

Lines changed: 21 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ use crate::context::GuardianSenderMessages;
2525
use crate::state::ActiveTurn;
2626
use crate::state::TurnState;
2727
use crate::tasks::RegularTask;
28+
use codex_features::Feature;
2829
use codex_history::CodexHarnessMetadata;
2930
use codex_history::ResponseItemEnvelope;
3031
use codex_history::SenderUserMessages;
@@ -708,12 +709,15 @@ impl Session {
708709
if active_turn.is_some() {
709710
return None;
710711
}
711-
let needs_new_turn = self.input_queue.has_trigger_turn_mailbox_items().await;
712-
if !needs_new_turn && !self.has_outstanding_durable_sleep() {
712+
let has_trigger_turn = self.input_queue.has_trigger_turn_mailbox_items().await;
713+
if !has_trigger_turn
714+
&& !self.has_outstanding_durable_sleep()
715+
&& !self.should_wake_on_agent_messages().await
716+
{
713717
return None;
714718
}
715719
let active_turn = active_turn.get_or_insert_with(ActiveTurn::default);
716-
let previous_options = if needs_new_turn {
720+
let previous_options = if has_trigger_turn {
717721
Default::default()
718722
} else {
719723
self.state
@@ -727,6 +731,20 @@ impl Session {
727731
Some((Arc::clone(&active_turn.turn_state), previous_options))
728732
}
729733

734+
async fn should_wake_on_agent_messages(&self) -> bool {
735+
if !self.features.enabled(Feature::WakeOnAgentMessages)
736+
|| self.is_interrupted()
737+
|| !self.input_queue.has_pending_mailbox_items().await
738+
{
739+
return false;
740+
}
741+
742+
let state = self.state.lock().await;
743+
let config = &state.session_configuration;
744+
!config.session_source.is_non_root_agent()
745+
&& TurnStartKind::Automatic.permits_mode(config.step_settings.collaboration_mode.mode)
746+
}
747+
730748
pub(crate) async fn route_realtime_text_input(
731749
self: &Arc<Self>,
732750
text: String,

‎codex-rs/core/src/session/turn_input_tests.rs‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ use crate::session::turn_context::TurnContext;
77
use crate::state::TaskKind;
88
use crate::tasks::SessionTask;
99
use crate::tasks::SessionTaskResult;
10+
use codex_features::Feature;
1011
use codex_protocol::AgentPath;
1112
use codex_protocol::config_types::ApprovalsReviewer;
1213
use codex_protocol::config_types::CollaborationMode;
@@ -852,6 +853,51 @@ async fn start_only_rejects_pending_trigger_turn_without_injecting() {
852853
assert_eq!(session.collaboration_mode().await.mode, ModeKind::Plan);
853854
}
854855

856+
/// Queue-only agent mail respects stop, Plan mode, and inherited child configuration.
857+
#[test_case(SessionSource::Exec, true, ModeKind::Default; "interrupted_root")]
858+
#[test_case(SessionSource::Exec, false, ModeKind::Plan; "plan_root")]
859+
#[test_case(SessionSource::SubAgent(SubAgentSource::Review), false, ModeKind::Default; "inherited_child_flag")]
860+
#[tokio::test]
861+
async fn agent_message_wake_reservation_preserves_existing_boundaries(
862+
source: SessionSource,
863+
interrupted: bool,
864+
mode: ModeKind,
865+
) {
866+
let (mut session, _turn_context) = make_session_and_context().await;
867+
session
868+
.features
869+
.enable(Feature::WakeOnAgentMessages)
870+
.expect("enable agent-message wakes");
871+
{
872+
let mut state = session.state.lock().await;
873+
state.session_configuration.session_source = source;
874+
Arc::make_mut(&mut state.session_configuration.step_settings)
875+
.collaboration_mode
876+
.mode = mode;
877+
}
878+
if interrupted {
879+
session.mark_interrupted();
880+
}
881+
let communication = InterAgentCommunication::new(
882+
AgentPath::try_from("/root/worker").expect("valid worker path"),
883+
AgentPath::root(),
884+
Vec::new(),
885+
"late worker update".to_string(),
886+
/*trigger_turn*/ false,
887+
);
888+
session
889+
.input_queue
890+
.enqueue_mailbox_communication(communication.clone(), Default::default())
891+
.await;
892+
893+
assert!(session.reserve_pending_work_turn().await.is_none());
894+
assert!(session.active_turn.lock().await.is_none());
895+
assert_eq!(
896+
session.input_queue.drain_mailbox_input_items().await.0,
897+
vec![TurnInput::InterAgentCommunication(communication)]
898+
);
899+
}
900+
855901
#[tokio::test]
856902
async fn steer_only_requires_active_turn() {
857903
let (session, _turn_context, _rx) = make_session_and_context_with_rx().await;

‎codex-rs/core/src/tasks/mod.rs‎

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -449,8 +449,8 @@ impl Session {
449449

450450
/// Starts a regular turn when the session is idle and pending work is waiting.
451451
///
452-
/// Pending work includes mailbox mail marked with `trigger_turn`, or any mailbox mail while
453-
/// an outstanding durable sleep is attached to the thread.
452+
/// Mail can request a new task with `trigger_turn`, wake durable sleep, or wake a root with
453+
/// `wake_on_agent_messages` enabled.
454454
///
455455
/// This helper generates a fresh sub-id for the synthetic turn before delegating to the
456456
/// explicit-sub-id variant.
@@ -466,8 +466,7 @@ impl Session {
466466
/// Starts a regular turn with the provided sub-id when pending work should wake an idle
467467
/// session.
468468
///
469-
/// The turn is created only when the session is idle and mailbox mail either requests a turn
470-
/// or can wake an outstanding durable sleep.
469+
/// Starts an idle turn for explicit follow-ups, durable sleep, or enabled message wakes.
471470
pub(crate) async fn maybe_start_turn_for_pending_work_with_sub_id(
472471
self: &Arc<Self>,
473472
sub_id: String,
@@ -499,7 +498,7 @@ impl Session {
499498
if !input.iter().any(
500499
|item| matches!(item, TurnInput::InterAgentCommunication(mail) if mail.trigger_turn),
501500
) {
502-
// Queue-only mail wakes durable sleep without selecting a new task's settings.
501+
// Queue-only wakes retain the previous turn context without selecting a new task.
503502
start_options = previous_options;
504503
start_options.cyber_access_program = self
505504
.reference_context_item()

‎codex-rs/core/tests/suite/agent_control_tests.rs‎

Lines changed: 91 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,9 @@ use codex_protocol::ThreadId;
3131
use codex_protocol::error::CodexErr;
3232
use codex_protocol::error::CodexErrorDetails;
3333
use codex_protocol::error::Result as CodexResult;
34+
use codex_protocol::protocol::InterAgentCommunication;
3435
use codex_protocol::protocol::MultiAgentVersion;
36+
use codex_protocol::protocol::Op;
3537
use codex_protocol::protocol::SessionSource;
3638
use codex_protocol::protocol::SubAgentSource;
3739
use codex_protocol::protocol::TokenUsage;
@@ -44,6 +46,8 @@ use core_test_support::test_codex::test_codex;
4446
use core_test_support::wait_for_event;
4547
use futures::future::BoxFuture;
4648
use pretty_assertions::assert_eq;
49+
use serde_json::json;
50+
use test_case::test_case;
4751
use wiremock::MockServer;
4852

4953
struct TestAgentControl {
@@ -214,16 +218,21 @@ impl AgentControl for TestAgentControl {
214218

215219
async fn test_with_host_control(
216220
server: &MockServer,
221+
extra_features: &[Feature],
217222
) -> anyhow::Result<(TestCodex, Arc<TestAgentControl>)> {
218223
let thread_id = ThreadId::new();
219224
let controller = Arc::new(TestAgentControl::new(thread_id));
225+
let extra_features = extra_features.to_vec();
220226
let test = test_codex()
221-
.with_config(|config| {
227+
.with_config(move |config| {
222228
for feature in [
223229
Feature::Collab,
224230
Feature::MultiAgentV2,
225231
Feature::AgentMessageBoard,
226-
] {
232+
]
233+
.into_iter()
234+
.chain(extra_features.iter().copied())
235+
{
227236
config
228237
.features
229238
.enable(feature)
@@ -251,7 +260,7 @@ async fn test_with_host_control(
251260
async fn host_threads_preserve_lineage_settings_and_resume_routing() -> anyhow::Result<()> {
252261
skip_if_no_network!(Ok(()));
253262
let server = responses::start_mock_server().await;
254-
let (test, controller) = test_with_host_control(&server).await?;
263+
let (test, controller) = test_with_host_control(&server, &[]).await?;
255264
let root_id = test.session_configured.thread_id;
256265
let mut child_config = test.config.clone();
257266
child_config.service_tier = None;
@@ -490,7 +499,7 @@ async fn host_factory_follows_thread_lifecycle() -> anyhow::Result<()> {
490499
async fn collaboration_tools_dispatch_to_the_host_controller() -> anyhow::Result<()> {
491500
skip_if_no_network!(Ok(()));
492501
let server = responses::start_mock_server().await;
493-
let (test, _controller) = test_with_host_control(&server).await?;
502+
let (test, _controller) = test_with_host_control(&server, &[]).await?;
494503
let mut events = vec![responses::ev_response_created("tools")];
495504
for (call, name, arguments) in [
496505
(
@@ -541,7 +550,7 @@ use codex_protocol::protocol::InternalSessionSource;
541550
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
542551
async fn host_mailbox_notifies_waiting_agent_and_delivers_once() -> anyhow::Result<()> {
543552
let server = responses::start_mock_server().await;
544-
let (test, controller) = test_with_host_control(&server).await?;
553+
let (test, controller) = test_with_host_control(&server, &[]).await?;
545554
responses::mount_sse_once(
546555
&server,
547556
responses::sse(vec![
@@ -609,3 +618,80 @@ async fn host_mailbox_notifies_waiting_agent_and_delivers_once() -> anyhow::Resu
609618
test.codex.shutdown_and_wait().await?;
610619
Ok(())
611620
}
621+
622+
/// Both delivery paths wake an opted-in completed root exactly once.
623+
#[test_case(false; "direct_op")]
624+
#[test_case(true; "host_mailbox")]
625+
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
626+
async fn agent_mail_wakes_completed_root_when_enabled(
627+
via_host_mailbox: bool,
628+
) -> anyhow::Result<()> {
629+
let server = responses::start_mock_server().await;
630+
let (test, controller) =
631+
test_with_host_control(&server, &[Feature::WakeOnAgentMessages]).await?;
632+
let mut final_answer = responses::ev_assistant_message("first-answer", "initial work finished");
633+
final_answer["item"]["phase"] = json!("final_answer");
634+
let response = responses::mount_sse_sequence(
635+
&server,
636+
vec![
637+
responses::sse(vec![
638+
responses::ev_response_created("initial"),
639+
final_answer,
640+
responses::ev_completed("initial"),
641+
]),
642+
responses::sse(vec![
643+
responses::ev_response_created("agent-mail"),
644+
responses::ev_assistant_message("mail-answer", "received the agent update"),
645+
responses::ev_completed("agent-mail"),
646+
]),
647+
],
648+
)
649+
.await;
650+
test.submit_text_turn("Finish the initial work.").await?;
651+
assert_eq!(response.requests().len(), 1);
652+
653+
let message = InterAgentCommunication::new(
654+
AgentPath::root().join("worker").expect("valid worker path"),
655+
AgentPath::root(),
656+
Vec::new(),
657+
"Late agent update".into(),
658+
/*trigger_turn*/ false,
659+
);
660+
if via_host_mailbox {
661+
controller.mail.lock().expect("mail lock").push(message);
662+
controller.mailbox_state.send_replace(true);
663+
} else {
664+
test.codex
665+
.submit(Op::InterAgentCommunication {
666+
communication: message,
667+
start_options: Default::default(),
668+
})
669+
.await?;
670+
}
671+
wait_for_event(&test.codex, |event| {
672+
matches!(event, EventMsg::TurnComplete(_))
673+
})
674+
.await;
675+
test.codex.shutdown_and_wait().await?;
676+
677+
let requests = response.requests();
678+
assert_eq!(requests.len(), 2);
679+
assert_eq!(
680+
responses::received_responses_requests(&server).await.len(),
681+
2
682+
);
683+
assert_ne!(
684+
requests[0].body_json()["client_metadata"]["turn_id"],
685+
requests[1].body_json()["client_metadata"]["turn_id"]
686+
);
687+
let messages = requests[1].inputs_of_type("agent_message");
688+
assert_eq!(messages.len(), 1);
689+
assert_eq!(messages[0]["author"], "/root/worker");
690+
assert_eq!(messages[0]["recipient"], "/root");
691+
assert_eq!(
692+
messages[0]["content"],
693+
json!([{"type": "input_text", "text": "Late agent update"}])
694+
);
695+
assert!(controller.mail.lock().expect("mail lock").is_empty());
696+
Ok(())
697+
}

0 commit comments

Comments
 (0)