diff --git a/codex-rs/core-api/src/lib.rs b/codex-rs/core-api/src/lib.rs index 6e3f009d97..97f2e8afbc 100644 --- a/codex-rs/core-api/src/lib.rs +++ b/codex-rs/core-api/src/lib.rs @@ -29,6 +29,9 @@ 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::AgentTreeShutdownFailure; +pub use codex_core::AgentTreeShutdownFailureReason; +pub use codex_core::AgentTreeShutdownReport; 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 952632070b..a8d57e19ce 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/codex-rs/core/src/agent/control.rs @@ -466,7 +466,7 @@ impl LocalAgentControl { let Ok(membership) = self.runtime.admit_start() else { return; }; - let teardown = membership.into_teardown_guard(); + let teardown = membership.into_teardown_guard("completion_watcher", Some(child_thread_id)); let control = self.clone(); let watcher = async move { let status = match control.subscribe_status(child_thread_id).await { diff --git a/codex-rs/core/src/agent/control/residency.rs b/codex-rs/core/src/agent/control/residency.rs index 7f2ad89af8..dd53eb6f2c 100644 --- a/codex-rs/core/src/agent/control/residency.rs +++ b/codex-rs/core/src/agent/control/residency.rs @@ -8,6 +8,7 @@ use crate::thread_manager::ThreadManagerState; use codex_protocol::ThreadId; use codex_protocol::error::AgentErrorContext; use codex_protocol::error::CodexErr; +use codex_protocol::error::CodexErrKind; use codex_protocol::error::CodexErrorDetails; use codex_protocol::error::Result as CodexResult; use codex_protocol::protocol::MultiAgentVersion; @@ -184,11 +185,15 @@ 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 teardown = membership + .clone() + .into_teardown_guard("resident_eviction", Some(candidate_thread_id)); let eviction = tokio::spawn(async move { let _residency_guard = residency_guard; candidate_thread.ensure_rollout_materialized().await; if let Err(err) = candidate_thread.shutdown_and_wait().await { + teardown + .record_shutdown_failure("stop_resident", CodexErrKind::from(&err).into()); warn!( "failed to shut down v2 resident thread before unloading {candidate_thread_id}: {err}" ); diff --git a/codex-rs/core/src/agent/control/runtime.rs b/codex-rs/core/src/agent/control/runtime.rs index 2ebc34a5f2..0ea885ab83 100644 --- a/codex-rs/core/src/agent/control/runtime.rs +++ b/codex-rs/core/src/agent/control/runtime.rs @@ -8,6 +8,9 @@ use crate::agent::api::AgentControl; use crate::agent::registry::AgentRegistry; use crate::config::RolloutBudgetConfig; use crate::rollout_budget::RolloutBudget; +use crate::thread_manager::AgentTreeShutdownFailure; +use crate::thread_manager::AgentTreeShutdownFailureReason; +use crate::thread_manager::AgentTreeShutdownReport; use crate::thread_manager::ThreadIdGenerator; use crate::thread_manager::ThreadManagerState; use arc_swap::ArcSwapOption; @@ -18,33 +21,86 @@ use codex_protocol::error::AgentErrorContext; use codex_protocol::error::CodexErr; use codex_protocol::error::Result as CodexResult; use std::sync::Arc; +use std::sync::Mutex; use std::sync::OnceLock; +use std::sync::PoisonError; 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; +use uuid::Uuid; -#[derive(Debug, Default)] +// Avoid retaining an unbounded number of failures during a long-lived tree's lifetime. +const MAX_RETAINED_SHUTDOWN_FAILURES: usize = 64; + +#[derive(Debug)] pub(crate) struct AgentTreeShutdownState { members: TaskTracker, - failed: AtomicBool, + report: Mutex, +} + +impl Default for AgentTreeShutdownState { + fn default() -> Self { + Self { + members: TaskTracker::new(), + report: Mutex::new(AgentTreeShutdownReport { + tree_id: Uuid::now_v7(), + failures: Vec::new(), + omitted_failures: 0, + }), + } + } } 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(()) + self.wait_detailed() + .await + .map_err(|_| CodexErr::Fatal("agent tree shutdown did not complete cleanly".to_owned())) } - fn record_failure(&self) { - self.failed.store(true, Ordering::Release); + pub(crate) async fn wait_detailed(&self) -> Result<(), AgentTreeShutdownReport> { + self.members.wait().await; + // Reporters hold a membership until after recording failures. Once the tracker drains, + // no admitted operation can add a failure after the snapshot. + let report = self + .report + .lock() + .unwrap_or_else(PoisonError::into_inner) + .clone(); + if report.failures.is_empty() && report.omitted_failures == 0 { + Ok(()) + } else { + Err(report) + } + } + + fn record_failure(&self, failure: AgentTreeShutdownFailure) { + let tree_id = { + let mut report = self.report.lock().unwrap_or_else(PoisonError::into_inner); + if report.failures.len() < MAX_RETAINED_SHUTDOWN_FAILURES { + report.failures.push(failure.clone()); + } else { + report.omitted_failures = report.omitted_failures.saturating_add(/*rhs*/ 1); + } + report.tree_id + }; + let (reason, phase, error_kind) = match &failure.reason { + AgentTreeShutdownFailureReason::OperationFailed { phase, error_kind } => { + ("operation_failed", Some(*phase), Some(*error_kind)) + } + AgentTreeShutdownFailureReason::GuardAbandoned => ("guard_abandoned", None, None), + }; + let thread_id = failure.thread_id.map(tracing::field::display); + tracing::warn!( + %tree_id, + operation = failure.operation, + phase, + thread_id, + reason, + error_kind, + "agent tree shutdown failure recorded" + ); } } @@ -55,9 +111,15 @@ pub(crate) struct AgentTreeMembership { } impl AgentTreeMembership { - pub(crate) fn into_teardown_guard(self) -> AgentTreeTeardownGuard { + pub(crate) fn into_teardown_guard( + self, + operation: &'static str, + thread_id: Option, + ) -> AgentTreeTeardownGuard { AgentTreeTeardownGuard { membership: self, + operation, + thread_id, completed: false, } } @@ -66,16 +128,35 @@ impl AgentTreeMembership { /// Marks tree shutdown as failed if teardown work exits without completing. pub(crate) struct AgentTreeTeardownGuard { membership: AgentTreeMembership, + operation: &'static str, + thread_id: Option, completed: bool, } impl AgentTreeTeardownGuard { - pub(crate) fn clone_for_teardown(&self) -> Self { - self.membership.clone().into_teardown_guard() + pub(crate) fn clone_for_teardown( + &self, + operation: &'static str, + thread_id: Option, + ) -> Self { + self.membership + .clone() + .into_teardown_guard(operation, thread_id) } - pub(crate) fn record_shutdown_failure(&self) { - self.membership.state.record_failure(); + pub(crate) fn set_thread_id(&mut self, thread_id: ThreadId) { + self.thread_id = Some(thread_id); + } + + pub(crate) fn record_shutdown_failure(&self, phase: &'static str, error_kind: &'static str) { + self.membership + .state + .record_failure(AgentTreeShutdownFailure::operation_failed( + self.operation, + phase, + self.thread_id, + error_kind, + )); } pub(crate) fn complete(mut self) { @@ -86,7 +167,14 @@ impl AgentTreeTeardownGuard { impl Drop for AgentTreeTeardownGuard { fn drop(&mut self) { if !self.completed { - self.record_shutdown_failure(); + // Record before membership is dropped, so a waiter cannot snapshot first. + self.membership + .state + .record_failure(AgentTreeShutdownFailure { + operation: self.operation, + thread_id: self.thread_id, + reason: AgentTreeShutdownFailureReason::GuardAbandoned, + }); } } } @@ -213,8 +301,10 @@ impl LocalAgentRuntime { Arc::clone(&self.shutdown_state) } - pub(crate) fn record_shutdown_failure(&self) { - self.shutdown_state.record_failure(); + /// The caller must hold tree membership until this record has been written. This keeps + /// the completed shutdown report stable for every waiter. + pub(crate) fn record_shutdown_failure(&self, failure: AgentTreeShutdownFailure) { + self.shutdown_state.record_failure(failure); } pub(crate) fn generate_thread_id(&self) -> ThreadId { diff --git a/codex-rs/core/src/agent/control/runtime_tests.rs b/codex-rs/core/src/agent/control/runtime_tests.rs index ea37e18f22..2884575ab3 100644 --- a/codex-rs/core/src/agent/control/runtime_tests.rs +++ b/codex-rs/core/src/agent/control/runtime_tests.rs @@ -1,29 +1,250 @@ use super::*; use crate::thread_manager::default_thread_id_generator; +use codex_protocol::error::CodexErrorDetails; +use codex_thread_store::ThreadStoreError; use futures::future; +use pretty_assertions::assert_eq; +use tokio::sync::Barrier; +use tokio::sync::oneshot; -#[tokio::test] -async fn aborted_teardown_fails_tree_shutdown() { - let runtime = LocalAgentRuntime::new( +fn runtime() -> LocalAgentRuntime { + LocalAgentRuntime::new( Weak::default(), default_thread_id_generator(), /*rollout_budget*/ None, - ); + ) +} + +#[tokio::test] +async fn aborted_teardown_reports_operation_and_thread() { + let runtime = runtime(); + let thread_id = ThreadId::new(); let teardown = runtime .admit_start() .expect("teardown should be admitted") - .into_teardown_guard(); + .into_teardown_guard("session_startup", Some(thread_id)); + let (ready_tx, ready_rx) = oneshot::channel(); let task = tokio::spawn(async move { let _teardown = teardown; + let _ = ready_tx.send(()); future::pending::<()>().await; }); + ready_rx.await.expect("teardown task should start"); let shutdown = runtime.request_shutdown(); task.abort(); task.await.expect_err("teardown task should be aborted"); - shutdown - .wait() + let report = shutdown + .wait_detailed() .await .expect_err("aborted teardown must fail tree shutdown"); + assert_eq!( + report, + AgentTreeShutdownReport { + tree_id: report.tree_id, + failures: vec![AgentTreeShutdownFailure { + operation: "session_startup", + thread_id: Some(thread_id), + reason: AgentTreeShutdownFailureReason::GuardAbandoned, + }], + omitted_failures: 0, + } + ); + assert_eq!( + report.to_string(), + format!( + "agent tree shutdown did not complete cleanly: tree_id={}, failures=[{{operation=session_startup, thread_id={thread_id}, reason=guard_abandoned}}]; omitted_failures=0", + report.tree_id, + ) + ); + let error = shutdown + .wait() + .await + .expect_err("legacy wait must fail after an abandoned teardown"); + assert!(matches!( + error.details(), + CodexErrorDetails::Fatal(message) if message == "agent tree shutdown did not complete cleanly" + )); +} + +#[tokio::test] +async fn completed_teardown_does_not_fail_tree_shutdown() { + let runtime = runtime(); + let teardown = runtime + .admit_start() + .expect("teardown should be admitted") + .into_teardown_guard("session_loop", /*thread_id*/ None); + let shutdown = runtime.request_shutdown(); + teardown.complete(); + + assert_eq!(shutdown.wait_detailed().await, Ok(())); + assert!(shutdown.wait().await.is_ok()); +} + +#[tokio::test] +async fn report_display_includes_safe_cause_without_raw_error_payload() { + let runtime = runtime(); + let thread_id = ThreadId::new(); + let teardown = runtime + .admit_start() + .expect("teardown should be admitted") + .into_teardown_guard("session_startup", Some(thread_id)); + let raw_payload = "sensitive thread-store error payload"; + let error = ThreadStoreError::Internal { + message: raw_payload.to_owned(), + }; + teardown.record_shutdown_failure( + "discard_persistence", + crate::thread_manager::thread_store_error_kind(&error), + ); + let shutdown = runtime.request_shutdown(); + teardown.complete(); + + let report = shutdown + .wait_detailed() + .await + .expect_err("recorded failure must fail shutdown"); + let displayed = report.to_string(); + assert_eq!( + displayed, + format!( + "agent tree shutdown did not complete cleanly: tree_id={}, failures=[{{operation=session_startup, phase=discard_persistence, thread_id={thread_id}, reason=operation_failed, error_kind=thread_store_internal}}]; omitted_failures=0", + report.tree_id, + ) + ); + assert!(!displayed.contains(raw_payload)); +} + +#[tokio::test] +async fn cloned_guard_reports_its_own_context_before_releasing_membership() { + let runtime = runtime(); + let original = runtime + .admit_start() + .expect("teardown should be admitted") + .into_teardown_guard("startup", /*thread_id*/ None); + let mut session = original.clone_for_teardown("session_loop", /*thread_id*/ None); + let thread_id = ThreadId::new(); + session.set_thread_id(thread_id); + let shutdown = runtime.request_shutdown(); + original.complete(); + + let wait = shutdown.wait_detailed(); + tokio::pin!(wait); + assert!(futures::poll!(&mut wait).is_pending()); + drop(session); + let report = wait + .await + .expect_err("abandoned session must fail shutdown"); + assert_eq!( + report, + AgentTreeShutdownReport { + tree_id: report.tree_id, + failures: vec![AgentTreeShutdownFailure { + operation: "session_loop", + thread_id: Some(thread_id), + reason: AgentTreeShutdownFailureReason::GuardAbandoned, + }], + omitted_failures: 0, + } + ); +} + +#[tokio::test] +async fn concurrent_failures_are_retained_for_every_waiter() { + let runtime = runtime(); + let barrier = Arc::new(Barrier::new(/*n*/ 3)); + let thread_ids = [ + ThreadId::from_u128(/*value*/ 1), + ThreadId::from_u128(/*value*/ 2), + ]; + let mut tasks = Vec::new(); + for thread_id in thread_ids { + let teardown = runtime + .admit_start() + .expect("teardown should be admitted") + .into_teardown_guard("child_spawn", Some(thread_id)); + let barrier = Arc::clone(&barrier); + tasks.push(tokio::spawn(async move { + barrier.wait().await; + teardown.record_shutdown_failure("discard_persistence", "thread_store_internal"); + teardown.complete(); + })); + } + let shutdown = runtime.request_shutdown(); + let first_waiter = { + let shutdown = Arc::clone(&shutdown); + tokio::spawn(async move { shutdown.wait_detailed().await }) + }; + let second_waiter = { + let shutdown = Arc::clone(&shutdown); + tokio::spawn(async move { shutdown.wait_detailed().await }) + }; + barrier.wait().await; + for task in tasks { + task.await.expect("cleanup task should finish"); + } + let mut first = first_waiter + .await + .expect("first waiter should finish") + .expect_err("recorded failures must fail shutdown"); + let second = second_waiter + .await + .expect("second waiter should finish") + .expect_err("recorded failures must fail shutdown"); + assert_eq!(first, second); + assert_eq!(shutdown.wait_detailed().await, Err(first.clone())); + + first + .failures + .sort_by_key(|failure| failure.thread_id.map(|thread_id| thread_id.to_string())); + assert_eq!( + first, + AgentTreeShutdownReport { + tree_id: first.tree_id, + failures: thread_ids + .into_iter() + .map(|thread_id| AgentTreeShutdownFailure::operation_failed( + "child_spawn", + "discard_persistence", + Some(thread_id), + "thread_store_internal", + )) + .collect(), + omitted_failures: 0, + } + ); +} + +#[tokio::test] +async fn failure_report_counts_omitted_failures() { + let runtime = runtime(); + let teardown = runtime + .admit_start() + .expect("teardown should be admitted") + .into_teardown_guard("startup", /*thread_id*/ None); + let failure = AgentTreeShutdownFailure::operation_failed( + "startup", + "stop_session", + /*thread_id*/ None, + "internal", + ); + for _ in 0..MAX_RETAINED_SHUTDOWN_FAILURES + 3 { + runtime.record_shutdown_failure(failure.clone()); + } + let shutdown = runtime.request_shutdown(); + teardown.complete(); + + let report = shutdown + .wait_detailed() + .await + .expect_err("recorded failures must fail shutdown"); + assert_eq!( + report, + AgentTreeShutdownReport { + tree_id: report.tree_id, + failures: vec![failure; MAX_RETAINED_SHUTDOWN_FAILURES], + omitted_failures: 3, + } + ); } diff --git a/codex-rs/core/src/agent/control/spawn_guard.rs b/codex-rs/core/src/agent/control/spawn_guard.rs index 509fee1ce0..d8befedd40 100644 --- a/codex-rs/core/src/agent/control/spawn_guard.rs +++ b/codex-rs/core/src/agent/control/spawn_guard.rs @@ -2,8 +2,11 @@ use super::AgentTreeMembership; use crate::thread_manager::ThreadManagerState; +use crate::thread_manager::thread_store_error_kind; +use codex_agent_graph_store::AgentGraphStoreError; use codex_agent_graph_store::ThreadSpawnEdgeStatus; use codex_protocol::ThreadId; +use codex_protocol::error::CodexErrKind; use codex_thread_store::ThreadStoreError; use std::sync::Arc; use tokio::task::JoinHandle; @@ -56,19 +59,23 @@ impl Drop for PendingSpawn { return; }; let state = Arc::clone(&self.state); - let teardown = membership.into_teardown_guard(); + let teardown = membership.into_teardown_guard("child_spawn_cleanup", Some(child)); 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(); + teardown + .record_shutdown_failure("stop_child", CodexErrKind::from(&error).into()); warn!("failed to stop 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(); + teardown.record_shutdown_failure( + "discard_child_persistence", + thread_store_error_kind(&error), + ); warn!("failed to discard cancelled child spawn: {error}"); } } @@ -83,7 +90,13 @@ impl Drop for PendingSpawn { .set_thread_spawn_edge_status(child, ThreadSpawnEdgeStatus::Closed) .await { - teardown.record_shutdown_failure(); + let error_kind = match &error { + AgentGraphStoreError::InvalidRequest { .. } => { + "agent_graph_store_invalid_request" + } + AgentGraphStoreError::Internal { .. } => "agent_graph_store_internal", + }; + teardown.record_shutdown_failure("close_spawn_edge", error_kind); 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 01a9eadfea..801d4d9f4b 100644 --- a/codex-rs/core/src/agent/control_tests.rs +++ b/codex-rs/core/src/agent/control_tests.rs @@ -4271,7 +4271,14 @@ async fn completion_watcher_does_not_hide_tree_shutdown_failure() { child_thread_id.to_string(), /*child_agent_path*/ None, ); - harness.control.runtime.record_shutdown_failure(); + harness.control.runtime.record_shutdown_failure( + crate::thread_manager::AgentTreeShutdownFailure::operation_failed( + "completion_watcher_test", + "test", + Some(parent_thread_id), + "test_error", + ), + ); let shutdown = harness.control.runtime.request_shutdown(); timeout(Duration::from_secs(5), shutdown.wait()) diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index e7753db4e5..a2b44e3f1f 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -148,6 +148,9 @@ mod thread_manager; pub(crate) mod web_search; pub(crate) mod windows_sandbox_read_grants; pub use thread_manager::AgentTreeShutdown; +pub use thread_manager::AgentTreeShutdownFailure; +pub use thread_manager::AgentTreeShutdownFailureReason; +pub use thread_manager::AgentTreeShutdownReport; 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 14f721d172..c569b04ec9 100644 --- a/codex-rs/core/src/session/handlers.rs +++ b/codex-rs/core/src/session/handlers.rs @@ -24,6 +24,8 @@ use crate::tasks::CompactTask; use crate::tasks::UserShellCommandMode; use crate::tasks::UserShellCommandTask; use crate::tasks::execute_user_shell_command; +use crate::thread_manager::AgentTreeShutdownFailure; +use crate::thread_manager::thread_store_error_kind; use codex_protocol::error::CodexErr; use codex_protocol::protocol::CodexErrorInfo; use codex_protocol::protocol::ErrorEvent; @@ -310,7 +312,14 @@ 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(); + sess.services.local_agent_runtime.record_shutdown_failure( + AgentTreeShutdownFailure::operation_failed( + "session_shutdown", + "code_mode", + Some(sess.thread_id), + "code_mode_error", + ), + ); warn!("failed to shutdown code mode session: {err}"); } sess.stop_mcp_prewarm_worker().await; @@ -356,7 +365,14 @@ 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(); + sess.services.local_agent_runtime.record_shutdown_failure( + AgentTreeShutdownFailure::operation_failed( + "session_shutdown", + "close_persistence", + Some(sess.thread_id), + thread_store_error_kind(&e), + ), + ); warn!("failed to shutdown thread persistence: {e}"); let event = Event { id: sub_id.clone(), @@ -704,7 +720,14 @@ 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(); + sess.services.local_agent_runtime.record_shutdown_failure( + AgentTreeShutdownFailure::operation_failed( + "session_shutdown", + "close_persistence_after_channel_closed", + Some(sess.thread_id), + thread_store_error_kind(&err), + ), + ); 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 e7832e1e9f..63d6929068 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -966,7 +966,7 @@ impl Session { // This task will run until Op::Shutdown is received. let tree_teardown = startup .as_ref() - .and_then(|startup| startup.session_teardown()); + .and_then(|startup| startup.session_teardown(thread_id)); 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) diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 0335827ff7..263f08f2f1 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -1846,7 +1846,7 @@ impl Session { next_internal_sub_id: AtomicU64::new(0), }); if let Some(startup) = &startup { - let _ = startup.session.set(Arc::clone(&sess)); + startup.set_session(Arc::clone(&sess)); } if let Some(network_policy_decider_session) = network_policy_decider_session { let mut guard = network_policy_decider_session.write().await; diff --git a/codex-rs/core/src/session/startup.rs b/codex-rs/core/src/session/startup.rs index 5338036fae..75cf4d5f51 100644 --- a/codex-rs/core/src/session/startup.rs +++ b/codex-rs/core/src/session/startup.rs @@ -10,6 +10,9 @@ use std::sync::PoisonError; use crate::agent::control::AgentTreeMembership; use crate::agent::control::AgentTreeTeardownGuard; +use crate::thread_manager::thread_store_error_kind; +use codex_protocol::ThreadId; +use codex_protocol::error::CodexErrKind; use codex_protocol::protocol::Op; use codex_thread_store::LiveThreadInitGuard; use futures::future::BoxFuture; @@ -32,16 +35,25 @@ impl SessionStartup { } pub(crate) fn hold_membership(&self, membership: AgentTreeMembership) { - let previous = self - .lock_teardown() - .replace(membership.into_teardown_guard()); + let teardown = membership.into_teardown_guard("session_startup", /*thread_id*/ None); + let previous = self.lock_teardown().replace(teardown); 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 set_session(&self, session: Arc) { + let thread_id = session.thread_id; + let mut teardown = self.lock_teardown(); + if self.session.set(session).is_ok() + && let Some(teardown) = teardown.as_mut() + { + teardown.set_thread_id(thread_id); + } + } + + pub(crate) fn session_teardown(&self, thread_id: ThreadId) -> Option { + let teardown = self.lock_teardown(); + let teardown = teardown.as_ref()?; + Some(teardown.clone_for_teardown("session_loop", Some(thread_id))) } pub(crate) fn release_membership(&self) { @@ -51,7 +63,12 @@ impl SessionStartup { } pub(crate) async fn cleanup(&self) { - let teardown = self.lock_teardown().take(); + let mut teardown = self.lock_teardown().take(); + if let Some(teardown) = teardown.as_mut() + && let Some(session) = self.session.get() + { + teardown.set_thread_id(session.thread_id); + } 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. @@ -59,7 +76,10 @@ impl SessionStartup { let _ = io.submit(Op::Interrupt).await; if let Err(error) = io.shutdown_and_wait().await { if let Some(teardown) = teardown.as_ref() { - teardown.record_shutdown_failure(); + teardown.record_shutdown_failure( + "stop_session_io", + CodexErrKind::from(&error).into(), + ); } tracing::warn!("failed to stop cancelled session init: {error}"); } @@ -70,7 +90,10 @@ impl SessionStartup { let mut persistence = std::mem::take(&mut *self.persistence.lock().await); if let Err(error) = persistence.discard().await { if let Some(teardown) = teardown.as_ref() { - teardown.record_shutdown_failure(); + teardown.record_shutdown_failure( + "discard_persistence", + thread_store_error_kind(&error), + ); } tracing::warn!( "failed to discard thread persistence for failed session init: {error}" diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 4bd5bff358..ea4dee313e 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -3,6 +3,10 @@ mod shared_instructions; mod shutdown; pub use shutdown::AgentTreeShutdown; +pub use shutdown::AgentTreeShutdownFailure; +pub use shutdown::AgentTreeShutdownFailureReason; +pub use shutdown::AgentTreeShutdownReport; +pub(crate) use shutdown::thread_store_error_kind; use crate::CodexAppsToolsCache; use crate::agent::LocalAgentControl; diff --git a/codex-rs/core/src/thread_manager/shutdown.rs b/codex-rs/core/src/thread_manager/shutdown.rs index f009ab5139..11f89bf39a 100644 --- a/codex-rs/core/src/thread_manager/shutdown.rs +++ b/codex-rs/core/src/thread_manager/shutdown.rs @@ -4,7 +4,103 @@ use super::ThreadManager; use crate::agent::control::AgentTreeShutdownState; use codex_protocol::ThreadId; use codex_protocol::error::Result as CodexResult; +use codex_thread_store::ThreadStoreError; +use std::fmt; use std::sync::Arc; +use uuid::Uuid; + +/// Return a stable, payload-free category for a thread-store error. +pub(crate) fn thread_store_error_kind(error: &ThreadStoreError) -> &'static str { + match error { + ThreadStoreError::ThreadNotFound { .. } => "thread_store_not_found", + ThreadStoreError::InvalidRequest { .. } => "thread_store_invalid_request", + ThreadStoreError::Conflict { .. } => "thread_store_conflict", + ThreadStoreError::Unsupported { .. } => "thread_store_unsupported", + ThreadStoreError::Internal { .. } => "thread_store_internal", + } +} + +/// A safely loggable reason that a tracked part of an agent tree failed to shut down. +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum AgentTreeShutdownFailureReason { + /// An operation returned an error at a specific step. The kind is a stable, + /// non-sensitive classification. + OperationFailed { + phase: &'static str, + error_kind: &'static str, + }, + /// A tracked operation was dropped before it could report that its work completed. + GuardAbandoned, +} + +/// One failure observed during the lifetime or teardown of an agent tree. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct AgentTreeShutdownFailure { + pub operation: &'static str, + pub thread_id: Option, + pub reason: AgentTreeShutdownFailureReason, +} + +impl AgentTreeShutdownFailure { + /// Classifies an operation failure without retaining arbitrary error messages. + pub fn operation_failed( + operation: &'static str, + phase: &'static str, + thread_id: Option, + error_kind: &'static str, + ) -> Self { + Self { + operation, + thread_id, + reason: AgentTreeShutdownFailureReason::OperationFailed { phase, error_kind }, + } + } +} + +/// The bounded set of failures observed for one local agent-tree instance. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct AgentTreeShutdownReport { + /// Correlation ID for this runtime instance; this is not a thread's ID. + pub tree_id: Uuid, + /// Failures retained in recording order. Each failure is also logged when recorded. + pub failures: Vec, + /// Number of additional failures omitted after the retention limit was reached. + pub omitted_failures: usize, +} + +impl fmt::Display for AgentTreeShutdownReport { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!( + f, + "agent tree shutdown did not complete cleanly: tree_id={}, failures=[", + self.tree_id + )?; + for (index, failure) in self.failures.iter().enumerate() { + if index > 0 { + write!(f, ", ")?; + } + write!(f, "{{operation={}", failure.operation)?; + if let AgentTreeShutdownFailureReason::OperationFailed { phase, .. } = &failure.reason { + write!(f, ", phase={phase}")?; + } + if let Some(thread_id) = failure.thread_id { + write!(f, ", thread_id={thread_id}")?; + } + match &failure.reason { + AgentTreeShutdownFailureReason::OperationFailed { error_kind, .. } => { + write!(f, ", reason=operation_failed, error_kind={error_kind}")?; + } + AgentTreeShutdownFailureReason::GuardAbandoned => { + write!(f, ", reason=guard_abandoned")?; + } + } + write!(f, "}}")?; + } + write!(f, "]; omitted_failures={}", self.omitted_failures) + } +} + +impl std::error::Error for AgentTreeShutdownReport {} /// Completion handle for one exact local agent tree. #[derive(Clone, Debug)] @@ -21,6 +117,12 @@ impl AgentTreeShutdown { pub async fn wait(&self) -> CodexResult<()> { self.state.wait().await } + + /// Waits for the tree and returns its recorded failure report if cleanup was unsuccessful. + /// Each call observes the same completed report. Dropping the future does not cancel shutdown. + pub async fn wait_detailed(&self) -> Result<(), AgentTreeShutdownReport> { + self.state.wait_detailed().await + } } impl ThreadManager { diff --git a/codex-rs/core/tests/suite/agent_execution.rs b/codex-rs/core/tests/suite/agent_execution.rs index 898fa44891..41b1b6e5a0 100644 --- a/codex-rs/core/tests/suite/agent_execution.rs +++ b/codex-rs/core/tests/suite/agent_execution.rs @@ -1,5 +1,7 @@ use anyhow::Context; use anyhow::Result; +use codex_core::AgentTreeShutdownFailure; +use codex_core::AgentTreeShutdownFailureReason; use codex_core::windows_sandbox::WindowsSandboxLevelExt; use codex_exec_server::CreateDirectoryOptions; use codex_features::Feature; @@ -701,7 +703,11 @@ async fn tree_shutdown_waits_for_children_without_affecting_other_roots() -> Res error.details(), CodexErrorDetails::InvalidRequest(_) )); - tokio::time::timeout(Duration::from_secs(/*secs*/ 10), shutdown.wait()).await??; + tokio::time::timeout(Duration::from_secs(/*secs*/ 10), shutdown.wait_detailed()) + .await? + .expect("tree shutdown should succeed"); + // Legacy waiters observe the same successful completion after a detailed wait. + 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!( @@ -809,9 +815,33 @@ async fn tree_shutdown_reports_persistence_writer_failure() -> Result<()> { .thread_manager .request_agent_tree_shutdown(root_id) .await?; - let result = tokio::time::timeout(Duration::from_secs(/*secs*/ 10), shutdown.wait()).await?; + let report = tokio::time::timeout(Duration::from_secs(/*secs*/ 10), shutdown.wait_detailed()) + .await? + .expect_err("shutdown should report the failed persistence writer"); - assert!(result.is_err()); + assert_eq!( + report.failures, + vec![AgentTreeShutdownFailure { + operation: "session_shutdown", + thread_id: Some(root_id), + reason: AgentTreeShutdownFailureReason::OperationFailed { + phase: "close_persistence", + error_kind: "thread_store_not_found", + }, + }] + ); + assert_eq!(report.omitted_failures, 0); + assert_eq!( + shutdown + .wait_detailed() + .await + .expect_err("failure persists"), + report + ); + shutdown + .wait() + .await + .expect_err("legacy wait must still report the failure"); Ok(()) }