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 { .. })