core: make thread activity lifecycle atomic

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

@@ -792,7 +792,7 @@ impl Codex {
         }
         let is_shutdown = matches!(&sub.op, Op::Shutdown);
         let _send_guard = self.session.submission_send_lock.lock().await;
-        let Some(reservation) = self.session.submission_lifecycle.try_reserve(is_shutdown) else {
+        let Some(mut reservation) = self.session.try_reserve_submission(is_shutdown) else {
             if is_shutdown {
                 return Ok(());
             }
@@ -800,6 +800,9 @@ impl Codex {
                 "thread is shutting down".to_string(),
             ));
         };
+        if !reservation.prepare_delivery() {
+            return Err(CodexErr::InternalAgentDied);
+        }
         if self.tx_sub.send(sub).await.is_err() {
             reservation.release_after_failed_delivery();
             return Err(CodexErr::InternalAgentDied);
@@ -825,9 +828,12 @@ impl Codex {
         if self.session.submission_lifecycle.is_closing() {
             return Ok(true);
         }
-        let Some(reservation) = self.session.try_claim_idle_shutdown().await else {
+        let Some(mut reservation) = self.session.try_claim_idle_shutdown().await else {
             return Ok(false);
         };
+        if !reservation.prepare_idle_shutdown_delivery() {
+            return Ok(false);
+        }
         let sub = Submission {
             id: new_submission_id(),
             op: Op::Shutdown,
@@ -1053,11 +1059,11 @@ impl Session {
 
     pub(crate) async fn try_claim_idle_shutdown(
         &self,
-    ) -> Option<session::SubmissionReservation<'_>> {
+    ) -> Option<session::SessionActivityReservation<'_>> {
         if self.has_idle_shutdown_blocker().await {
             return None;
         }
-        let reservation = self.submission_lifecycle.try_reserve_idle_shutdown()?;
+        let reservation = self.try_reserve_idle_shutdown()?;
         if self.has_idle_shutdown_blocker().await {
             return None;
         }

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

@@ -8,6 +8,7 @@ use crate::shell_snapshot::ShellSnapshot;
 use crate::skills::SkillError;
 use crate::state::ActiveTurn;
 use crate::thread_activity::ThreadActivityHandle;
+use crate::thread_activity::ThreadActivityReservation;
 use codex_extension_api::ExtensionDataInit;
 use codex_login::auth::AgentIdentityAuthPolicy;
 use codex_protocol::SessionId;
@@ -108,12 +109,37 @@ pub(crate) struct SubmissionReservation<'a> {
     active: bool,
 }
 
+pub(crate) struct SessionActivityReservation<'a> {
+    submission: SubmissionReservation<'a>,
+    thread: ThreadActivityReservation,
+}
+
+impl SessionActivityReservation<'_> {
+    pub(crate) fn prepare_delivery(&mut self) -> bool {
+        self.thread.prepare_delivery()
+    }
+
+    pub(crate) fn prepare_idle_shutdown_delivery(&mut self) -> bool {
+        self.thread.prepare_idle_shutdown_delivery()
+    }
+
+    pub(crate) fn commit(self) {
+        self.submission.commit();
+        self.thread.commit();
+    }
+
+    pub(crate) fn release_after_failed_delivery(self) {
+        self.submission.release_after_failed_delivery();
+        self.thread.release_after_failed_delivery();
+    }
+}
+
 impl SubmissionReservation<'_> {
     pub(crate) fn commit(mut self) {
         self.active = false;
     }
 
-    pub(crate) fn release_after_failed_delivery(mut self) {
+    fn release_after_failed_delivery(mut self) {
         if self.active {
             self.lifecycle.release(/*clear_closing*/ false);
             self.active = false;
@@ -572,12 +598,28 @@ impl Session {
         self.thread_id
     }
 
-    pub(crate) fn try_reserve_activity(&self) -> Option<SubmissionReservation<'_>> {
-        self.submission_lifecycle.try_reserve(/*close*/ false)
+    pub(crate) fn try_reserve_submission(
+        &self,
+        close: bool,
+    ) -> Option<SessionActivityReservation<'_>> {
+        let thread = self.thread_activity.try_reserve(close)?;
+        let submission = self.submission_lifecycle.try_reserve(close)?;
+        Some(SessionActivityReservation { submission, thread })
+    }
+
+    pub(crate) fn try_reserve_activity(&self) -> Option<SessionActivityReservation<'_>> {
+        self.try_reserve_submission(/*close*/ false)
+    }
+
+    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()?;
+        Some(SessionActivityReservation { submission, thread })
     }
 
     pub(crate) fn finish_submission(&self) {
         self.submission_lifecycle.finish_submission();
+        self.thread_activity.finish_submission();
     }
 
     /// Returns the identity shared by the root thread and all descendant threads.

codex-rs/core/src/thread_activity.rs

@@ -20,7 +20,10 @@ struct ThreadActivityNode {
     // A lineage edge follows the stable logical thread ID across runtime generations.
     parent_thread_id: Option<ThreadId>,
     active: usize,
+    committed: usize,
     closing: bool,
+    closed: bool,
+    unpublished: bool,
     handle_live: bool,
     initializing: bool,
 }
@@ -35,6 +38,15 @@ pub(crate) struct ThreadActivityHandle {
     generation: u64,
 }
 
+pub(crate) struct ThreadActivityReservation {
+    gate: Arc<ThreadActivityGate>,
+    thread_id: ThreadId,
+    generation: u64,
+    clear_closing_on_drop: bool,
+    delivery_prepared: bool,
+    active: bool,
+}
+
 impl ThreadActivityGate {
     pub(crate) fn register(
         self: &Arc<Self>,
@@ -66,7 +78,10 @@ impl ThreadActivityGate {
                 generation,
                 parent_thread_id,
                 active: 1,
+                committed: 0,
                 closing: false,
+                closed: false,
+                unpublished: false,
                 handle_live: true,
                 initializing: true,
             },
@@ -124,6 +139,170 @@ impl ThreadActivityGate {
         })
     }
 
+    fn validate_ancestors(
+        state: &ThreadActivityState,
+        mut parent_thread_id: Option<ThreadId>,
+        allow_closing: bool,
+    ) -> bool {
+        let mut visited = HashSet::new();
+        while let Some(parent_id) = parent_thread_id
+            && visited.insert(parent_id)
+        {
+            let Some(node) = state.nodes.get(&parent_id) else {
+                return true;
+            };
+            if (node.closing || node.initializing) && !allow_closing {
+                return false;
+            }
+            parent_thread_id = node.parent_thread_id;
+        }
+        parent_thread_id.is_none()
+    }
+
+    fn try_reserve(
+        self: &Arc<Self>,
+        thread_id: ThreadId,
+        generation: u64,
+        close: bool,
+    ) -> Option<ThreadActivityReservation> {
+        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.closing
+            || (node.initializing && !close)
+            || !Self::validate_ancestors(&state, node.parent_thread_id, close)
+        {
+            return None;
+        }
+        let node = state.nodes.get_mut(&thread_id)?;
+        node.active = node.active.checked_add(1)?;
+        node.closing = close;
+        Some(ThreadActivityReservation {
+            gate: Arc::clone(self),
+            thread_id,
+            generation,
+            clear_closing_on_drop: close,
+            delivery_prepared: false,
+            active: true,
+        })
+    }
+
+    fn try_reserve_idle_shutdown(
+        self: &Arc<Self>,
+        thread_id: ThreadId,
+        generation: u64,
+    ) -> Option<ThreadActivityReservation> {
+        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.closing
+            || !Self::validate_ancestors(&state, node.parent_thread_id, /*allow_closing*/ true)
+            || node.active != 0
+            || Self::has_active_descendant(&state, thread_id)
+        {
+            return None;
+        }
+        let node = state.nodes.get_mut(&thread_id)?;
+        node.active = 1;
+        node.closing = true;
+        Some(ThreadActivityReservation {
+            gate: Arc::clone(self),
+            thread_id,
+            generation,
+            clear_closing_on_drop: true,
+            delivery_prepared: false,
+            active: true,
+        })
+    }
+
+    fn release(&self, thread_id: ThreadId, generation: u64, clear_closing: bool) {
+        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.active == 0 {
+            return;
+        }
+        node.active -= 1;
+        if clear_closing && !node.closed {
+            node.closing = false;
+        }
+        if node.active == 0 && !node.handle_live {
+            Self::prune_inactive_orphaned(&mut state, thread_id, generation);
+        }
+    }
+
+    fn prepare_delivery(&self, thread_id: ThreadId, generation: u64) -> bool {
+        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 false;
+        };
+        if node.generation != generation || node.closed {
+            return false;
+        }
+        let Some(committed) = node.committed.checked_add(1) else {
+            return false;
+        };
+        if committed > node.active {
+            return false;
+        }
+        node.committed = committed;
+        true
+    }
+
+    fn rollback_delivery(&self, thread_id: ThreadId, generation: u64, clear_closing: bool) {
+        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.committed == 0 || node.active == 0 {
+            return;
+        }
+        node.committed -= 1;
+        node.active -= 1;
+        if clear_closing && !node.closed {
+            node.closing = false;
+        }
+        if node.active == 0 && !node.handle_live {
+            Self::prune_inactive_orphaned(&mut state, thread_id, generation);
+        }
+    }
+
+    fn finish_submission(&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.committed == 0 {
+            return;
+        }
+        node.committed -= 1;
+        node.active = node.active.saturating_sub(1);
+        if node.active == 0 && !node.handle_live {
+            Self::prune_inactive_orphaned(&mut state, thread_id, generation);
+        }
+    }
+
     fn mark_closed(&self, thread_id: ThreadId, generation: u64) {
         let mut state = self
             .state
@@ -134,9 +313,55 @@ impl ThreadActivityGate {
         };
         if node.generation == generation {
             node.closing = true;
+            node.closed = true;
+            if node.unpublished {
+                node.unpublished = false;
+                node.active = node.active.saturating_sub(1);
+            }
         }
     }
 
