core: track thread publication lifecycle

codex-rs/core/src/agent/control.rs

@@ -15,6 +15,8 @@ use crate::session::emit_subagent_session_started;
 use crate::session_prefix::format_inter_agent_completion_message;
 use crate::session_prefix::format_subagent_context_line;
 use crate::session_prefix::format_subagent_notification_message;
+use crate::thread_activity::ThreadActivityGate;
+use crate::thread_activity::ThreadActivityHandle;
 use crate::thread_manager::ResumeThreadWithHistoryOptions;
 use crate::thread_manager::ThreadManagerState;
 use crate::thread_rollout_truncation::truncate_rollout_to_last_n_fork_turns;
@@ -103,6 +105,7 @@ pub(crate) struct AgentControl {
     state: Arc<AgentRegistry>,
     v2_residency: Arc<V2Residency>,
     agent_execution_limiter: Arc<AgentExecutionLimiter>,
+    thread_activity_gate: Arc<ThreadActivityGate>,
     /// Session-scoped state shared by the root thread and every cloned sub-agent control handle.
     rollout_budget: Arc<RolloutBudget>,
 }
@@ -113,8 +116,13 @@ impl AgentControl {
         manager: Weak<ThreadManagerState>,
         rollout_budget: Option<RolloutBudgetConfig>,
     ) -> Self {
+        let thread_activity_gate = manager
+            .upgrade()
+            .map(|manager| manager.thread_activity_gate())
+            .unwrap_or_default();
         let control = Self {
             manager,
+            thread_activity_gate,
             ..Default::default()
         };
         if let Some(rollout_budget) = rollout_budget {
@@ -133,6 +141,15 @@ impl AgentControl {
         self.session_id
     }
 
+    pub(crate) fn register_thread_activity(
+        &self,
+        thread_id: ThreadId,
+        parent_thread_id: Option<ThreadId>,
+    ) -> Option<ThreadActivityHandle> {
+        self.thread_activity_gate
+            .register(thread_id, parent_thread_id)
+    }
+
     pub(crate) fn rollout_budget(&self) -> &RolloutBudget {
         self.rollout_budget.as_ref()
     }

codex-rs/core/src/codex_delegate.rs

@@ -137,6 +137,7 @@ pub(crate) async fn run_codex_thread_interactive(
     }))
     .or_cancel(&cancel_token)
     .await??;
+    codex.session.thread_activity.mark_initialized();
     let thread_config = codex.thread_config_snapshot().await;
     let client_metadata = parent_session.app_server_client_metadata().await;
     emit_subagent_session_started(

codex-rs/core/src/lib.rs

@@ -111,6 +111,7 @@ pub use codex_protocol::config_types::ModelProviderAuthInfo;
 mod event_mapping;
 pub mod review_format;
 pub use codex_prompts as review_prompts;
+mod thread_activity;
 mod thread_manager;
 pub(crate) mod web_search;
 pub(crate) mod windows_sandbox_read_grants;

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

@@ -858,6 +858,7 @@ pub(super) async fn submission_loop(
             warn!("failed to shutdown thread persistence after submission channel closed: {err}");
         }
     }
+    sess.thread_activity.mark_closed();
     debug!("Agent loop exited");
 }
 

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

@@ -7,6 +7,7 @@ use crate::environment_selection::TurnEnvironmentSnapshot;
 use crate::shell_snapshot::ShellSnapshot;
 use crate::skills::SkillError;
 use crate::state::ActiveTurn;
+use crate::thread_activity::ThreadActivityHandle;
 use codex_extension_api::ExtensionDataInit;
 use codex_login::auth::AgentIdentityAuthPolicy;
 use codex_protocol::SessionId;
