core: preserve thread activity through completion

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<AgentStatus>,
+    _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<SessionSource>,
         child_reference: String,
         child_agent_path: Option<AgentPath>,
+        completion_child: Option<CompletionWatcherChild>,
     ) {
         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>) -> AgentStatus {
+    let mut current = status.borrow().clone();
+    while !is_final(&current) {
+        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()

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!(
+                &notification_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!(
+                &notification_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(

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);

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);

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<session::SessionActivityReservation<'_>> {

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<ThreadActivityReservation> {
+        self.thread_activity.try_reserve(/*close*/ false)
+    }
+
+    pub(crate) fn try_reserve_dispatched_turn_activity(&self) -> Option<ThreadActivityReservation> {
+        self.thread_activity.try_reserve_dispatched_turn()
+    }
+
     pub(crate) fn try_reserve_idle_shutdown(&self) -> Option<SessionActivityReservation<'_>> {
         let thread = self.thread_activity.try_reserve_idle_shutdown()?;
         let submission = self.submission_lifecycle.try_reserve_idle_shutdown()?;

codex-rs/core/src/session/tests.rs

@@ -10194,7 +10194,7 @@ async fn try_start_turn_if_idle_rejects_active_turn_without_injecting() {
         Vec::<TurnInput>::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]

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<RunningTask>,
     pub(crate) turn_state: Arc<Mutex<TurnState>>,
+    pub(crate) thread_activity: Option<ThreadActivityReservation>,
 }
 
 /// 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,
         }
     }
 }