+    fn mark_unpublished(&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.closed || node.unpublished {
+            return;
+        }
+        let Some(active) = node.active.checked_add(1) else {
+            return;
+        };
+        node.active = active;
+        node.unpublished = true;
+    }
+
+    fn prepare_idle_shutdown_delivery(&self, thread_id: ThreadId, generation: u64) -> bool {
+        let mut state = self
+            .state
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner);
+        let Some(node) = state.nodes.get(&thread_id) else {
+            return false;
+        };
+        if node.generation != generation
+            || !node.closing
+            || node.closed
+            || node.active != 1
+            || Self::has_active_descendant(&state, thread_id)
+        {
+            return false;
+        }
+        let Some(node) = state.nodes.get_mut(&thread_id) else {
+            return false;
+        };
+        node.committed = 1;
+        true
+    }
+
     fn mark_initialized(&self, thread_id: ThreadId, generation: u64) {
         let mut state = self
             .state
@@ -187,9 +412,27 @@ impl ThreadActivityHandle {
         self.gate.mark_initialized(self.thread_id, self.generation);
     }
 
+    pub(crate) fn try_reserve(&self, close: bool) -> Option<ThreadActivityReservation> {
+        self.gate
+            .try_reserve(self.thread_id, self.generation, close)
+    }
+
+    pub(crate) fn try_reserve_idle_shutdown(&self) -> Option<ThreadActivityReservation> {
+        self.gate
+            .try_reserve_idle_shutdown(self.thread_id, self.generation)
+    }
+
+    pub(crate) fn finish_submission(&self) {
+        self.gate.finish_submission(self.thread_id, self.generation);
+    }
+
     pub(crate) fn mark_closed(&self) {
         self.gate.mark_closed(self.thread_id, self.generation);
     }
+
+    pub(crate) fn mark_unpublished(&self) {
+        self.gate.mark_unpublished(self.thread_id, self.generation);
+    }
 }
 
 impl Drop for ThreadActivityHandle {
@@ -204,9 +447,20 @@ impl Drop for ThreadActivityHandle {
         };
         if node.generation == self.generation {
             node.handle_live = false;
-            node.active = 0;
             node.closing = true;
-            node.initializing = false;
+            node.closed = true;
+            // The loop cannot finish committed submissions after its session handle disappears.
+            // Live RAII reservations (notably parent completion delivery) must drain themselves.
+            node.active = node.active.saturating_sub(node.committed);
+            node.committed = 0;
+            if node.initializing {
+                node.active = node.active.saturating_sub(1);
+                node.initializing = false;
+            }
+            if node.unpublished {
+                node.active = node.active.saturating_sub(1);
+                node.unpublished = false;
+            }
             ThreadActivityGate::prune_inactive_orphaned(
                 &mut state,
                 self.thread_id,
@@ -216,6 +470,65 @@ impl Drop for ThreadActivityHandle {
     }
 }
 
+impl ThreadActivityReservation {
+    pub(crate) fn prepare_delivery(&mut self) -> bool {
+        if !self.active || self.delivery_prepared {
+            return false;
+        }
+        if self.gate.prepare_delivery(self.thread_id, self.generation) {
+            self.delivery_prepared = true;
+            true
+        } else {
+            false
+        }
+    }
+
+    pub(crate) fn prepare_idle_shutdown_delivery(&mut self) -> bool {
+        if !self.active || self.delivery_prepared {
+            return false;
+        }
+        if self
+            .gate
+            .prepare_idle_shutdown_delivery(self.thread_id, self.generation)
+        {
+            self.delivery_prepared = true;
+            true
+        } else {
+            false
+        }
+    }
+
+    pub(crate) fn commit(mut self) {
+        if self.delivery_prepared {
+            self.active = false;
+        }
+    }
+
+    pub(crate) fn release_after_failed_delivery(mut self) {
+        self.rollback(/*clear_closing*/ false);
+    }
+
+    fn rollback(&mut self, clear_closing: bool) {
+        if !self.active {
+            return;
+        }
+        if self.delivery_prepared {
+            self.gate
+                .rollback_delivery(self.thread_id, self.generation, clear_closing);
+        } else {
+            self.gate
+                .release(self.thread_id, self.generation, clear_closing);
+        }
+        self.active = false;
+    }
+}
+
+impl Drop for ThreadActivityReservation {
+    fn drop(&mut self) {
+        self.rollback(self.clear_closing_on_drop);
+    }
+}
+
 #[cfg(test)]
 #[path = "thread_activity_tests.rs"]
 mod tests;

codex-rs/core/src/thread_activity_tests.rs

@@ -77,3 +77,90 @@ fn dropped_tree_handles_are_pruned_leaf_to_root() {
     drop(child);
     assert!(gate.state.lock().expect("gate state").nodes.is_empty());
 }
+
+#[test]
+fn unpublication_holds_activity_through_a_provisional_close() {
+    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));
+
+    let provisional_close = child
+        .try_reserve_idle_shutdown()
+        .expect("reserve provisional child shutdown");
+    child.mark_unpublished();
+    drop(provisional_close);
+
+    assert!(parent.try_reserve_idle_shutdown().is_none());
+
+    child.mark_closed();
+    assert!(parent.try_reserve_idle_shutdown().is_some());
+    assert!(child.try_reserve(/*close*/ false).is_none());
+
+    let child = register(&gate, ThreadId::new(), Some(parent_id));
+    let mut parent_close = parent
+        .try_reserve_idle_shutdown()
+        .expect("reserve provisional parent shutdown");
+    child.mark_unpublished();
+    assert!(!parent_close.prepare_idle_shutdown_delivery());
+    drop(parent_close);
+    assert!(parent.try_reserve_idle_shutdown().is_none());
+}
+
+#[test]
+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 parent = register(&gate, parent_id, /*parent_thread_id*/ None);
+    let child = register(&gate, child_id, Some(parent_id));
+
+    let child_activity = child
+        .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");
+    assert!(parent.try_reserve_idle_shutdown().is_none());
+    drop(parent);
+    assert!(gate.register(parent_id, /*parent_thread_id*/ None).is_err());
+    drop(parent_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());
+    drop(shutdown);
+
+    let mut failed_shutdown = parent
+        .try_reserve(/*close*/ true)
+        .expect("reserve explicit shutdown");
+    assert!(failed_shutdown.prepare_delivery());
+    failed_shutdown.release_after_failed_delivery();
+    {
+        let state = gate.state.lock().expect("gate state");
+        let node = state.nodes.get(&parent_id).expect("parent node");
+        assert_eq!((node.active, node.committed, node.closing), (0, 0, true));
+    }
+}
+
+#[test]
+fn receiver_can_finish_before_sender_observes_delivery() {
+    let gate = Arc::new(ThreadActivityGate::default());
+    let thread_id = ThreadId::new();
+    let thread = register(&gate, thread_id, /*parent_thread_id*/ None);
+    let mut delivery = thread
+        .try_reserve(/*close*/ false)
+        .expect("reserve submission delivery");
+    assert!(delivery.prepare_delivery());
+
+    thread.finish_submission();
+    delivery.commit();
+
+    let state = gate.state.lock().expect("gate state");
+    let node = state.nodes.get(&thread_id).expect("thread node");
+    assert_eq!((node.active, node.committed), (0, 0));
+}