@@ -145,6 +146,7 @@ pub(crate) struct Session {
     pub(crate) conversation: Arc<RealtimeConversationManager>,
     pub(crate) active_turn: Mutex<Option<ActiveTurn>>,
     pub(crate) submission_lifecycle: SubmissionLifecycle,
+    pub(crate) thread_activity: ThreadActivityHandle,
     pub(crate) submission_send_lock: Mutex<()>,
     pub(crate) input_queue: InputQueue,
     pub(crate) guardian_review_session: GuardianReviewSessionManager,
@@ -630,6 +632,7 @@ impl Session {
         session_configuration.forked_from_thread_id = forked_from_id;
         let parent_thread_id = session_configuration
             .parent_thread_id
+            .or_else(|| session_configuration.session_source.parent_thread_id())
             .or_else(|| initial_history.get_resumed_parent_thread_id());
         session_configuration.parent_thread_id = parent_thread_id;
         let multi_agent_version = multi_agent_version.map(OnceLock::from).unwrap_or_default();
@@ -1238,6 +1241,10 @@ impl Session {
                 tool_search_handler_cache: Default::default(),
                 turn_environments: Arc::clone(&turn_environments),
             };
+            let thread_activity = services
+                .agent_control
+                .register_thread_activity(thread_id, parent_thread_id)
+                .ok_or_else(|| anyhow::anyhow!("thread or ancestor is shutting down"))?;
             let sess = Arc::new(Session {
                 thread_id,
                 installation_id,
@@ -1251,6 +1258,7 @@ impl Session {
                 conversation: Arc::new(RealtimeConversationManager::new()),
                 active_turn: Mutex::new(None),
                 submission_lifecycle: SubmissionLifecycle::default(),
+                thread_activity,
                 submission_send_lock: Mutex::new(()),
                 input_queue: InputQueue::new(),
                 guardian_review_session: GuardianReviewSessionManager::default(),

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

@@ -5507,6 +5507,10 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) {
         "turn_id".to_string(),
         skills_snapshot,
     );
+    let thread_activity = Arc::new(crate::thread_activity::ThreadActivityGate::default())
+        .register(thread_id, /*parent_thread_id*/ None)
+        .expect("test thread activity registration");
+    thread_activity.mark_initialized();
     let session = Session {
         thread_id,
         installation_id: "11111111-1111-4111-8111-111111111111".to_string(),
@@ -5520,6 +5524,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) {
         conversation: Arc::new(RealtimeConversationManager::new()),
         active_turn: Mutex::new(None),
         submission_lifecycle: SubmissionLifecycle::default(),
+        thread_activity,
         submission_send_lock: Mutex::new(()),
         input_queue: super::input_queue::InputQueue::new(),
         guardian_review_session: crate::guardian::GuardianReviewSessionManager::default(),
@@ -7783,6 +7788,10 @@ where
         "turn_id".to_string(),
         skills_snapshot,
     ));
+    let thread_activity = Arc::new(crate::thread_activity::ThreadActivityGate::default())
+        .register(thread_id, /*parent_thread_id*/ None)
+        .expect("test thread activity registration");
+    thread_activity.mark_initialized();
     let session = Arc::new(Session {
         thread_id,
         installation_id: "11111111-1111-4111-8111-111111111111".to_string(),
@@ -7796,6 +7805,7 @@ where
         conversation: Arc::new(RealtimeConversationManager::new()),
         active_turn: Mutex::new(None),
         submission_lifecycle: SubmissionLifecycle::default(),
+        thread_activity,
         submission_send_lock: Mutex::new(()),
         input_queue: super::input_queue::InputQueue::new(),
         guardian_review_session: crate::guardian::GuardianReviewSessionManager::default(),

codex-rs/core/src/thread_activity.rs

@@ -0,0 +1,216 @@
+use codex_protocol::ThreadId;
+use std::collections::HashMap;
+use std::collections::HashSet;
+use std::sync::Arc;
+use std::sync::Mutex;
+
+#[derive(Default)]
+pub(crate) struct ThreadActivityGate {
+    state: Mutex<ThreadActivityState>,
+}
+
+#[derive(Default)]
+struct ThreadActivityState {
+    next_generation: u64,
+    nodes: HashMap<ThreadId, ThreadActivityNode>,
+}
+
+struct ThreadActivityNode {
+    generation: u64,
+    parent_thread_id: Option<ThreadId>,
+    active: usize,
+    closing: bool,
+    handle_live: bool,
+    initializing: bool,
+}
+
+pub(crate) struct ThreadActivityHandle {
+    gate: Arc<ThreadActivityGate>,
+    thread_id: ThreadId,
+    generation: u64,
+}
+
+impl ThreadActivityGate {
+    pub(crate) fn register(
+        self: &Arc<Self>,
+        thread_id: ThreadId,
+        parent_thread_id: Option<ThreadId>,
+    ) -> Option<ThreadActivityHandle> {
+        let mut state = self
+            .state
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner);
+
+        if state
+            .nodes
+            .get(&thread_id)
+            .is_some_and(|node| !node.closing || node.active > 0)
+            || parent_thread_id.is_some_and(|parent_thread_id| {
+                Self::ancestor_blocks_registration(&state, parent_thread_id, thread_id)
+            })
+            || Self::has_active_descendant(&state, thread_id)
+        {
+            return None;
+        }
+
+        state.next_generation = state.next_generation.wrapping_add(1);
+        let generation = state.next_generation;
+        state.nodes.insert(
+            thread_id,
+            ThreadActivityNode {
+                generation,
+                parent_thread_id,
+                active: 1,
+                closing: false,
+                handle_live: true,
+                initializing: true,
+            },
+        );
+        Some(ThreadActivityHandle {
+            gate: Arc::clone(self),
+            thread_id,
+            generation,
+        })
+    }
+
+    fn ancestor_blocks_registration(
+        state: &ThreadActivityState,
+        mut ancestor_id: ThreadId,
+        descendant_id: ThreadId,
+    ) -> bool {
+        let mut visited = HashSet::new();
+        while visited.insert(ancestor_id) {
+            if ancestor_id == descendant_id {
+                return true;
+            }
+            let Some(node) = state.nodes.get(&ancestor_id) else {
+                return false;
+            };
+            if node.closing {
+                return true;
+            }
+            let Some(parent_thread_id) = node.parent_thread_id else {
+                return false;
+            };
+            ancestor_id = parent_thread_id;
+        }
+        true
+    }
+
+    fn has_active_descendant(state: &ThreadActivityState, ancestor_id: ThreadId) -> bool {
+        state.nodes.iter().any(|(thread_id, node)| {
+            if *thread_id == ancestor_id || node.active == 0 {
+                return false;
+            }
+            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 true;
+                }
+                parent_thread_id = state
+                    .nodes
+                    .get(&parent_id)
+                    .and_then(|node| node.parent_thread_id);
+            }
+            false
+        })
+    }
+
+    fn mark_closed(&self, thread_id: ThreadId, generation: u64) {
+        let mut state = self
+            .state
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner);
+        let Some(node) = state.nodes.get_mut(&thread_id) else {
+            return;
+        };
+        if node.generation == generation {
+            node.closing = true;
+        }
+    }
+
+    fn mark_initialized(&self, thread_id: ThreadId, generation: u64) {
+        let mut state = self
+            .state
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner);
+        let Some(node) = state.nodes.get_mut(&thread_id) else {
+            return;
+        };
+        if node.generation == generation && node.initializing {
+            node.initializing = false;
+            node.active = node.active.saturating_sub(1);
+        }
+    }
+
+    fn prune_inactive_orphaned(
+        state: &mut ThreadActivityState,
+        mut thread_id: ThreadId,
+        mut generation: u64,
+    ) {
+        loop {
+            let has_children = state
+                .nodes
+                .values()
+                .any(|node| node.parent_thread_id == Some(thread_id));
+            let Some(node) = state.nodes.get(&thread_id) else {
+                return;
+            };
+            if node.generation != generation || node.handle_live || node.active != 0 || has_children
+            {
+                return;
+            }
+            let parent_thread_id = node.parent_thread_id;
+            state.nodes.remove(&thread_id);
+            let Some(parent_id) = parent_thread_id else {
+                return;
+            };
+            let Some(parent) = state.nodes.get(&parent_id) else {
+                return;
+            };
+            thread_id = parent_id;
+            generation = parent.generation;
+        }
+    }
+}
+
+impl ThreadActivityHandle {
+    pub(crate) fn mark_initialized(&self) {
+        self.gate.mark_initialized(self.thread_id, self.generation);
+    }
+
+    pub(crate) fn mark_closed(&self) {
+        self.gate.mark_closed(self.thread_id, self.generation);
+    }
+}
+
+impl Drop for ThreadActivityHandle {
+    fn drop(&mut self) {
+        let mut state = self
+            .gate
+            .state
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner);
+        let Some(node) = state.nodes.get_mut(&self.thread_id) else {
+            return;
+        };
+        if node.generation == self.generation {
+            node.handle_live = false;
+            node.active = 0;
+            node.closing = true;
+            node.initializing = false;
+            ThreadActivityGate::prune_inactive_orphaned(
+                &mut state,
+                self.thread_id,
+                self.generation,
+            );
+        }
+    }
+}
+
+#[cfg(test)]
+#[path = "thread_activity_tests.rs"]
+mod tests;

