From 8e44ad94b0519fbca39ed452bb2e1a6b43ad76a4 Mon Sep 17 00:00:00 2001 From: Owen Lin Date: Thu, 1 Oct 2026 00:37:02 +0000 Subject: [PATCH] Add coordinated shutdown for local agent trees (#49814) ## What changed Expose `ThreadManager::request_agent_tree_shutdown` and an `AgentTreeShutdown::wait` handle. Requests reject new starts in the shared runtime, cancel pending startup, and signal every session in that tree, including delegates outside the thread registry. Waiting covers admitted operations and session teardown, and reports cleanup or persistence writer failures. Other root trees remain unaffected; registry removal and persistence handoff remain caller responsibilities. Retain tree membership through startup, cancelled spawn cleanup, residency eviction, and completion watchers. Release persistence resources before returning startup failures, and stop completion watchers during tree shutdown without reporting shutdown as a child result. ## Testing Add coverage for tree isolation, children removed from the registry, delegate teardown, cancellation during required MCP initialization, aborted teardown, persistence writer failures, and successful retry after a failed resume releases its writer. GitOrigin-RevId: a109e25adf23d20baf7595891099947050c6b6b5 --- codex-rs/core-api/src/lib.rs | 1 + codex-rs/core/src/agent/control.rs | 25 ++- codex-rs/core/src/agent/control/residency.rs | 12 +- .../core/src/agent/control/residency_tests.rs | 30 ++- codex-rs/core/src/agent/control/runtime.rs | 111 ++++++++++ .../core/src/agent/control/runtime_tests.rs | 29 +++ codex-rs/core/src/agent/control/spawn.rs | 17 +- .../core/src/agent/control/spawn_guard.rs | 37 +++- codex-rs/core/src/agent/control_tests.rs | 27 +++ codex-rs/core/src/codex_delegate.rs | 30 ++- codex-rs/core/src/codex_delegate_tests.rs | 33 +++ codex-rs/core/src/lib.rs | 1 + codex-rs/core/src/session/handlers.rs | 18 +- codex-rs/core/src/session/mod.rs | 18 +- codex-rs/core/src/session/session.rs | 6 +- codex-rs/core/src/session/startup.rs | 91 +++++++- codex-rs/core/src/thread_manager.rs | 53 ++++- codex-rs/core/src/thread_manager/shutdown.rs | 46 ++++ codex-rs/core/tests/suite/agent_execution.rs | 196 ++++++++++++++++++ .../core/tests/suite/startup_cancellation.rs | 53 +++++ codex-rs/thread-store/src/live_thread.rs | 10 +- .../src/local/tests/acquisition_tests.rs | 2 +- 22 files changed, 795 insertions(+), 51 deletions(-) create mode 100644 codex-rs/core/src/agent/control/runtime_tests.rs create mode 100644 codex-rs/core/src/thread_manager/shutdown.rs diff --git a/codex-rs/core-api/src/lib.rs b/codex-rs/core-api/src/lib.rs index 0c7549b0c0..8ac16704a7 100644 --- 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 index 565d36519a..8ab0fca3fb 100644 --- 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 index f1041c9c4e..edaa97215b 100644 --- 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, config: &Config, + membership: &AgentTreeMembership, protected_thread_id: Option, ) -> CodexResult { 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, manager: &Arc, capacity: usize, + membership: &AgentTreeMembership, protected_thread_id: Option, ) -> CodexResult { 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, manager: &Arc, + membership: &AgentTreeMembership, protected_thread_id: Option, ) -> bool { // Keep shutting-down workers counted until removal. Each runtime's write guard @@ -178,6 +182,7 @@ 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; @@ -185,6 +190,7 @@ impl V2Residency { 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 index cfa375c813..9c850dc601 100644 --- 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 index 287883785f..5a621e7307 100644 --- 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, + _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>>, pub(super) registry: Arc, pub(super) residency: Arc, + /// Shared by every session in this tree, including private delegates. + pub(crate) shutdown: CancellationToken, + shutdown_state: Arc, } 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 { + 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 { + 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 new file mode 100644 index 0000000000..ea37e18f22 --- /dev/null +++ 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"); +} diff --git a/codex-rs/core/src/agent/control/spawn.rs b/codex-rs/core/src/agent/control/spawn.rs index 7ca849717b..e832dbddf7 100644 --- a/codex-rs/core/src/agent/control/spawn.rs +++ b/codex-rs/core/src/agent/control/spawn.rs @@ -318,6 +318,7 @@ impl LocalAgentControl { thread_id: ThreadId, parent: Option>, ) -> CodexResult<()> { + let membership = self.runtime.admit_start()?; let state = self.runtime.upgrade()?; let owner_thread_id = parent.as_ref().map(|parent| parent.session.thread_id); if let Some(parent) = &parent { @@ -589,7 +590,7 @@ impl LocalAgentControl { // Reserving a slot can evict an idle nested parent. Capture its instructions // alongside its authority so the child does not depend on a later live lookup. let residency_slot = self - .reserve_v2_residency_slot(&state, &config, Some(thread_id)) + .reserve_v2_residency_slot(&state, &config, &membership, Some(thread_id)) .await?; match state @@ -638,6 +639,7 @@ impl LocalAgentControl { session_source: Option, options: SpawnAgentOptions, ) -> CodexResult<(LiveAgent, ThreadConfigSnapshot)> { + let membership = self.runtime.admit_start()?; let spawn_started_at = Instant::now(); let state = self.runtime.upgrade()?; let multi_agent_version = state @@ -661,7 +663,12 @@ impl LocalAgentControl { let (residency_slot, residency_reservation) = if spawn_uses_v2_residency { let residency_reservation_started_at = Instant::now(); let residency_slot = self - .reserve_v2_residency_slot(&state, &config, /*protected_thread_id*/ None) + .reserve_v2_residency_slot( + &state, + &config, + &membership, + /*protected_thread_id*/ None, + ) .await?; ( Some(residency_slot), @@ -780,7 +787,8 @@ impl LocalAgentControl { } }; agent_metadata.agent_id = Some(new_thread.thread_id); - let mut pending_spawn = PendingSpawn::new(Arc::clone(&state), new_thread.thread_id); + let mut pending_spawn = + PendingSpawn::new(Arc::clone(&state), new_thread.thread_id, membership); if let Some(SessionSource::SubAgent( subagent_source @ SubAgentSource::ThreadSpawn { @@ -871,7 +879,7 @@ impl LocalAgentControl { if let Some(residency_slot) = residency_slot { residency_slot.commit(new_thread.thread_id); } - pending_spawn.disarm(); + let _membership = pending_spawn.disarm(); // Notify a new thread has been created. This notification will be processed by clients // to subscribe or drain this newly created thread. @@ -1261,6 +1269,7 @@ impl LocalAgentControl { thread_id: ThreadId, session_source: SessionSource, ) -> CodexResult { + let _membership = self.runtime.admit_start()?; let root_depth = thread_spawn_depth(&session_source).unwrap_or(0); let (resumed_thread_id, resumed_multi_agent_version) = Box::pin( self.resume_single_agent_from_rollout(config.clone(), thread_id, session_source), diff --git a/codex-rs/core/src/agent/control/spawn_guard.rs b/codex-rs/core/src/agent/control/spawn_guard.rs index 7df61e64ac..509fee1ce0 100644 --- a/codex-rs/core/src/agent/control/spawn_guard.rs +++ b/codex-rs/core/src/agent/control/spawn_guard.rs @@ -1,23 +1,29 @@ //! Owns a spawned child until its initial input is accepted. +use super::AgentTreeMembership; use crate::thread_manager::ThreadManagerState; use codex_agent_graph_store::ThreadSpawnEdgeStatus; use codex_protocol::ThreadId; +use codex_thread_store::ThreadStoreError; use std::sync::Arc; use tokio::task::JoinHandle; use tracing::warn; pub(super) struct PendingSpawn { state: Arc, - child: Option, + pending: Option<(ThreadId, AgentTreeMembership)>, edge_write: Option>, } impl PendingSpawn { - pub(super) fn new(state: Arc, child: ThreadId) -> Self { + pub(super) fn new( + state: Arc, + child: ThreadId, + membership: AgentTreeMembership, + ) -> Self { Self { state, - child: Some(child), + pending: Some((child, membership)), edge_write: None, } } @@ -36,27 +42,36 @@ impl PendingSpawn { self.edge_write = None; } - pub(super) fn disarm(mut self) { - self.child = None; + pub(super) fn disarm(mut self) -> AgentTreeMembership { + let Some((_child, membership)) = self.pending.take() else { + unreachable!("pending spawn must own agent-tree membership"); + }; + membership } } impl Drop for PendingSpawn { fn drop(&mut self) { - let Some(child) = self.child.take() else { + let Some((child, membership)) = self.pending.take() else { return; }; let state = Arc::clone(&self.state); + let teardown = membership.into_teardown_guard(); let edge_write = self.edge_write.take(); drop(tokio::spawn(async move { if let Some(thread) = state.remove_thread(&child).await { if let Err(error) = thread.shutdown_and_wait().await { + teardown.record_shutdown_failure(); warn!("failed to stop cancelled child spawn: {error}"); } - if let Some(live_thread) = thread.session.live_thread() - && let Err(error) = live_thread.discard().await - { - warn!("failed to discard cancelled child spawn: {error}"); + if let Some(live_thread) = thread.session.live_thread() { + match live_thread.discard().await { + Ok(()) | Err(ThreadStoreError::ThreadNotFound { .. }) => {} + Err(error) => { + teardown.record_shutdown_failure(); + warn!("failed to discard cancelled child spawn: {error}"); + } + } } } // A pending Open write must finish before cleanup writes Closed. @@ -68,8 +83,10 @@ impl Drop for PendingSpawn { .set_thread_spawn_edge_status(child, ThreadSpawnEdgeStatus::Closed) .await { + teardown.record_shutdown_failure(); warn!("failed to close cancelled child spawn edge: {error}"); } + teardown.complete(); })); } } diff --git a/codex-rs/core/src/agent/control_tests.rs b/codex-rs/core/src/agent/control_tests.rs index 5d483f3853..669236dc02 100644 --- a/codex-rs/core/src/agent/control_tests.rs +++ b/codex-rs/core/src/agent/control_tests.rs @@ -4344,6 +4344,33 @@ async fn completion_watcher_notifies_parent_when_child_is_missing() { ); } +#[tokio::test] +async fn completion_watcher_does_not_hide_tree_shutdown_failure() { + let harness = AgentControlHarness::new().await; + let (parent_thread_id, _parent_thread) = harness.start_thread().await; + let (child_thread_id, _child_thread) = harness.start_thread().await; + + harness.control.maybe_start_completion_watcher( + child_thread_id, + Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id, + depth: 1, + agent_path: None, + agent_nickname: None, + agent_role: Some("explorer".to_string()), + })), + child_thread_id.to_string(), + /*child_agent_path*/ None, + ); + harness.control.runtime.record_shutdown_failure(); + let shutdown = harness.control.runtime.request_shutdown(); + + timeout(Duration::from_secs(5), shutdown.wait()) + .await + .expect("completion watcher should stop during tree shutdown") + .expect_err("recorded tree shutdown failure should be returned"); +} + #[tokio::test] async fn spawn_thread_subagent_gets_random_nickname_in_session_source() { let harness = AgentControlHarness::new().await; diff --git a/codex-rs/core/src/codex_delegate.rs b/codex-rs/core/src/codex_delegate.rs index 9d6b9656ba..04e5d14e00 100644 --- a/codex-rs/core/src/codex_delegate.rs +++ b/codex-rs/core/src/codex_delegate.rs @@ -27,6 +27,8 @@ use crate::session::SessionIo; use crate::session::SessionSpawnArgs; use crate::session::emit_subagent_session_started; use crate::session::session::Session; +use crate::session::startup::SessionStartup; +use crate::session::startup::SessionStartupGuard; use crate::session::turn_context::TurnContext; use codex_history::InitialHistory; use codex_login::AuthManager; @@ -73,6 +75,10 @@ pub(crate) async fn run_codex_thread_interactive( let conversation_history = initial_history.unwrap_or(InitialHistory::New); let forked_from_thread_id = conversation_history.forked_from_id(); + let runtime = parent_session.services.local_agent_runtime.clone(); + let startup = Arc::new(SessionStartup::default()); + startup.hold_membership(runtime.admit_start()?); + let startup_guard = SessionStartupGuard::new(Arc::clone(&startup)); let instructions = parent_session.inherited_instructions().await; let session_source = SessionSource::SubAgent(subagent_source.clone()); let is_guardian_reviewer = crate::guardian::is_basic_session_source(&session_source); @@ -83,8 +89,8 @@ pub(crate) async fn run_codex_thread_interactive( }; let mut thread_extension_init = codex_extension_api::ExtensionDataInit::default(); thread_extension_init.insert(isolation); - let (session, io) = Session::spawn(SessionSpawnArgs { - startup: None, + let spawn_result = Session::spawn(SessionSpawnArgs { + startup: Some(Arc::clone(&startup)), config, allow_provider_model_fallback: false, instructions, @@ -116,7 +122,7 @@ pub(crate) async fn run_codex_thread_interactive( originator: parent_ctx.originator.clone(), agent_control: crate::agent::control::AgentControlInit::Provided { control: Arc::clone(&parent_session.services.agent_control), - runtime: parent_session.services.local_agent_runtime.clone(), + runtime, }, dynamic_tools: Vec::new(), metrics_service_name: None, @@ -139,7 +145,18 @@ pub(crate) async fn run_codex_thread_interactive( windows_sandbox_proxy_settings_mode, }) .or_cancel(&cancel_token) - .await??; + .await; + let (session, io) = match spawn_result { + Ok(Ok(spawned)) => spawned, + Ok(Err(error)) => { + startup_guard.cleanup().await; + return Err(error); + } + Err(error) => { + startup_guard.cleanup().await; + return Err(error.into()); + } + }; let thread_config = session.thread_config_snapshot().await; let client_metadata = parent_session.app_server_client_metadata().await; emit_subagent_session_started( @@ -151,7 +168,10 @@ pub(crate) async fn run_codex_thread_interactive( thread_config, subagent_source, ); - Ok((session, forward_session_io(Arc::new(io), cancel_token))) + let caller_io = forward_session_io(Arc::new(io), cancel_token); + startup.release_membership(); + startup_guard.disarm(); + Ok((session, caller_io)) } /// Keeps delegate IO cancellation identical for standalone and manager-owned reviewers. diff --git a/codex-rs/core/src/codex_delegate_tests.rs b/codex-rs/core/src/codex_delegate_tests.rs index f8762c98ce..3923eb4fd5 100644 --- a/codex-rs/core/src/codex_delegate_tests.rs +++ b/codex-rs/core/src/codex_delegate_tests.rs @@ -14,6 +14,7 @@ use codex_protocol::protocol::McpStartupUpdateEvent; use codex_protocol::protocol::RawResponseItemEvent; use codex_protocol::protocol::TurnAbortReason; use codex_protocol::protocol::TurnAbortedEvent; +use futures::FutureExt; use pretty_assertions::assert_eq; use std::sync::Arc; use std::sync::atomic::AtomicUsize; @@ -404,3 +405,35 @@ async fn run_codex_thread_interactive_rejects_approval_policy_that_can_prompt() ) )); } + +/// Private delegates participate in tree shutdown even though they are not manager-visible. +#[tokio::test] +async fn tree_shutdown_waits_for_private_delegate() { + let (parent, context, _events) = + crate::session::tests::make_session_and_context_with_rx().await; + let mut config = context.config.as_ref().clone(); + config.permissions.approval_policy = Constrained::allow_only(AskForApproval::Never); + let (_, io) = run_codex_thread_interactive( + config, + Arc::clone(&parent.services.auth_manager), + Arc::clone(&parent.services.models_manager), + Arc::clone(&parent), + Arc::clone(&context), + context.initial_environments.clone(), + CancellationToken::new(), + SubAgentSource::Review, + codex_extension_api::SessionIsolation::Isolated, + /*initial_history*/ None, + crate::session::GitEnrichmentPolicy::Fresh, + codex_sandboxing::WindowsSandboxProxySettingsMode::Reconcile, + ) + .await + .expect("start private delegate"); + + let shutdown = parent.services.local_agent_runtime.request_shutdown(); + timeout(Duration::from_secs(/*secs*/ 10), shutdown.wait()) + .await + .expect("tree shutdown should wait for the private delegate") + .expect("private delegate should shut down cleanly"); + assert!(io.session_loop_termination.now_or_never().is_some()); +} diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index 45ff58ab6d..25edb95267 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -145,6 +145,7 @@ pub use codex_prompts as review_prompts; mod thread_manager; pub(crate) mod web_search; pub(crate) mod windows_sandbox_read_grants; +pub use thread_manager::AgentTreeShutdown; pub use thread_manager::ForkSnapshot; pub use thread_manager::InternalSessionParent; pub use thread_manager::NewThread; diff --git a/codex-rs/core/src/session/handlers.rs b/codex-rs/core/src/session/handlers.rs index 5114a80b24..a30c195427 100644 --- a/codex-rs/core/src/session/handlers.rs +++ b/codex-rs/core/src/session/handlers.rs @@ -309,6 +309,7 @@ pub(super) async fn shutdown_session_runtime(sess: &Arc) { .terminate_all_processes() .await; if let Err(err) = sess.services.code_mode_service.shutdown().await { + sess.services.local_agent_runtime.record_shutdown_failure(); warn!("failed to shutdown code mode session: {err}"); } sess.stop_mcp_prewarm_worker().await; @@ -354,6 +355,7 @@ pub async fn shutdown(sess: &Arc, sub_id: String) -> bool { if let Some(live_thread) = sess.live_thread() && let Err(e) = live_thread.shutdown().await { + sess.services.local_agent_runtime.record_shutdown_failure(); warn!("failed to shutdown thread persistence: {e}"); let event = Event { id: sub_id.clone(), @@ -422,9 +424,20 @@ pub(super) async fn submission_loop( config: Arc, rx_sub: Receiver, ) { - // To break out of this loop, send Op::Shutdown. + // Session shutdown and tree shutdown both use the existing teardown handler. let mut shutdown_received = false; - while let Ok(sub) = rx_sub.recv().await { + loop { + let sub = tokio::select! { + biased; + _ = sess.services.local_agent_runtime.shutdown.cancelled() => { + shutdown_received = shutdown(&sess, super::new_submission_id()).await; + break; + } + sub = rx_sub.recv() => match sub { + Ok(sub) => sub, + Err(_) => break, + }, + }; if matches!(sub.op, Op::ResolveElicitation { .. }) { debug!(submission_id = %sub.id, operation = sub.op.kind(), "Submission"); } else { @@ -653,6 +666,7 @@ pub(super) async fn submission_loop( if let Some(live_thread) = sess.live_thread() && let Err(err) = live_thread.shutdown().await { + sess.services.local_agent_runtime.record_shutdown_failure(); warn!("failed to shutdown thread persistence after submission channel closed: {err}"); } } diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 4708d2921e..1d1944095a 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -512,6 +512,7 @@ impl Session { pub(crate) fn spawn( args: SessionSpawnArgs, ) -> BoxFuture<'static, CodexResult<(Arc, SessionIo)>> { + let tree_shutdown = args.agent_control.runtime().shutdown.clone(); Box::pin(async move { let parent_trace = match args.parent_trace { Some(trace) => { @@ -528,12 +529,16 @@ impl Session { if let Some(trace) = parent_trace.as_ref() { let _ = set_parent_from_w3c_trace_context(&thread_spawn_span, trace); } - Self::spawn_internal(SessionSpawnArgs { + let spawn = Self::spawn_internal(SessionSpawnArgs { parent_trace, ..args }) - .instrument(thread_spawn_span) - .await + .instrument(thread_spawn_span); + tokio::select! { + biased; + _ = tree_shutdown.cancelled() => Err(CodexErr::TurnAborted), + result = spawn => result, + } }) } @@ -931,11 +936,17 @@ impl Session { let thread_id = session.thread_id; // This task will run until Op::Shutdown is received. + let tree_teardown = startup + .as_ref() + .and_then(|startup| startup.session_teardown()); let session_for_loop = Arc::clone(&session); let session_loop_handle = tokio::spawn(storage_originator.scope(async move { submission_loop(session_for_loop, configured_config, rx_sub) .instrument(info_span!("session_loop", thread_id = %thread_id)) .await; + if let Some(tree_teardown) = tree_teardown { + tree_teardown.complete(); + } })); let io = SessionIo { tx_sub, @@ -946,6 +957,7 @@ impl Session { if let Some(startup) = startup { let _ = startup.io.set(io.clone()); + startup.persistence.lock().await.commit(); } Ok((session, io)) } diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index c012ae02c5..878bffa8f2 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -1966,7 +1966,11 @@ impl Session { Ok(sess) } Err(err) => { - live_thread_init.discard().await; + if let Err(error) = live_thread_init.discard().await { + tracing::warn!( + "failed to discard thread persistence for failed session init: {error}" + ); + } Err(err) } } diff --git a/codex-rs/core/src/session/startup.rs b/codex-rs/core/src/session/startup.rs index cc1a18064c..5338036fae 100644 --- a/codex-rs/core/src/session/startup.rs +++ b/codex-rs/core/src/session/startup.rs @@ -1,11 +1,18 @@ -//! Retains persistence acquisition and resources throughout managed startup. -//! The thread manager drops initialization, then joins acquisition and resource cleanup here. +//! Retains agent-tree membership and resources throughout managed startup. +//! Cleanup releases membership only after acquisition and partial session teardown finish; +//! abandoning that ownership marks tree shutdown as failed. use std::sync::Arc; +use std::sync::Mutex as StdMutex; +use std::sync::MutexGuard; use std::sync::OnceLock; +use std::sync::PoisonError; +use crate::agent::control::AgentTreeMembership; +use crate::agent::control::AgentTreeTeardownGuard; use codex_protocol::protocol::Op; use codex_thread_store::LiveThreadInitGuard; +use futures::future::BoxFuture; use tokio::sync::Mutex; use super::SessionIo; @@ -13,25 +20,101 @@ use super::session::Session; #[derive(Default)] pub(crate) struct SessionStartup { + teardown: StdMutex>, pub(crate) persistence: Mutex, pub(crate) session: OnceLock>, pub(crate) io: OnceLock, } impl SessionStartup { + fn lock_teardown(&self) -> MutexGuard<'_, Option> { + self.teardown.lock().unwrap_or_else(PoisonError::into_inner) + } + + pub(crate) fn hold_membership(&self, membership: AgentTreeMembership) { + let previous = self + .lock_teardown() + .replace(membership.into_teardown_guard()); + assert!(previous.is_none(), "agent-tree membership already set"); + } + + pub(crate) fn session_teardown(&self) -> Option { + self.lock_teardown() + .as_ref() + .map(AgentTreeTeardownGuard::clone_for_teardown) + } + + pub(crate) fn release_membership(&self) { + if let Some(teardown) = self.lock_teardown().take() { + teardown.complete(); + } + } + pub(crate) async fn cleanup(&self) { + let teardown = self.lock_teardown().take(); if let Some(io) = self.io.get() { // The session loop owns persistence now. Preserve its shutdown semantics even // if registration or the caller's handoff was interrupted after the loop started. self.persistence.lock().await.commit(); let _ = io.submit(Op::Interrupt).await; - let _ = io.shutdown_and_wait().await; + if let Err(error) = io.shutdown_and_wait().await { + if let Some(teardown) = teardown.as_ref() { + teardown.record_shutdown_failure(); + } + tracing::warn!("failed to stop cancelled session init: {error}"); + } } else { if let Some(session) = self.session.get() { super::handlers::shutdown_session_runtime(session).await; } let mut persistence = std::mem::take(&mut *self.persistence.lock().await); - persistence.discard().await; + if let Err(error) = persistence.discard().await { + if let Some(teardown) = teardown.as_ref() { + teardown.record_shutdown_failure(); + } + tracing::warn!( + "failed to discard thread persistence for failed session init: {error}" + ); + } + } + if let Some(teardown) = teardown { + teardown.complete(); } } } + +/// Cleans up a cancelled session start before releasing its agent-tree membership. +pub(crate) struct SessionStartupGuard { + startup: Option>, +} + +impl SessionStartupGuard { + pub(crate) fn new(startup: Arc) -> Self { + Self { + startup: Some(startup), + } + } + + pub(crate) fn disarm(mut self) { + self.startup = None; + } + + pub(crate) fn cleanup(mut self) -> BoxFuture<'static, ()> { + Box::pin(async move { + if let Some(startup) = self.startup.take() { + startup.cleanup().await; + } + }) + } +} + +impl Drop for SessionStartupGuard { + fn drop(&mut self) { + let Some(startup) = self.startup.take() else { + return; + }; + drop(tokio::spawn(async move { + startup.cleanup().await; + })); + } +} diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 187e1d793e..9b6ec56ff8 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -1,5 +1,8 @@ mod managed; mod shared_instructions; +mod shutdown; + +pub use shutdown::AgentTreeShutdown; use crate::CodexAppsToolsCache; use crate::agent::LocalAgentControl; @@ -23,6 +26,8 @@ use crate::session::SessionIo; use crate::session::SessionSpawnArgs; use crate::session::resolve_multi_agent_version; use crate::session::session::Session; +use crate::session::startup::SessionStartup; +use crate::session::startup::SessionStartupGuard; use crate::tasks::InterruptedTurnHistoryMarker; use crate::tasks::interrupted_turn_history_marker; use crate::thread_startup_metadata::ThreadStartupMetadata; @@ -2038,7 +2043,18 @@ impl ThreadManagerState { } /// Spawn a new thread with optional history and register it with the manager. - async fn spawn_thread(&self, request: ThreadSpawnRequest) -> CodexResult { + async fn spawn_thread(&self, mut request: ThreadSpawnRequest) -> CodexResult { + let runtime = request.agent_control.runtime().clone(); + let membership = runtime.admit_start()?; + let owns_startup = request.startup.is_none(); + let startup_state = Arc::clone( + request + .startup + .get_or_insert_with(|| Arc::new(SessionStartup::default())), + ); + startup_state.hold_membership(membership); + let startup_guard = + owns_startup.then(|| SessionStartupGuard::new(Arc::clone(&startup_state))); let ThreadSpawnRequest { startup, options, @@ -2133,10 +2149,13 @@ impl ThreadManagerState { resumed.conversation_id ))); } - drop(threads); let session_configured = thread .startup_metadata() .to_session_configured_event(initial_history.get_event_msgs()); + startup_state.release_membership(); + if let Some(startup_guard) = startup_guard { + startup_guard.disarm(); + } return Ok(NewThread { thread_id: resumed.conversation_id, session_configured, @@ -2268,7 +2287,7 @@ impl ThreadManagerState { }; let attachment_source = forked_from_thread_id.filter(|_| matches!(&initial_history, InitialHistory::Forked(_))); - let (session, io) = Session::spawn(SessionSpawnArgs { + let spawn_result = Session::spawn(SessionSpawnArgs { startup, config, allow_provider_model_fallback, @@ -2329,7 +2348,16 @@ impl ThreadManagerState { }, windows_sandbox_proxy_settings_mode, }) - .await?; + .await; + let (session, io) = match spawn_result { + Ok(spawned) => spawned, + Err(error) => { + if let Some(startup_guard) = startup_guard { + startup_guard.cleanup().await; + } + return Err(error); + } + }; if let Some(source_thread_id) = attachment_source && session.live_thread().is_some() && self.thread_store.supports_thread_attachments() @@ -2356,9 +2384,18 @@ impl ThreadManagerState { { session.services.mcp_runtime.enable_full_access_form_input(); } - let new_thread = self + let finalize_result = self .finalize_thread_spawn(session, io, tracked_session_source, initial_host_config) - .await?; + .await; + let new_thread = match finalize_result { + Ok(new_thread) => new_thread, + Err(error) => { + if let Some(startup_guard) = startup_guard { + startup_guard.cleanup().await; + } + return Err(error); + } + }; new_thread.thread.emit_thread_ready_lifecycle().await; if source_changed_during_startup.load(Ordering::Acquire) { new_thread.thread.session.request_mcp_runtime_refresh(); @@ -2366,6 +2403,10 @@ impl ThreadManagerState { if is_resumed_thread { new_thread.thread.emit_thread_resume_lifecycle().await; } + startup_state.release_membership(); + if let Some(startup_guard) = startup_guard { + startup_guard.disarm(); + } Ok(new_thread) } diff --git a/codex-rs/core/src/thread_manager/shutdown.rs b/codex-rs/core/src/thread_manager/shutdown.rs new file mode 100644 index 0000000000..f009ab5139 --- /dev/null +++ b/codex-rs/core/src/thread_manager/shutdown.rs @@ -0,0 +1,46 @@ +//! Fences and signals one local agent tree, then exposes its teardown result. + +use super::ThreadManager; +use crate::agent::control::AgentTreeShutdownState; +use codex_protocol::ThreadId; +use codex_protocol::error::Result as CodexResult; +use std::sync::Arc; + +/// Completion handle for one exact local agent tree. +#[derive(Clone, Debug)] +#[must_use = "agent-tree shutdown is not complete until wait() finishes"] +pub struct AgentTreeShutdown { + state: Arc, +} + +impl AgentTreeShutdown { + /// Waits for every operation and session admitted before the shutdown fence to finish. + /// Returns an error if cleanup or persistence writer termination failed. Registry removal and + /// caller-owned persistence handoff remain with the caller. Dropping this future does not + /// cancel shutdown. + pub async fn wait(&self) -> CodexResult<()> { + self.state.wait().await + } +} + +impl ThreadManager { + /// Requests shutdown of a loaded thread and every session sharing its local runtime. + /// + /// Returns after fencing new starts and signalling admitted sessions. Callers can wait on the + /// returned handle before completing registry cleanup or a persistence handoff. This fences + /// only starts sharing the current runtime; callers must separately serialize top-level loads + /// or resumes of the same thread ID through that handoff. + pub async fn request_agent_tree_shutdown( + &self, + thread_id: ThreadId, + ) -> CodexResult { + let thread = self.get_thread(thread_id).await?; + Ok(AgentTreeShutdown { + state: thread + .session + .services + .local_agent_runtime + .request_shutdown(), + }) + } +} diff --git a/codex-rs/core/tests/suite/agent_execution.rs b/codex-rs/core/tests/suite/agent_execution.rs index 480b5bb532..3d1c99dfcb 100644 --- a/codex-rs/core/tests/suite/agent_execution.rs +++ b/codex-rs/core/tests/suite/agent_execution.rs @@ -1,3 +1,4 @@ +use anyhow::Context; use anyhow::Result; use codex_core::windows_sandbox::WindowsSandboxLevelExt; use codex_exec_server::CreateDirectoryOptions; @@ -10,6 +11,7 @@ use codex_protocol::protocol::EnvironmentConfigState; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::ThreadSettingsOverrides; use codex_protocol::protocol::TurnEnvironmentSelections; +use core_test_support::apps_test_server::AppsTestServer; use core_test_support::responses::ResponseMock; use core_test_support::responses::ev_assistant_message; use core_test_support::responses::ev_completed; @@ -18,9 +20,11 @@ use core_test_support::responses::ev_response_created; use core_test_support::responses::mount_sse_once_match; use core_test_support::responses::sse; use core_test_support::responses::start_mock_server; +use core_test_support::skip_if_no_network; use core_test_support::submit_thread_settings; use core_test_support::test_codex::test_codex; use core_test_support::wait_for_event; +use futures::FutureExt; use pretty_assertions::assert_eq; use serde_json::json; use std::time::Duration; @@ -618,3 +622,195 @@ async fn v2_residency_reload_preserves_inherited_environment_and_tools( #[path = "agent_eviction_tests.rs"] mod eviction_tests; + +/// Tree shutdown waits for all admitted members, even after registry removal, without crossing +/// into another root's tree. +#[tokio::test] +async fn tree_shutdown_waits_for_children_without_affecting_other_roots() -> Result<()> { + use codex_core::StartThreadOptions; + use codex_protocol::error::CodexErrorDetails; + use codex_protocol::protocol::InternalSessionSource; + use codex_protocol::protocol::SessionSource; + + let server = start_mock_server().await; + mount_root_collaboration_call( + &server, + FIRST_PROMPT, + "first-call", + "spawn_agent", + json!({"message": FIRST_TASK, "task_name": "first", "fork_turns": "none"}), + ) + .await; + mount_completed_worker(&server, FIRST_TASK, "first-call").await; + let test = test_codex() + .with_model("gpt-5.6-sol") + .with_config(|config| { + config.features.enable(Feature::Collab).unwrap(); + config.features.enable(Feature::MultiAgentV2).unwrap(); + }) + .build_with_auto_env(&server) + .await?; + test.submit_turn(FIRST_PROMPT).await?; + let root_id = test.session_configured.thread_id; + let worker_id = test + .thread_manager + .list_thread_ids() + .await + .into_iter() + .find(|id| *id != root_id) + .expect("spawned worker"); + let worker = test.thread_manager.get_thread(worker_id).await?; + wait_for_event(&worker, |event| matches!(event, EventMsg::TurnComplete(_))).await; + let internal = test + .thread_manager + .spawn_internal_session( + root_id, + StartThreadOptions { + session_source: Some(SessionSource::Internal(InternalSessionSource::Guardian)), + ..StartThreadOptions::new(test.config.clone()) + }, + ) + .await?; + let other = test + .thread_manager + .start_thread(StartThreadOptions::new(test.config.clone())) + .await?; + + assert!(test.codex.wait_until_terminated().now_or_never().is_none()); + let shutdown = test + .thread_manager + .request_agent_tree_shutdown(worker_id) + .await?; + // A live child can leave the registry before its owner finishes teardown. + test.thread_manager.remove_thread(&worker_id).await; + let error = match test + .thread_manager + .spawn_internal_session( + root_id, + StartThreadOptions { + session_source: Some(SessionSource::Internal(InternalSessionSource::Guardian)), + ..StartThreadOptions::new(test.config.clone()) + }, + ) + .await + { + Ok(_) => panic!("agent-tree shutdown should reject a new internal session"), + Err(error) => error, + }; + assert!(matches!( + error.details(), + CodexErrorDetails::InvalidRequest(_) + )); + tokio::time::timeout(Duration::from_secs(/*secs*/ 10), shutdown.wait()).await??; + assert!(test.codex.wait_until_terminated().now_or_never().is_some()); + assert!(worker.wait_until_terminated().now_or_never().is_some()); + assert!( + internal + .thread + .wait_until_terminated() + .now_or_never() + .is_some() + ); + assert!( + other + .thread + .wait_until_terminated() + .now_or_never() + .is_none() + ); + // A root that has never used collaboration can also receive the signal. + let other_shutdown = test + .thread_manager + .request_agent_tree_shutdown(other.thread_id) + .await?; + tokio::time::timeout(Duration::from_secs(/*secs*/ 10), other_shutdown.wait()).await??; + assert!( + other + .thread + .wait_until_terminated() + .now_or_never() + .is_some() + ); + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn tree_shutdown_cancels_admitted_session_startup() -> Result<()> { + use codex_core::StartThreadOptions; + use codex_protocol::error::CodexErrorDetails; + use codex_protocol::protocol::InternalSessionSource; + use codex_protocol::protocol::SessionSource; + + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let test = test_codex().build_with_auto_env(&server).await?; + let root_id = test.session_configured.thread_id; + let mcp_server = start_mock_server().await; + let (http_server, control) = AppsTestServer::mount_with_startup_control(&mcp_server).await?; + let release = control.hold_next_successful_initialize(); + let mut config = test.config.clone(); + let mut servers = config.mcp_servers.get().clone(); + servers.insert( + "stalled".to_owned(), + serde_json::from_value(json!({ + "url": format!("{}/api/codex/ps/mcp", http_server.chatgpt_base_url), + "http_headers": { "Authorization": "Bearer synthetic-test-token" }, + "required": true, + "startup_timeout_sec": 120, + }))?, + ); + config.mcp_servers.set(servers)?; + let mut start = Box::pin(test.thread_manager.spawn_internal_session( + root_id, + StartThreadOptions { + session_source: Some(SessionSource::Internal(InternalSessionSource::Guardian)), + ..StartThreadOptions::new(config) + }, + )); + + tokio::select! { + _ = &mut start => panic!("startup must wait for the required MCP server"), + result = tokio::time::timeout(Duration::from_secs(10), async { + while control.initialize_attempts() == 0 { + tokio::task::yield_now().await; + } + }) => result.context("required MCP initialization did not begin")?, + } + let shutdown = test + .thread_manager + .request_agent_tree_shutdown(root_id) + .await?; + let (start_result, shutdown_result) = tokio::time::timeout(Duration::from_secs(10), async { + tokio::join!(&mut start, shutdown.wait()) + }) + .await + .context("tree shutdown must cancel admitted startup")?; + shutdown_result?; + let error = match start_result { + Ok(_) => panic!("cancelled internal session startup must not succeed"), + Err(error) => error, + }; + assert!(matches!(error.details(), CodexErrorDetails::TurnAborted)); + drop(release); + test.codex.shutdown_and_wait().await?; + Ok(()) +} + +#[tokio::test] +async fn tree_shutdown_reports_persistence_writer_failure() -> Result<()> { + let server = start_mock_server().await; + let test = test_codex().build_with_auto_env(&server).await?; + let root_id = test.session_configured.thread_id; + + // Close the real local writer first so Core's shutdown observes a writer failure. + test.thread_store.shutdown_thread(root_id).await?; + let shutdown = test + .thread_manager + .request_agent_tree_shutdown(root_id) + .await?; + let result = tokio::time::timeout(Duration::from_secs(/*secs*/ 10), shutdown.wait()).await?; + + assert!(result.is_err()); + Ok(()) +} diff --git a/codex-rs/core/tests/suite/startup_cancellation.rs b/codex-rs/core/tests/suite/startup_cancellation.rs index 5901e1feb5..5a86101f07 100644 --- a/codex-rs/core/tests/suite/startup_cancellation.rs +++ b/codex-rs/core/tests/suite/startup_cancellation.rs @@ -17,10 +17,12 @@ use codex_extension_api::McpServerContributionContext; use codex_extension_api::McpServerContributor; use codex_rollout::RolloutRecorder; use codex_thread_store::LocalThreadStore; +use core_test_support::apps_test_server::AppsTestServer; use core_test_support::responses::start_mock_server; use core_test_support::skip_if_no_network; use core_test_support::test_codex::test_codex; use pretty_assertions::assert_eq; +use serde_json::json; use tokio::sync::Notify; use tokio::time::timeout; @@ -113,3 +115,54 @@ async fn cancelled_resume_releases_writer_while_mcp_startup_is_pending() -> Resu resumed.thread.shutdown_and_wait().await?; Ok(()) } + +#[tokio::test] +async fn failed_resume_releases_writer_before_returning() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let test = test_codex() + .with_config(|config| config.experimental_thread_store = ThreadStoreConfig::Local) + .build_with_auto_env(&server) + .await?; + let thread_id = test.session_configured.thread_id; + let environments = test.codex.environment_selections().await; + test.codex.ensure_rollout_materialized().await; + let rollout_path = test.codex.rollout_path().context("thread rollout")?; + test.codex.shutdown_and_wait().await?; + test.thread_manager.remove_thread(&thread_id).await; + let history = RolloutRecorder::get_rollout_history(&rollout_path).await?; + + let mcp_server = start_mock_server().await; + let (http_server, control) = AppsTestServer::mount_with_startup_control(&mcp_server).await?; + control.fail_next_initialize_attempts(/*attempts*/ 1); + let mut config = test.config.clone(); + let mut servers = config.mcp_servers.get().clone(); + servers.insert( + "fails-once".to_owned(), + serde_json::from_value(json!({ + "url": format!("{}/api/codex/ps/mcp", http_server.chatgpt_base_url), + "http_headers": { "Authorization": "Bearer synthetic-test-token" }, + "required": true, + }))?, + ); + config.mcp_servers.set(servers)?; + let resume_options = || StartThreadOptions { + initial_history: history.clone(), + environments: Some(environments.clone()), + ..StartThreadOptions::new(config.clone()) + }; + + assert!( + test.thread_manager + .start_thread(resume_options()) + .await + .is_err(), + "first resume should fail MCP initialization" + ); + let resumed = test.thread_manager.start_thread(resume_options()).await?; + + assert_eq!(resumed.thread_id, thread_id); + resumed.thread.shutdown_and_wait().await?; + Ok(()) +} diff --git a/codex-rs/thread-store/src/live_thread.rs b/codex-rs/thread-store/src/live_thread.rs index f9bb507fbc..ab18df6d8e 100644 --- a/codex-rs/thread-store/src/live_thread.rs +++ b/codex-rs/thread-store/src/live_thread.rs @@ -96,14 +96,12 @@ impl LiveThreadInitGuard { self.live_thread = None; } - pub async fn discard(&mut self) { - let _ = self.finish_acquisition().await; + pub async fn discard(&mut self) -> ThreadStoreResult<()> { + self.finish_acquisition().await?; let Some(live_thread) = self.live_thread.take() else { - return; + return Ok(()); }; - if let Err(err) = live_thread.discard().await { - warn!("failed to discard thread persistence for failed session init: {err}"); - } + live_thread.discard().await } } diff --git a/codex-rs/thread-store/src/local/tests/acquisition_tests.rs b/codex-rs/thread-store/src/local/tests/acquisition_tests.rs index db02660571..a67ddff77c 100644 --- a/codex-rs/thread-store/src/local/tests/acquisition_tests.rs +++ b/codex-rs/thread-store/src/local/tests/acquisition_tests.rs @@ -43,7 +43,7 @@ async fn cancelled_acquisition_waits_for_writer_handoff_before_discarding() { "cleanup must retain the unfinished acquisition" ); release.send(()).expect("creation still owned"); - cleanup.await; + cleanup.await.expect("discard acquired writer"); assert!(matches!( store.flush_thread(thread_id).await, Err(ThreadStoreError::ThreadNotFound { .. })