codex-rs/core/src/thread_manager.rs

@@ -885,7 +885,7 @@ impl ThreadManager {
     /// as `Arc<CodexThread>`, it is possible that other references to it exist elsewhere.
     /// Returns the thread if the thread was found and removed.
     pub async fn remove_thread(&self, thread_id: &ThreadId) -> Option<Arc<CodexThread>> {
-        self.state.threads.write().await.remove(thread_id)
+        self.state.remove_thread(thread_id).await
     }
 
     /// Tries to shut down all tracked threads concurrently within the provided timeout.
@@ -924,6 +924,9 @@ impl ThreadManager {
 
         let mut tracked_threads = self.state.threads.write().await;
         for thread_id in &report.completed {
+            if let Some(thread) = tracked_threads.get(thread_id) {
+                thread.codex.session.thread_activity.mark_unpublished();
+            }
             tracked_threads.remove(thread_id);
         }
 
@@ -1168,7 +1171,11 @@ impl ThreadManagerState {
 
     /// Remove a thread from the manager by ID, returning it when present.
     pub(crate) async fn remove_thread(&self, thread_id: &ThreadId) -> Option<Arc<CodexThread>> {
-        self.threads.write().await.remove(thread_id)
+        let mut threads = self.threads.write().await;
+        if let Some(thread) = threads.get(thread_id) {
+            thread.codex.session.thread_activity.mark_unpublished();
+        }
+        threads.remove(thread_id)
     }
 
     pub(crate) async fn effective_multi_agent_version_for_spawn(
@@ -1554,6 +1561,7 @@ impl ThreadManagerState {
                         thread,
                     });
                 }
+                thread.codex.session.thread_activity.mark_unpublished();
                 threads.remove(&resumed.conversation_id);
             }
         }