From b133fd4368e93dea7463b4342d8b36d9c7d3efcf Mon Sep 17 00:00:00 2001 From: Winston Howes Date: Mon, 6 Jul 2026 18:46:56 -0700 Subject: [PATCH] core: make thread activity lifecycle atomic --- codex-rs/core/src/session/mod.rs | 14 +- codex-rs/core/src/session/session.rs | 48 +++- codex-rs/core/src/thread_activity.rs | 317 ++++++++++++++++++++- codex-rs/core/src/thread_activity_tests.rs | 87 ++++++ codex-rs/core/src/thread_manager.rs | 12 +- 5 files changed, 467 insertions(+), 11 deletions(-) diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 70fea7bc58b..404075ccf20 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/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> { + ) -> Option> { 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; } diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 4e0a92d83f6..a7bb9b233d6 100644 --- a/codex-rs/core/src/session/session.rs +++ b/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> { - self.submission_lifecycle.try_reserve(/*close*/ false) + pub(crate) fn try_reserve_submission( + &self, + close: bool, + ) -> Option> { + 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> { + self.try_reserve_submission(/*close*/ false) + } + + 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()?; + 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. diff --git a/codex-rs/core/src/thread_activity.rs b/codex-rs/core/src/thread_activity.rs index 7c8c0e1e0a8..b2f0b4297b0 100644 --- a/codex-rs/core/src/thread_activity.rs +++ b/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, 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, + thread_id: ThreadId, + generation: u64, + clear_closing_on_drop: bool, + delivery_prepared: bool, + active: bool, +} + impl ThreadActivityGate { pub(crate) fn register( self: &Arc, @@ -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, + 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, + thread_id: ThreadId, + generation: u64, + close: bool, + ) -> 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.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, + 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.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,7 +313,53 @@ 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) { @@ -187,9 +412,27 @@ impl ThreadActivityHandle { self.gate.mark_initialized(self.thread_id, self.generation); } + pub(crate) fn try_reserve(&self, close: bool) -> Option { + self.gate + .try_reserve(self.thread_id, self.generation, close) + } + + pub(crate) fn try_reserve_idle_shutdown(&self) -> Option { + 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; diff --git a/codex-rs/core/src/thread_activity_tests.rs b/codex-rs/core/src/thread_activity_tests.rs index a952e590a4e..b5b0f8d518c 100644 --- a/codex-rs/core/src/thread_activity_tests.rs +++ b/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)); +} diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 968f1f3c295..26cf0ff9411 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -885,7 +885,7 @@ impl ThreadManager { /// as `Arc`, 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> { - 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> { - 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); } }