diff --git a/codex-rs/core/src/session/turn_input.rs b/codex-rs/core/src/session/turn_input.rs index 84b418a94d..7672b541a1 100644 --- a/codex-rs/core/src/session/turn_input.rs +++ b/codex-rs/core/src/session/turn_input.rs @@ -657,6 +657,36 @@ impl Session { } } + #[expect( + clippy::await_holding_invalid_type, + reason = "The admission lock keeps the previous turn stable while capturing its attribution." + )] + pub(crate) async fn reserve_pending_work_turn( + &self, + ) -> Option<(Arc>, TurnStartOptions)> { + let mut active_turn = self.active_turn.lock().await; + 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() { + return None; + } + let active_turn = active_turn.get_or_insert_with(ActiveTurn::default); + let previous_options = if needs_new_turn { + Default::default() + } else { + self.state + .lock() + .await + .turn_attribution + .as_ref() + .map(codex_history::TurnAttribution::start_options) + .unwrap_or_default() + }; + Some((Arc::clone(&active_turn.turn_state), previous_options)) + } + pub(crate) async fn route_realtime_text_input( self: &Arc, text: String, diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index 44baa9d60d..0d183a1793 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -457,20 +457,12 @@ impl Session { self: &Arc, sub_id: String, ) { - if !self.input_queue.has_pending_mailbox_items().await - || (!self.input_queue.has_trigger_turn_mailbox_items().await - && !self.has_outstanding_durable_sleep()) - { + if !self.input_queue.has_pending_mailbox_items().await { return; } - let turn_state = { - let mut active_turn = self.active_turn.lock().await; - if active_turn.is_some() { - return; - } - let active_turn = active_turn.get_or_insert_with(ActiveTurn::default); - Arc::clone(&active_turn.turn_state) + let Some((turn_state, previous_options)) = self.reserve_pending_work_turn().await else { + return; }; self.services @@ -493,6 +485,7 @@ impl Session { |item| matches!(item, TurnInput::InterAgentCommunication(mail) if mail.trigger_turn), ) { // Queue-only mail wakes durable sleep without selecting a new task's settings. + start_options = previous_options; start_options.cyber_access_program = self .reference_context_item() .await @@ -511,14 +504,18 @@ impl Session { turn_context.turn_metadata_state.set_turn_trigger(trigger); } if let Some(id) = start_options.parent_turn_id { - if let Some(initiating_agent_path) = input.iter().find_map(|item| { - let TurnInput::InterAgentCommunication(communication) = item else { - return None; - }; - communication - .trigger_turn - .then(|| communication.author.clone()) - }) { + if let Some(initiating_agent_path) = input + .iter() + .find_map(|item| { + let TurnInput::InterAgentCommunication(communication) = item else { + return None; + }; + communication + .trigger_turn + .then(|| communication.author.clone()) + }) + .or(start_options.initiating_agent_path) + { turn_context .turn_metadata_state .set_initiating_agent_path(initiating_agent_path); diff --git a/codex-rs/core/tests/suite/pending_input.rs b/codex-rs/core/tests/suite/pending_input.rs index fee9018175..fed5252322 100644 --- a/codex-rs/core/tests/suite/pending_input.rs +++ b/codex-rs/core/tests/suite/pending_input.rs @@ -27,6 +27,7 @@ use codex_protocol::models::AgentMessageInputContent; use codex_protocol::models::ContentItem; use codex_protocol::models::PermissionProfile; use codex_protocol::models::ResponseItem; +use codex_protocol::openai_models::ReasoningEffort; use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::CodexErrorInfo; use codex_protocol::protocol::EventMsg; @@ -518,14 +519,25 @@ impl codex_extension_api::ThreadLifecycleContributor } } +#[test_case(false, Some("automation"), None, false; "live_root")] +#[test_case(true, Some("automation"), None, false; "reloaded_root")] +#[test_case(true, Some("automation"), Some("original-parent"), false; "reloaded_delegated")] +#[test_case(true, None, None, false; "reloaded_unknown")] +#[test_case(false, Some("automation"), Some("original-parent"), true; "settings_while_asleep")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn queue_only_agent_mail_wakes_sleeping_root_with_previous_turn_context() { +async fn queue_only_agent_mail_wakes_sleeping_root_with_previous_turn_context( + reload: bool, + trigger: Option<&str>, + parent: Option<&str>, + update_settings: bool, +) -> anyhow::Result<()> { const CHILD_MESSAGE: &str = "worker completed"; let server = responses::start_mock_server().await; let requests = responses::mount_sse_sequence( &server, vec![ + responses::sse_completed("resp-earlier-human"), responses::sse_completed("resp-initial"), responses::sse_completed("resp-wake"), ], @@ -534,41 +546,98 @@ async fn queue_only_agent_mail_wakes_sleeping_root_with_previous_turn_context() let mut extensions = codex_extension_api::ExtensionRegistryBuilder::::new(); extensions.thread_lifecycle_contributor(Arc::new(SleepingRootExtension)); - let codex = test_codex() + let mut builder = test_codex() .with_model("gpt-5.4") .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) - .with_extensions(Arc::new(extensions.build())) - .build_with_auto_env(&server) - .await - .expect("build Codex test session") - .codex; + .with_extensions(Arc::new(extensions.build())); + let mut test = builder.build_with_auto_env(&server).await?; - codex + test.codex .start_or_steer_turn( TurnInputRequest::user_input(vec![UserInput::Text { - text: "wait for the worker".to_string(), + text: "establish an earlier human context".to_string(), text_elements: Vec::new(), }]) .on_start(TurnStartOptions { + turn_trigger: Some("composer".to_string()), + ..Default::default() + }), + ) + .await?; + wait_for_turn_complete(&test.codex).await; + // Automatic turns need not replace the model-context baseline on replay. + let TurnInputSubmission::Started { turn_id } = test + .codex + .start_or_steer_turn( + TurnInputRequest::user_input(Vec::new()).on_start(TurnStartOptions { + turn_trigger: trigger.map(str::to_owned), + parent_turn_id: parent.map(str::to_owned), + root_turn_id: parent.map(|_| "original-root".to_string()), cyber_access_program: Some(CyberAccessProgram::Standard), ..Default::default() }), ) - .await - .expect("start initial turn"); - wait_for_turn_complete(&codex).await; - enqueue_queue_only_agent_mail(&codex, CHILD_MESSAGE).await; - wait_for_turn_complete(&codex).await; + .await? + else { + panic!("expected an automatic turn"); + }; + wait_for_turn_complete(&test.codex).await; + if reload { + if let Some(url) = test.executor_environment().exec_server_url() { + builder = builder.with_exec_server_url(url); + } + test = builder.restart(&server, &test).await?; + } + if update_settings { + core_test_support::submit_thread_settings( + &test.codex, + ThreadSettingsOverrides { + effort: Some(Some(ReasoningEffort::Low)), + ..Default::default() + }, + ) + .await?; + } + enqueue_queue_only_agent_mail(&test.codex, CHILD_MESSAGE).await; + wait_for_turn_complete(&test.codex).await; + let captured = requests.requests(); + let metadata = captured + .iter() + .skip(1) + .map(|request| { + serde_json::from_str::( + &request.header("x-codex-turn-metadata").expect("metadata"), + ) + }) + .collect::, _>>()?; + let root = parent.map_or(turn_id.as_str(), |_| "original-root"); assert_eq!( - requests - .requests() + metadata .iter() - .map(|request| request.body_json()["access_programs"].clone()) + .map(|value| ( + value["turn_trigger"].as_str(), + value["parent_turn_id"].as_str(), + value["root_turn_id"].as_str(), + )) .collect::>(), - vec![json!({"cyber": "standard"}); 2], + vec![(trigger, parent, Some(root)); 2], ); - let history = codex + assert_ne!(metadata[0]["turn_id"], metadata[1]["turn_id"]); + // Live wake keeps the existing execution-setting behavior. Attribution after + // reload must use the automatic turn, not the earlier human context baseline. + if !reload { + assert_eq!( + captured + .iter() + .skip(1) + .map(|request| request.body_json()["access_programs"].clone()) + .collect::>(), + vec![json!({"cyber": "standard"}); 2], + ); + } + let history = test + .codex .load_history(/*include_archived*/ true) .await .expect("load persisted thread history"); @@ -587,6 +656,45 @@ async fn queue_only_agent_mail_wakes_sleeping_root_with_previous_turn_context() ) ) })); + + if parent.is_none() { + return Ok(()); + } + let followup = + responses::mount_sse_once(&server, responses::sse_completed("resp-followup")).await; + test.codex + .submit(Op::InterAgentCommunication { + communication: InterAgentCommunication::new( + AgentPath::try_from("/root/worker").expect("worker path should parse"), + AgentPath::root(), + Vec::new(), + "start different work".to_string(), + /*trigger_turn*/ true, + ), + start_options: TurnStartOptions { + turn_trigger: Some("composer".to_string()), + parent_turn_id: Some("new-parent".to_string()), + root_turn_id: Some("new-root".to_string()), + ..Default::default() + }, + }) + .await?; + wait_for_turn_complete(&test.codex).await; + let followup: Value = serde_json::from_str( + &followup + .single_request() + .header("x-codex-turn-metadata") + .expect("metadata"), + )?; + assert_eq!( + ( + followup["turn_trigger"].as_str(), + followup["parent_turn_id"].as_str(), + followup["root_turn_id"].as_str(), + ), + (Some("composer"), Some("new-parent"), Some("new-root")), + ); + Ok(()) } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] diff --git a/codex-rs/core/tests/suite/subagent_notifications.rs b/codex-rs/core/tests/suite/subagent_notifications.rs index 04b3026a50..16ca5fa447 100644 --- a/codex-rs/core/tests/suite/subagent_notifications.rs +++ b/codex-rs/core/tests/suite/subagent_notifications.rs @@ -2794,8 +2794,12 @@ async fn plaintext_multi_agent_v2_completion_sends_agent_message( Ok(()) } +#[test_case(false; "live")] +#[test_case(true; "reloaded_sleep")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn multi_agent_v2_peer_followup_completion_notifies_initiating_turn() -> Result<()> { +async fn multi_agent_v2_peer_followup_completion_notifies_initiating_turn( + reload_sleeping_worker: bool, +) -> Result<()> { const SPAWN_WORKER_PROMPT: &str = "spawn the completion-routing worker"; const SPAWN_REQUESTER_PROMPT: &str = "spawn the completion-routing requester"; const READ_RESULT_PROMPT: &str = "read the completion-routing worker result"; @@ -2877,7 +2881,7 @@ async fn multi_agent_v2_peer_followup_completion_notifies_initiating_turn() -> R submit_turn_with_trigger(&test, SPAWN_WORKER_PROMPT, "automation_cron_scheduled").await?; let worker_thread_id = created_threads.recv().await?; - let worker_thread = test.thread_manager.get_thread(worker_thread_id).await?; + let mut worker_thread = test.thread_manager.get_thread(worker_thread_id).await?; wait_for_event(worker_thread.as_ref(), |event| { matches!(event, EventMsg::TurnComplete(_)) }) @@ -2995,7 +2999,7 @@ async fn multi_agent_v2_peer_followup_completion_notifies_initiating_turn() -> R }) .await; - let worker_followup_turn_id = worker_followup_request + let mut worker_followup_turn_id = worker_followup_request .requests() .into_iter() .find_map(|request| { @@ -3017,6 +3021,52 @@ async fn multi_agent_v2_peer_followup_completion_notifies_initiating_turn() -> R }) }) .expect("worker follow-up model request"); + if reload_sleeping_worker { + worker_thread.shutdown_and_wait().await?; + test.thread_manager.remove_thread(&worker_thread_id).await; + test.thread_manager + .ensure_multi_agent_v2_child_loaded(worker_thread_id) + .await?; + worker_thread = test.thread_manager.get_thread(worker_thread_id).await?; + // Simulate the extension restoring durable sleep; Core owns the + // queue-only wake and completion routing exercised below. + worker_thread + .thread_extension_data() + .insert(codex_extension_items::sleep::SleepItem { + id: "peer-worker-sleep".to_string(), + duration_ms: 60_000, + }); + let wake_request = mount_sse_once_match( + &server, + |request: &wiremock::Request| body_contains(request, "wake sleeping peer work"), + sse(vec![ + ev_response_created("resp-peer-worker-wake"), + ev_assistant_message("msg-peer-worker-wake", "peer follow-up finished"), + ev_completed("resp-peer-worker-wake"), + ]), + ) + .await; + worker_thread + .submit(Op::InterAgentCommunication { + communication: codex_protocol::protocol::InterAgentCommunication::new( + codex_protocol::AgentPath::root(), + codex_protocol::AgentPath::try_from("/root/worker").expect("worker path"), + Vec::new(), + "wake sleeping peer work".to_string(), + /*trigger_turn*/ false, + ), + start_options: TurnStartOptions::default(), + }) + .await?; + worker_followup_turn_id = + wait_for_event_match(worker_thread.as_ref(), |event| match event { + EventMsg::TurnComplete(completed) => Some(completed.turn_id.clone()), + _ => None, + }) + .await; + let request = wake_request.single_request(); + assert_parent_turn(&request.body_json(), Some(&requester_turn_id))?; + } let completed = timeout( Duration::from_secs(5), wait_for_event_match(requester_thread.as_ref(), |event| match event { @@ -3026,6 +3076,7 @@ async fn multi_agent_v2_peer_followup_completion_notifies_initiating_turn() -> R TurnItem::SubAgentActivity(activity) if activity.kind == SubAgentActivityKind::Completed && activity.agent_thread_id == worker_thread_id + && activity.id == format!("subagent-completed-{worker_followup_turn_id}") ) => { Some(completed.clone())