copyberry · Oct 1, 2026 · Oct 1, 2026
diff --git a/codex-rs/core-api/src/lib.rs b/codex-rs/core-api/src/lib.rs
@@ -28,6 +28,7 @@ pub use codex_config::types::TuiKeymap;
 pub use codex_config::types::TuiNotificationSettings;
 pub use codex_config::types::TuiPetAnchor;
 pub use codex_config::types::UriBasedFileOpener;
+pub use codex_core::AgentTreeShutdown;
 pub use codex_core::CodexAppsToolsCache;
 pub use codex_core::CodexThread;
 pub use codex_core::CodexThreadSettingsOverrides;

diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs
@@ -58,6 +58,9 @@ use tracing::warn;
 use uuid::Uuid;
 
 pub(crate) use self::runtime::AgentControlInit;
+pub(crate) use self::runtime::AgentTreeMembership;
+pub(crate) use self::runtime::AgentTreeShutdownState;
+pub(crate) use self::runtime::AgentTreeTeardownGuard;
 pub(crate) use self::runtime::LocalAgentRuntime;
 pub(crate) use self::watch::StatusSubscription;
 
@@ -436,12 +439,20 @@ impl LocalAgentControl {
         else {
             return;
         };
+        let Ok(membership) = self.runtime.admit_start() else {
+            return;
+        };
+        let teardown = membership.into_teardown_guard();
         let control = self.clone();
-        tokio::spawn(async move {
+        let watcher = async move {
             let status = match control.subscribe_status(child_thread_id).await {
                 Ok(mut updates) => {
                     let mut final_status = None;
-                    while let Some(Ok(snapshot)) = updates.next().await {
+                    while let Some(Ok(snapshot)) = tokio::select! {
+                        biased;
+                        _ = control.runtime.shutdown.cancelled() => return,
+                        update = updates.next() => update,
+                    } {
                         if let Some(status) = snapshot.status()
                             && is_final(status)
                         {
@@ -456,7 +467,11 @@ impl LocalAgentControl {
                 }
                 Err(_) => control.get_status(child_thread_id).await,
             };
-            if !is_final(&status) {
+            // Tree shutdown is a lifecycle handoff, not a child result for the parent.
+            if !is_final(&status)
+                || (matches!(&status, AgentStatus::Shutdown)
+                    && control.runtime.shutdown.is_cancelled())
+            {
                 return;
             }
 
@@ -516,6 +531,10 @@ impl LocalAgentControl {
                     status,
                 ))
                 .await;
+        };
+        tokio::spawn(async move {
+            watcher.await;
+            teardown.complete();
         });
     }
 

diff --git a/codex-rs/core/src/agent/control/residency.rs b/codex-rs/core/src/agent/control/residency.rs
@@ -1,5 +1,6 @@
 use super::LocalAgentControl;
 use super::LocalAgentRuntime;
+use super::runtime::AgentTreeMembership;
 use crate::agent::AgentStatus;
 use crate::codex_thread::CodexThread;
 use crate::config::Config;
@@ -51,13 +52,14 @@ impl LocalAgentControl {
         &self,
         state: &Arc<ThreadManagerState>,
         config: &Config,
+        membership: &AgentTreeMembership,
         protected_thread_id: Option<ThreadId>,
     ) -> CodexResult<V2ResidencySlot> {
         let capacity = config
             .effective_agent_max_threads(MultiAgentVersion::V2)
             .unwrap_or(usize::MAX);
         Arc::clone(&self.runtime.residency)
-            .reserve_slot(state, capacity, protected_thread_id)
+            .reserve_slot(state, capacity, membership, protected_thread_id)
             .await
     }
 
@@ -101,6 +103,7 @@ impl V2Residency {
         self: Arc<Self>,
         manager: &Arc<ThreadManagerState>,
         capacity: usize,
+        membership: &AgentTreeMembership,
         protected_thread_id: Option<ThreadId>,
     ) -> CodexResult<V2ResidencySlot> {
         loop {
@@ -111,7 +114,7 @@ impl V2Residency {
                 });
             }
             if !self
-                .try_unload_one_resident(manager, protected_thread_id)
+                .try_unload_one_resident(manager, membership, protected_thread_id)
                 .await
             {
                 return Err(CodexErr::new(CodexErrorDetails::AgentLimitReached {
@@ -136,6 +139,7 @@ impl V2Residency {
     async fn try_unload_one_resident(
         self: &Arc<Self>,
         manager: &Arc<ThreadManagerState>,
+        membership: &AgentTreeMembership,
         protected_thread_id: Option<ThreadId>,
     ) -> bool {
         // Keep shutting-down workers counted until removal. Each runtime's write guard
@@ -178,13 +182,15 @@ impl V2Residency {
             // must keep delivery excluded and capacity reserved through registry removal.
             let manager = Arc::clone(manager);
             let residency = Arc::clone(self);
+            let teardown = membership.clone().into_teardown_guard();
             let eviction = tokio::spawn(async move {
                 let _residency_guard = residency_guard;
                 candidate_thread.ensure_rollout_materialized().await;
                 if let Err(err) = candidate_thread.shutdown_and_wait().await {
                     warn!(
                         "failed to shut down v2 resident thread before unloading {candidate_thread_id}: {err}"
                     );
+                    teardown.complete();
                     return false;
                 }
                 let environments = candidate_thread.environment_selections().await;
@@ -193,6 +199,7 @@ impl V2Residency {
                     .get(&candidate_thread_id)
                     .is_some_and(|registered| !Arc::ptr_eq(registered, &candidate_thread))
                 {
+                    teardown.complete();
                     return false;
                 }
                 candidate_thread
@@ -204,6 +211,7 @@ impl V2Residency {
                 // Keep publication excluded until both entries have been removed.
                 threads.remove(&candidate_thread_id);
                 residency.remove(candidate_thread_id);
+                teardown.complete();
                 true
             });
             match eviction.await {

diff --git a/codex-rs/core/src/agent/control/residency_tests.rs b/codex-rs/core/src/agent/control/residency_tests.rs
@@ -42,9 +42,15 @@ async fn residency_slot_reservation_unloads_oldest_idle_v2_agent() {
         .runtime
         .upgrade()
         .expect("thread manager should be live");
+    let membership = control.runtime.admit_start().expect("admit residency work");
 
     let first_slot = control
-        .reserve_v2_residency_slot(&state, &config, /*protected_thread_id*/ None)
+        .reserve_v2_residency_slot(
+            &state,
+            &config,
+            &membership,
+            /*protected_thread_id*/ None,
+        )
         .await
         .expect("first resident slot");
     let first =
@@ -53,7 +59,12 @@ async fn residency_slot_reservation_unloads_oldest_idle_v2_agent() {
     mark_thread_completed(first.thread.as_ref()).await;
 
     let second_slot = control
-        .reserve_v2_residency_slot(&state, &config, /*protected_thread_id*/ None)
+        .reserve_v2_residency_slot(
+            &state,
+            &config,
+            &membership,
+            /*protected_thread_id*/ None,
+        )
         .await
         .expect("second resident slot should evict the first idle agent");
     match manager.get_thread(first.thread_id).await {
@@ -93,9 +104,15 @@ async fn interrupted_v2_agent_is_lost_after_residency_eviction() {
         .runtime
         .upgrade()
         .expect("thread manager should be live");
+    let membership = control.runtime.admit_start().expect("admit residency work");
 
     let first_slot = control
-        .reserve_v2_residency_slot(&state, &config, /*protected_thread_id*/ None)
+        .reserve_v2_residency_slot(
+            &state,
+            &config,
+            &membership,
+            /*protected_thread_id*/ None,
+        )
         .await
         .expect("first resident slot");
     let first =
@@ -104,7 +121,12 @@ async fn interrupted_v2_agent_is_lost_after_residency_eviction() {
     mark_thread_interrupted(first.thread.as_ref()).await;
 
     let second_slot = control
-        .reserve_v2_residency_slot(&state, &config, /*protected_thread_id*/ None)
+        .reserve_v2_residency_slot(
+            &state,
+            &config,
+            &membership,
+            /*protected_thread_id*/ None,
+        )
         .await
         .expect("second resident slot should evict the first interrupted idle agent");
     match manager.get_thread(first.thread_id).await {

diff --git a/codex-rs/core/src/agent/control/runtime.rs b/codex-rs/core/src/agent/control/runtime.rs
@@ -14,9 +14,81 @@ use arc_swap::ArcSwapOption;
 use codex_extension_api::ThreadInstructionsProvider;
 use codex_protocol::SessionId;
 use codex_protocol::ThreadId;
+use codex_protocol::error::CodexErr;
+use codex_protocol::error::Result as CodexResult;
 use std::sync::Arc;
 use std::sync::OnceLock;
 use std::sync::Weak;
+use std::sync::atomic::AtomicBool;
+use std::sync::atomic::Ordering;
+use tokio_util::sync::CancellationToken;
+use tokio_util::task::TaskTracker;
+use tokio_util::task::task_tracker::TaskTrackerToken;
+
+#[derive(Debug, Default)]
+pub(crate) struct AgentTreeShutdownState {
+    members: TaskTracker,
+    failed: AtomicBool,
+}
+
+impl AgentTreeShutdownState {
+    pub(crate) async fn wait(&self) -> CodexResult<()> {
+        self.members.wait().await;
+        if self.failed.load(Ordering::Acquire) {
+            return Err(CodexErr::Fatal(
+                "agent tree shutdown did not complete cleanly".to_owned(),
+            ));
+        }
+        Ok(())
+    }
+
+    fn record_failure(&self) {
+        self.failed.store(true, Ordering::Release);
+    }
+}
+
+#[derive(Clone)]
+pub(crate) struct AgentTreeMembership {
+    state: Arc<AgentTreeShutdownState>,
+    _member: TaskTrackerToken,
+}
+
+impl AgentTreeMembership {
+    pub(crate) fn into_teardown_guard(self) -> AgentTreeTeardownGuard {
+        AgentTreeTeardownGuard {
+            membership: self,
+            completed: false,
+        }
+    }
+}
+
+/// Marks tree shutdown as failed if teardown work exits without completing.
+pub(crate) struct AgentTreeTeardownGuard {
+    membership: AgentTreeMembership,
+    completed: bool,
+}
+
+impl AgentTreeTeardownGuard {
+    pub(crate) fn clone_for_teardown(&self) -> Self {
+        self.membership.clone().into_teardown_guard()
+    }
+
+    pub(crate) fn record_shutdown_failure(&self) {
+        self.membership.state.record_failure();
+    }
+
+    pub(crate) fn complete(mut self) {
+        self.completed = true;
+    }
+}
+
+impl Drop for AgentTreeTeardownGuard {
+    fn drop(&mut self) {
+        if !self.completed {
+            self.record_shutdown_failure();
+        }
+    }
+}
 
 /// Local tree state, kept separate from the shared agent operation interface.
 #[derive(Clone)]
@@ -37,6 +109,9 @@ pub(crate) struct LocalAgentRuntime {
         Arc<OnceLock<Arc<dyn ThreadInstructionsProvider>>>,
     pub(super) registry: Arc<AgentRegistry>,
     pub(super) residency: Arc<V2Residency>,
+    /// Shared by every session in this tree, including private delegates.
+    pub(crate) shutdown: CancellationToken,
+    shutdown_state: Arc<AgentTreeShutdownState>,
 }
 
 impl LocalAgentRuntime {
@@ -50,6 +125,8 @@ impl LocalAgentRuntime {
             thread_id_generator,
             registry: Arc::default(),
             residency: Arc::default(),
+            shutdown: CancellationToken::new(),
+            shutdown_state: Arc::default(),
             agent_execution_limiter: Arc::default(),
             rollout_budget: Arc::default(),
             root_service_tier: Arc::new(ArcSwapOption::from(None)),
@@ -104,6 +181,36 @@ impl AgentControlInit {
 }
 
 impl LocalAgentRuntime {
+    pub(crate) fn admit_start(&self) -> CodexResult<AgentTreeMembership> {
+        if self.shutdown_state.members.is_closed() {
+            return Err(CodexErr::InvalidRequest(
+                "agent runtime is shutting down".to_owned(),
+            ));
+        }
+        let membership = AgentTreeMembership {
+            state: Arc::clone(&self.shutdown_state),
+            _member: self.shutdown_state.members.token(),
+        };
+        // Closing a TaskTracker does not reject new tokens. Recheck so a start racing with
+        // shutdown is either admitted before the fence or rejected after it.
+        if self.shutdown_state.members.is_closed() {
+            return Err(CodexErr::InvalidRequest(
+                "agent runtime is shutting down".to_owned(),
+            ));
+        }
+        Ok(membership)
+    }
+
+    pub(crate) fn request_shutdown(&self) -> Arc<AgentTreeShutdownState> {
+        self.shutdown_state.members.close();
+        self.shutdown.cancel();
+        Arc::clone(&self.shutdown_state)
+    }
+
+    pub(crate) fn record_shutdown_failure(&self) {
+        self.shutdown_state.record_failure();
+    }
+
     pub(crate) fn generate_thread_id(&self) -> ThreadId {
         (self.thread_id_generator)()
     }
@@ -128,3 +235,7 @@ impl LocalAgentRuntime {
         provider
     }
 }
+
+#[cfg(test)]
+#[path = "runtime_tests.rs"]
+mod tests;
diff --git a/codex-rs/core/src/agent/control/runtime_tests.rs b/codex-rs/core/src/agent/control/runtime_tests.rs
@@ -0,0 +1,29 @@
+use super::*;
+use crate::thread_manager::default_thread_id_generator;
+use futures::future;
+
+#[tokio::test]
+async fn aborted_teardown_fails_tree_shutdown() {
+    let runtime = LocalAgentRuntime::new(
+        Weak::default(),
+        default_thread_id_generator(),
+        /*rollout_budget*/ None,
+    );
+    let teardown = runtime
+        .admit_start()
+        .expect("teardown should be admitted")
+        .into_teardown_guard();
+    let task = tokio::spawn(async move {
+        let _teardown = teardown;
+        future::pending::<()>().await;
+    });
+    let shutdown = runtime.request_shutdown();
+
+    task.abort();
+    task.await.expect_err("teardown task should be aborted");
+
+    shutdown
+        .wait()
+        .await
+        .expect_err("aborted teardown must fail tree shutdown");
+}