diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index 7361c445ceb5..696982b0a131 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/codex-rs/core/src/agent/control.rs @@ -18,6 +18,7 @@ use crate::session_prefix::format_subagent_notification_message; use crate::thread_activity::ThreadActivityGate; use crate::thread_activity::ThreadActivityHandle; use crate::thread_activity::ThreadActivityRegistrationError; +use crate::thread_activity::ThreadActivityReservation; use crate::thread_manager::ResumeThreadWithHistoryOptions; use crate::thread_manager::ThreadManagerState; use crate::thread_rollout_truncation::truncate_rollout_to_last_n_fork_turns; @@ -81,6 +82,11 @@ pub(crate) struct LiveAgent { pub(crate) status: AgentStatus, } +struct CompletionWatcherChild { + status: watch::Receiver, + _activity: ThreadActivityReservation, +} + #[derive(Clone, Debug, Serialize, PartialEq, Eq)] pub(crate) struct ListedAgent { pub(crate) agent_name: String, @@ -155,6 +161,33 @@ impl AgentControl { self.rollout_budget.as_ref() } + pub(crate) async fn has_thread_spawn_descendant_idle_shutdown_blocker( + &self, + parent_thread_id: ThreadId, + ) -> bool { + let Ok(state) = self.upgrade() else { + return false; + }; + for descendant_id in self + .thread_activity_gate + .descendant_thread_ids(parent_thread_id) + { + let Ok(descendant) = state.get_thread(descendant_id).await else { + continue; + }; + if !is_final(&descendant.agent_status().await) + || descendant + .codex + .session + .has_local_idle_shutdown_blocker() + .await + { + return true; + } + } + false + } + /// Send rich user input items to an existing agent thread. pub(crate) async fn send_input( &self, @@ -479,6 +512,7 @@ impl AgentControl { session_source: Option, child_reference: String, child_agent_path: Option, + completion_child: Option, ) { let Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { parent_thread_id, .. @@ -488,19 +522,13 @@ impl AgentControl { }; let control = self.clone(); tokio::spawn(async move { - let status = match control.subscribe_status(child_thread_id).await { - Ok(mut status_rx) => { - let mut status = status_rx.borrow().clone(); - while !is_final(&status) { - if status_rx.changed().await.is_err() { - status = control.get_status(child_thread_id).await; - break; - } - status = status_rx.borrow().clone(); - } - status - } - Err(_) => control.get_status(child_thread_id).await, + let mut completion_child = completion_child; + let status = match completion_child.as_mut() { + Some(child) => wait_for_final_status(&mut child.status).await, + None => match control.subscribe_status(child_thread_id).await { + Ok(mut status) => wait_for_final_status(&mut status).await, + Err(_) => control.get_status(child_thread_id).await, + }, }; if !is_final(&status) { return; @@ -509,12 +537,13 @@ impl AgentControl { let Ok(state) = control.upgrade() else { return; }; - let child_thread = state.get_thread(child_thread_id).await.ok(); - let child_uses_multi_agent_v2 = match child_thread.as_ref() { - Some(child_thread) => { - child_thread.multi_agent_version() == Some(MultiAgentVersion::V2) - } - None => true, + let child_uses_multi_agent_v2 = match completion_child.as_ref() { + Some(_) => false, + None => state + .get_thread(child_thread_id) + .await + .ok() + .is_none_or(|child| child.multi_agent_version() == Some(MultiAgentVersion::V2)), }; if child_agent_path.is_some() && child_uses_multi_agent_v2 { let Some(child_agent_path) = child_agent_path.clone() else { @@ -769,6 +798,17 @@ fn agent_matches_prefix(agent_path: Option<&AgentPath>, prefix: &AgentPath) -> b }) } +async fn wait_for_final_status(status: &mut watch::Receiver) -> AgentStatus { + let mut current = status.borrow().clone(); + while !is_final(¤t) { + if status.changed().await.is_err() { + return AgentStatus::NotFound; + } + current = status.borrow().clone(); + } + current +} + pub(crate) fn render_input_preview(input: &[UserInput]) -> String { input .iter() diff --git a/codex-rs/core/src/agent/control/spawn.rs b/codex-rs/core/src/agent/control/spawn.rs index 9078434f1fe8..af71bac8bd95 100644 --- a/codex-rs/core/src/agent/control/spawn.rs +++ b/codex-rs/core/src/agent/control/spawn.rs @@ -328,6 +328,27 @@ impl AgentControl { (None, _, _) => Box::pin(state.spawn_new_thread(config.clone(), self.clone())).await?, }; agent_metadata.agent_id = Some(new_thread.thread_id); + let completion_child = if multi_agent_version != MultiAgentVersion::V2 + && matches!( + ¬ification_source, + Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { .. })) + ) { + let Some(completion_watcher) = + new_thread.thread.codex.session.try_reserve_turn_activity() + else { + let _ = state.remove_thread(&new_thread.thread_id).await; + let _ = new_thread.thread.shutdown_and_wait().await; + return Err(CodexErr::InvalidRequest( + "parent thread is shutting down".to_string(), + )); + }; + Some(CompletionWatcherChild { + status: new_thread.thread.subscribe_status(), + _activity: completion_watcher, + }) + } else { + None + }; reservation.commit(agent_metadata.clone()); if let Some(residency_slot) = residency_slot { residency_slot.commit(new_thread.thread_id); @@ -415,6 +436,7 @@ impl AgentControl { notification_source, child_reference, agent_metadata.agent_path.clone(), + completion_child, ); } @@ -737,6 +759,30 @@ impl AgentControl { inherited_exec_policy, }) .await?; + let completion_child = if multi_agent_version != MultiAgentVersion::V2 + && matches!( + ¬ification_source, + SessionSource::SubAgent(SubAgentSource::ThreadSpawn { .. }) + ) { + let Some(completion_watcher) = resumed_thread + .thread + .codex + .session + .try_reserve_turn_activity() + else { + let _ = state.remove_thread(&resumed_thread.thread_id).await; + let _ = resumed_thread.thread.shutdown_and_wait().await; + return Err(CodexErr::InvalidRequest( + "parent thread is shutting down".to_string(), + )); + }; + Some(CompletionWatcherChild { + status: resumed_thread.thread.subscribe_status(), + _activity: completion_watcher, + }) + } else { + None + }; let mut agent_metadata = agent_metadata; agent_metadata.agent_id = Some(resumed_thread.thread_id); reservation.commit(agent_metadata.clone()); @@ -754,6 +800,7 @@ impl AgentControl { Some(notification_source.clone()), child_reference, agent_metadata.agent_path.clone(), + completion_child, ); } self.persist_thread_spawn_edge_for_source( diff --git a/codex-rs/core/src/agent/control_tests.rs b/codex-rs/core/src/agent/control_tests.rs index 4741c3c101a3..4057e02c909e 100644 --- a/codex-rs/core/src/agent/control_tests.rs +++ b/codex-rs/core/src/agent/control_tests.rs @@ -1950,6 +1950,11 @@ async fn spawn_child_completion_notifies_parent_history() { .get_thread(child_thread_id) .await .expect("child thread should exist"); + harness + .manager + .remove_thread(&child_thread_id) + .await + .expect("child removal should succeed"); let _ = child_thread .submit(Op::Shutdown {}) .await @@ -2099,6 +2104,7 @@ async fn multi_agent_v2_completion_queues_message_for_direct_parent() { })), tester_path.to_string(), Some(tester_path.clone()), + /*completion_child*/ None, ); let tester_turn = tester_thread.codex.session.new_default_turn().await; tester_thread @@ -2187,6 +2193,7 @@ async fn completion_watcher_notifies_parent_when_child_is_missing() { })), child_thread_id.to_string(), /*child_agent_path*/ None, + /*completion_child*/ None, ); assert_eq!(wait_for_subagent_notification(&parent_thread).await, true); diff --git a/codex-rs/core/src/session/inject.rs b/codex-rs/core/src/session/inject.rs index 80f7b491a2a3..b4f65883319e 100644 --- a/codex-rs/core/src/session/inject.rs +++ b/codex-rs/core/src/session/inject.rs @@ -73,6 +73,12 @@ impl Session { input, )); }; + let Some(thread_activity) = self.try_reserve_turn_activity() else { + return Err(TryStartTurnIfIdleError::new( + TryStartTurnIfIdleRejectionReason::Busy, + input, + )); + }; let turn_state = { let mut active_turn = self.active_turn.lock().await; if active_turn.is_some() { @@ -82,6 +88,7 @@ impl Session { )); } let active_turn = active_turn.get_or_insert_with(ActiveTurn::default); + active_turn.thread_activity = Some(thread_activity); Arc::clone(&active_turn.turn_state) }; drop(reservation); diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 404075ccf204..8b3f3a899d12 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -1048,7 +1048,7 @@ fn push_prompt_fragment( } impl Session { - async fn has_idle_shutdown_blocker(&self) -> bool { + pub(crate) async fn has_local_idle_shutdown_blocker(&self) -> bool { matches!(self.agent_status.borrow().clone(), AgentStatus::Running) || self.conversation.running_state().await.is_some() || *self.services.elicitations.subscribe().borrow() @@ -1057,6 +1057,15 @@ impl Session { || !self.list_background_terminals().await.is_empty() } + async fn has_idle_shutdown_blocker(&self) -> bool { + self.has_local_idle_shutdown_blocker().await + || self + .services + .agent_control + .has_thread_spawn_descendant_idle_shutdown_blocker(self.thread_id) + .await + } + pub(crate) async fn try_claim_idle_shutdown( &self, ) -> Option> { diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 89db411437dc..66efbf5b55bc 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -612,6 +612,14 @@ impl Session { self.try_reserve_submission(/*close*/ false) } + pub(crate) fn try_reserve_turn_activity(&self) -> Option { + self.thread_activity.try_reserve(/*close*/ false) + } + + pub(crate) fn try_reserve_dispatched_turn_activity(&self) -> Option { + self.thread_activity.try_reserve_dispatched_turn() + } + pub(crate) fn try_reserve_idle_shutdown(&self) -> Option> { let thread = self.thread_activity.try_reserve_idle_shutdown()?; let submission = self.submission_lifecycle.try_reserve_idle_shutdown()?; diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index e7911e4b46bf..8175778b4f7a 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -10194,7 +10194,7 @@ async fn try_start_turn_if_idle_rejects_active_turn_without_injecting() { Vec::::new(), sess.input_queue.get_pending_input(&sess.active_turn).await ); - + assert!(sess.thread_activity.try_reserve_idle_shutdown().is_none()); sess.abort_all_tasks(TurnAbortReason::Interrupted).await; } @@ -10235,6 +10235,38 @@ async fn interrupt_waits_for_idle_turn_start_to_install_task() { .expect("idle turn should be accepted"); interrupt.await.expect("interrupt task should not panic"); assert!(sess.active_turn.lock().await.is_none()); + assert!(sess.thread_activity.try_reserve_idle_shutdown().is_some()); +} + +#[tokio::test] +#[expect( + clippy::await_holding_invalid_type, + reason = "the test holds session state to pause replacement at a deterministic boundary" +)] +async fn replacement_holds_thread_activity_through_startup() { + let (sess, tc, _rx) = make_session_and_context_with_rx().await; + let task = || NeverEndingTask { + kind: TaskKind::Regular, + listen_to_cancellation_token: true, + }; + sess.spawn_task(Arc::clone(&tc), Vec::new(), task()).await; + let replacement_context = sess.new_default_turn().await; + let state = sess.state.lock().await; + let replacement = tokio::spawn({ + let sess = Arc::clone(&sess); + async move { + sess.spawn_task(replacement_context, Vec::new(), task()) + .await + } + }); + while sess.active_turn.lock().await.is_some() { + tokio::task::yield_now().await; + } + assert!(sess.thread_activity.try_reserve_idle_shutdown().is_none()); + drop(state); + replacement.await.expect("replacement should not panic"); + sess.abort_all_tasks(TurnAbortReason::Interrupted).await; + assert!(sess.thread_activity.try_reserve_idle_shutdown().is_some()); } #[tokio::test] diff --git a/codex-rs/core/src/state/turn.rs b/codex-rs/core/src/state/turn.rs index d1e31b79949f..20b0408d05f8 100644 --- a/codex-rs/core/src/state/turn.rs +++ b/codex-rs/core/src/state/turn.rs @@ -22,6 +22,7 @@ use crate::agent::control::AgentExecutionGuard; use crate::session::TurnInputQueue; use crate::session::turn_context::TurnContext; use crate::tasks::AnySessionTask; +use crate::thread_activity::ThreadActivityReservation; use codex_protocol::models::AdditionalPermissionProfile; use codex_protocol::protocol::ReviewDecision; use codex_protocol::protocol::TokenUsage; @@ -30,6 +31,7 @@ use codex_protocol::protocol::TokenUsage; pub(crate) struct ActiveTurn { pub(crate) task: Option, pub(crate) turn_state: Arc>, + pub(crate) thread_activity: Option, } /// Whether mailbox deliveries should still be folded into the current turn. @@ -58,6 +60,7 @@ impl Default for ActiveTurn { Self { task: None, turn_state: Arc::new(Mutex::new(TurnState::default())), + thread_activity: None, } } } diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index 2dda4b2e29f5..82ad37fc92b0 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -322,9 +322,11 @@ impl Session { task: T, ) { let _turn_start = self.turn_start_lock.lock().await; - self.abort_all_tasks_locked(TurnAbortReason::Replaced).await; + let (_, replaced_turn_activity) = + self.abort_all_tasks_locked(TurnAbortReason::Replaced).await; self.clear_connector_selection().await; self.start_task_locked(turn_context, input, task).await; + drop(replaced_turn_activity); } pub(crate) async fn start_task_locked( @@ -333,6 +335,30 @@ impl Session { input: Vec, task: T, ) { + let turn_state = { + let mut active = self.active_turn.lock().await; + if active.is_none() { + let Some(thread_activity) = self.try_reserve_dispatched_turn_activity() else { + return; + }; + *active = Some(ActiveTurn { + thread_activity: Some(thread_activity), + ..Default::default() + }); + } + let Some(turn) = active.as_mut() else { + return; + }; + debug_assert!(turn.task.is_none()); + if turn.thread_activity.is_none() { + let Some(thread_activity) = self.try_reserve_dispatched_turn_activity() else { + return; + }; + turn.thread_activity = Some(thread_activity); + } + Arc::clone(&turn.turn_state) + }; + let task: Arc = Arc::new(task); let task_kind = task.kind(); let span_name = task.span_name(); @@ -354,14 +380,7 @@ impl Session { .lock() .await .clear_turn(&turn_context.sub_id); - let pending_items = self.input_queue.get_pending_input(&self.active_turn).await; - let turn_state = { - let mut active = self.active_turn.lock().await; - let turn = active.get_or_insert_with(ActiveTurn::default); - debug_assert!(turn.task.is_none()); - Arc::clone(&turn.turn_state) - }; turn_state.lock().await.token_usage_at_turn_start = token_usage_at_turn_start.clone(); self.input_queue .extend_pending_input_for_turn_state(turn_state.as_ref(), pending_items) @@ -371,8 +390,15 @@ impl Session { let turn_extension_data = Arc::clone(&turn_context.extension_data); let mut active = self.active_turn.lock().await; - let turn = active.get_or_insert_with(ActiveTurn::default); - debug_assert!(turn.task.is_none()); + let Some(turn) = active.as_mut() else { + return; + }; + if turn.task.is_some() + || turn.thread_activity.is_none() + || !Arc::ptr_eq(&turn.turn_state, &turn_state) + { + return; + } let agent_execution_guard = self.services.agent_control.execution_guard( turn_context.multi_agent_version, &turn_context.session_source, @@ -494,12 +520,18 @@ impl Session { let Some(reservation) = self.try_reserve_activity() else { return; }; + let Some(thread_activity) = self.try_reserve_turn_activity() else { + return; + }; { let mut active_turn = self.active_turn.lock().await; if active_turn.is_some() { return; } - *active_turn = Some(ActiveTurn::default()); + *active_turn = Some(ActiveTurn { + thread_activity: Some(thread_activity), + ..Default::default() + }); } drop(reservation); @@ -516,14 +548,21 @@ impl Session { )] pub async fn abort_all_tasks(self: &Arc, reason: TurnAbortReason) { let turn_start = self.turn_start_lock.lock().await; - let restart_pending = self.abort_all_tasks_locked(reason).await; + let (restart_pending, transition_activity) = self.abort_all_tasks_locked(reason).await; drop(turn_start); if restart_pending { self.maybe_start_turn_for_pending_work().await; } + drop(transition_activity); } - async fn abort_all_tasks_locked(self: &Arc, reason: TurnAbortReason) -> bool { + async fn abort_all_tasks_locked( + self: &Arc, + reason: TurnAbortReason, + ) -> ( + bool, + Option, + ) { let mut aborted_turn = false; let mut active_turn_to_clear = None; let mut turn_context = None; @@ -548,7 +587,12 @@ impl Session { // in-flight approval wait can surface as a model-visible rejection before TurnAborted. self.input_queue.clear_pending(active_turn).await; } - reason == TurnAbortReason::Interrupted && aborted_turn + let transition_activity = + active_turn_to_clear.and_then(|mut active_turn| active_turn.thread_activity.take()); + ( + reason == TurnAbortReason::Interrupted && aborted_turn, + transition_activity, + ) } #[expect( @@ -620,15 +664,18 @@ impl Session { .turn_metadata_state .cancel_git_enrichment_task(); - let turn_state = { + let completed_turn = { let mut active = self.active_turn.lock().await; active.as_mut().and_then(|active_turn| { let task = active_turn.task.take()?; task.handle.detach(); - Some(Arc::clone(&active_turn.turn_state)) + Some(( + Arc::clone(&active_turn.turn_state), + active_turn.thread_activity.take(), + )) }) }; - let Some(turn_state) = turn_state else { + let Some((turn_state, thread_activity)) = completed_turn else { return; }; let pending_input = self @@ -818,6 +865,7 @@ impl Session { }) }; self.send_event(turn_context.as_ref(), event).await; + drop(thread_activity); self.services .guardian_rejection_circuit_breaker .lock() diff --git a/codex-rs/core/src/thread_activity.rs b/codex-rs/core/src/thread_activity.rs index b2f0b4297b01..a89870f392c8 100644 --- a/codex-rs/core/src/thread_activity.rs +++ b/codex-rs/core/src/thread_activity.rs @@ -139,6 +139,36 @@ impl ThreadActivityGate { }) } + pub(crate) fn descendant_thread_ids(&self, ancestor_id: ThreadId) -> Vec { + let state = self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + state + .nodes + .iter() + .filter_map(|(thread_id, node)| { + if *thread_id == ancestor_id { + return None; + } + let mut parent_thread_id = node.parent_thread_id; + let mut visited = HashSet::new(); + while let Some(parent_id) = parent_thread_id + && visited.insert(parent_id) + { + if parent_id == ancestor_id { + return Some(*thread_id); + } + parent_thread_id = state + .nodes + .get(&parent_id) + .and_then(|node| node.parent_thread_id); + } + None + }) + .collect() + } + fn validate_ancestors( state: &ThreadActivityState, mut parent_thread_id: Option, @@ -191,6 +221,42 @@ impl ThreadActivityGate { }) } + fn try_reserve_dispatched_turn( + self: &Arc, + thread_id: ThreadId, + generation: u64, + ) -> Option { + let mut state = self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let node = state.nodes.get(&thread_id)?; + if node.generation != generation + || node.closed + || node.initializing + // One committed slot is the queued Shutdown itself; a turn may start only while an + // earlier FIFO submission is also still committed. + || (node.closing && node.committed <= 1) + || !Self::validate_ancestors( + &state, + node.parent_thread_id, + /*allow_closing*/ false, + ) + { + return None; + } + let node = state.nodes.get_mut(&thread_id)?; + node.active = node.active.checked_add(1)?; + Some(ThreadActivityReservation { + gate: Arc::clone(self), + thread_id, + generation, + clear_closing_on_drop: false, + delivery_prepared: false, + active: true, + }) + } + fn try_reserve_idle_shutdown( self: &Arc, thread_id: ThreadId, @@ -417,6 +483,11 @@ impl ThreadActivityHandle { .try_reserve(self.thread_id, self.generation, close) } + pub(crate) fn try_reserve_dispatched_turn(&self) -> Option { + self.gate + .try_reserve_dispatched_turn(self.thread_id, self.generation) + } + pub(crate) fn try_reserve_idle_shutdown(&self) -> Option { self.gate .try_reserve_idle_shutdown(self.thread_id, self.generation) diff --git a/codex-rs/core/src/thread_activity_tests.rs b/codex-rs/core/src/thread_activity_tests.rs index b5b0f8d518cc..eca5d9d1042b 100644 --- a/codex-rs/core/src/thread_activity_tests.rs +++ b/codex-rs/core/src/thread_activity_tests.rs @@ -112,27 +112,28 @@ fn unpublication_holds_activity_through_a_provisional_close() { fn activity_reservations_are_atomic_and_survive_session_teardown() { let gate = Arc::new(ThreadActivityGate::default()); let parent_id = ThreadId::new(); - let child_id = ThreadId::new(); + let intermediate_id = ThreadId::new(); + let leaf_id = ThreadId::new(); let parent = register(&gate, parent_id, /*parent_thread_id*/ None); - let child = register(&gate, child_id, Some(parent_id)); + let intermediate = register(&gate, intermediate_id, Some(parent_id)); + let leaf = register(&gate, leaf_id, Some(intermediate_id)); - let child_activity = child + let completion_activity = leaf .try_reserve(/*close*/ false) - .expect("reserve child activity"); - assert!(parent.try_reserve_idle_shutdown().is_none()); - drop(child_activity); - let parent_activity = parent - .try_reserve(/*close*/ false) - .expect("reserve parent activity"); + .expect("reserve child completion watcher"); assert!(parent.try_reserve_idle_shutdown().is_none()); + drop(intermediate); + let descendants = gate.descendant_thread_ids(parent_id); + assert!(descendants.contains(&intermediate_id)); + assert!(descendants.contains(&leaf_id)); drop(parent); assert!(gate.register(parent_id, /*parent_thread_id*/ None).is_err()); - drop(parent_activity); + drop(completion_activity); let parent = register(&gate, parent_id, /*parent_thread_id*/ None); let shutdown = parent .try_reserve_idle_shutdown() .expect("reserve idle shutdown"); - assert!(child.try_reserve(/*close*/ false).is_none()); + assert!(leaf.try_reserve(/*close*/ false).is_none()); drop(shutdown); let mut failed_shutdown = parent @@ -164,3 +165,29 @@ fn receiver_can_finish_before_sender_observes_delivery() { let node = state.nodes.get(&thread_id).expect("thread node"); assert_eq!((node.active, node.committed), (0, 0)); } + +#[test] +fn accepted_submission_can_start_before_queued_shutdown() { + let gate = Arc::new(ThreadActivityGate::default()); + let thread_id = ThreadId::new(); + let thread = register(&gate, thread_id, /*parent_thread_id*/ None); + + let mut user_submission = thread + .try_reserve(/*close*/ false) + .expect("reserve user submission"); + assert!(user_submission.prepare_delivery()); + user_submission.commit(); + let mut queued_shutdown = thread + .try_reserve(/*close*/ true) + .expect("reserve later shutdown"); + assert!(queued_shutdown.prepare_delivery()); + queued_shutdown.commit(); + + let turn = thread + .try_reserve_dispatched_turn() + .expect("accepted user submission starts before shutdown"); + thread.finish_submission(); + assert!(thread.try_reserve_dispatched_turn().is_none()); + drop(turn); + thread.finish_submission(); +} diff --git a/codex-rs/core/tests/suite/agent_execution.rs b/codex-rs/core/tests/suite/agent_execution.rs index 69db20cc57f6..a2530c05f6c8 100644 --- a/codex-rs/core/tests/suite/agent_execution.rs +++ b/codex-rs/core/tests/suite/agent_execution.rs @@ -1,13 +1,19 @@ use anyhow::Result; +use codex_core::sandboxing::SandboxPermissions; use codex_features::Feature; +use codex_protocol::models::PermissionProfile; +use codex_protocol::protocol::AskForApproval; +use codex_protocol::protocol::EventMsg; use core_test_support::responses::ev_assistant_message; use core_test_support::responses::ev_completed; use core_test_support::responses::ev_function_call_with_namespace; use core_test_support::responses::ev_response_created; +use core_test_support::responses::ev_shell_command_call_with_args; use core_test_support::responses::mount_sse_once_match; use core_test_support::responses::sse; use core_test_support::responses::start_mock_server; use core_test_support::test_codex::test_codex; +use core_test_support::wait_for_event; use pretty_assertions::assert_eq; use serde_json::json; use std::time::Duration; @@ -131,3 +137,95 @@ async fn v2_nested_spawn_checks_shared_active_execution_capacity() -> Result<()> Ok(()) } + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn idle_parent_cannot_shutdown_while_unpublished_child_waits_for_approval() -> Result<()> { + let server = start_mock_server().await; + let spawn_args = serde_json::to_string(&json!({ + "message": FIRST_TASK, + "task_name": "first", + }))?; + mount_sse_once_match( + &server, + |request: &wiremock::Request| body_contains(request, FIRST_PROMPT), + sse(vec![ + ev_response_created("parent-spawn-response"), + ev_function_call_with_namespace( + "spawn-call", + MULTI_AGENT_V2_NAMESPACE, + "spawn_agent", + &spawn_args, + ), + ev_completed("parent-spawn-response"), + ]), + ) + .await; + let shell_args = json!({ + "command": "printf child", + "sandbox_permissions": SandboxPermissions::RequireEscalated, + "justification": "hold the child at a deterministic approval boundary", + }); + mount_sse_once_match( + &server, + |request: &wiremock::Request| { + body_contains(request, FIRST_TASK) && !has_function_call_output(request, "spawn-call") + }, + sse(vec![ + ev_response_created("child-response"), + ev_shell_command_call_with_args("child-shell", &shell_args), + ev_completed("child-response"), + ]), + ) + .await; + mount_sse_once_match( + &server, + |request: &wiremock::Request| has_function_call_output(request, "spawn-call"), + sse(vec![ + ev_response_created("parent-followup-response"), + ev_assistant_message("parent-followup-message", "spawned"), + ev_completed("parent-followup-response"), + ]), + ) + .await; + + let mut builder = test_codex().with_model("koffing").with_config(|config| { + for feature in [Feature::Collab, Feature::MultiAgentV2] { + config + .features + .enable(feature) + .expect("test config should allow feature update"); + } + }); + let test = builder.build(&server).await?; + test.submit_turn_with_approval_and_permission_profile( + FIRST_PROMPT, + AskForApproval::OnRequest, + PermissionProfile::read_only(), + ) + .await?; + + let child_thread_id = test + .thread_manager + .list_thread_ids() + .await + .into_iter() + .find(|thread_id| *thread_id != test.session_configured.thread_id) + .expect("spawned child thread"); + let child = test.thread_manager.get_thread(child_thread_id).await?; + wait_for_event(&child, |event| { + matches!(event, EventMsg::ExecApprovalRequest(_)) + }) + .await; + assert!(!test.codex.try_shutdown_if_idle().await?); + let child = test + .thread_manager + .remove_thread(&child_thread_id) + .await + .expect("remove child while preserving its lifecycle"); + + assert!(!test.codex.try_shutdown_if_idle().await?); + + child.shutdown_and_wait().await?; + test.codex.shutdown_and_wait().await?; + Ok(()) +}