codex-rs/core/src/thread_activity_tests.rs

@@ -0,0 +1,75 @@
+use super::*;
+use codex_protocol::ThreadId;
+use std::sync::Arc;
+
+fn register(
+    gate: &Arc<ThreadActivityGate>,
+    thread_id: ThreadId,
+    parent_thread_id: Option<ThreadId>,
+) -> ThreadActivityHandle {
+    let handle = gate
+        .register(thread_id, parent_thread_id)
+        .expect("register thread");
+    handle.mark_initialized();
+    handle
+}
+
+#[test]
+fn registration_rejects_parent_cycles() {
+    let gate = Arc::new(ThreadActivityGate::default());
+    let first_id = ThreadId::new();
+    let missing_parent_id = ThreadId::new();
+    let _first = register(&gate, first_id, Some(missing_parent_id));
+
+    assert!(gate.register(missing_parent_id, Some(first_id)).is_none());
+    let self_parent_id = ThreadId::new();
+    assert!(
+        gate.register(self_parent_id, Some(self_parent_id))
+            .is_none()
+    );
+}
+
+#[test]
+fn registration_and_publication_preserve_thread_incarnations() {
+    let gate = Arc::new(ThreadActivityGate::default());
+    let parent_id = ThreadId::new();
+    let child_id = ThreadId::new();
+    let parent = register(&gate, parent_id, /*parent_thread_id*/ None);
+    let child = gate
+        .register(child_id, Some(parent_id))
+        .expect("register initializing child");
+
+    parent.mark_closed();
+    assert!(
+        gate.register(parent_id, /*parent_thread_id*/ None)
+            .is_none()
+    );
+    assert!(gate.register(ThreadId::new(), Some(parent_id)).is_none());
+
+    child.mark_initialized();
+    let replacement = gate
+        .register(parent_id, /*parent_thread_id*/ None)
+        .expect("register replacement parent");
+    replacement.mark_initialized();
+    drop(parent);
+    assert!(
+        gate.register(parent_id, /*parent_thread_id*/ None)
+            .is_none()
+    );
+}
+
+#[test]
+fn dropped_tree_handles_are_pruned_leaf_to_root() {
+    let gate = Arc::new(ThreadActivityGate::default());
+    let parent_id = ThreadId::new();
+    let child_id = ThreadId::new();
+    let parent = register(&gate, parent_id, /*parent_thread_id*/ None);
+    let child = register(&gate, child_id, Some(parent_id));
+
+    parent.mark_closed();
+    child.mark_closed();
+    drop(parent);
+    assert_eq!(gate.state.lock().expect("gate state").nodes.len(), 2);
+    drop(child);
+    assert!(gate.state.lock().expect("gate state").nodes.is_empty());
+}

codex-rs/core/src/thread_manager.rs

@@ -250,6 +250,7 @@ pub(crate) struct ThreadManagerState {
     user_instructions_provider: Arc<dyn UserInstructionsProvider>,
     thread_store: Arc<dyn ThreadStore>,
     agent_graph_store: Option<Arc<dyn AgentGraphStore>>,
+    thread_activity_gate: Arc<crate::thread_activity::ThreadActivityGate>,
     attestation_provider: Option<Arc<dyn AttestationProvider>>,
     external_time_provider: Option<Arc<dyn TimeProvider>>,
     session_source: SessionSource,
@@ -351,6 +352,9 @@ impl ThreadManager {
                 user_instructions_provider,
                 thread_store,
                 agent_graph_store,
+                thread_activity_gate: Arc::new(
+                    crate::thread_activity::ThreadActivityGate::default(),
+                ),
                 attestation_provider,
                 external_time_provider,
                 auth_manager,
@@ -459,6 +463,9 @@ impl ThreadManager {
                 ),
                 thread_store,
                 agent_graph_store,
+                thread_activity_gate: Arc::new(
+                    crate::thread_activity::ThreadActivityGate::default(),
+                ),
                 attestation_provider: None,
                 external_time_provider: None,
                 auth_manager,
@@ -1074,6 +1081,10 @@ impl ThreadManager {
 }
 
 impl ThreadManagerState {
+    pub(crate) fn thread_activity_gate(&self) -> Arc<crate::thread_activity::ThreadActivityGate> {
+        Arc::clone(&self.thread_activity_gate)
+    }
+
     pub(crate) fn agent_graph_store(&self) -> Option<Arc<dyn AgentGraphStore>> {
         self.agent_graph_store.clone()
     }
@@ -1646,6 +1657,7 @@ impl ThreadManagerState {
                     session_source,
                 ));
                 e.insert(thread.clone());
+                thread.codex.session.thread_activity.mark_initialized();
                 return Ok(NewThread {
                     thread_id,
                     thread,