diff --git a/codex-rs/core-api/src/lib.rs b/codex-rs/core-api/src/lib.rs index 222b8f2f49..71ecadcc53 100644 --- a/codex-rs/core-api/src/lib.rs +++ b/codex-rs/core-api/src/lib.rs @@ -41,6 +41,7 @@ pub use codex_core::EnvironmentNetworkPolicy; pub use codex_core::ForkSnapshot; pub use codex_core::InputWaitOutcome; pub use codex_core::LoadedAgentsMd; +pub use codex_core::LoadedSuspendedAgent; pub use codex_core::McpManager; pub use codex_core::NetworkDomainPermission; pub use codex_core::NetworkDomainPermissionEntry; @@ -59,6 +60,7 @@ pub use codex_core::StartupModels; pub use codex_core::StateDbHandle; pub use codex_core::SteerSubmission; pub use codex_core::SuspendTurnOutcome; +pub use codex_core::SuspendedAgentTree; pub use codex_core::ThreadManager; pub use codex_core::ThreadShutdownReport; pub use codex_core::ThreadStoreBuilder; diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index 6ef9c31d7f..aa227fee4d 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/codex-rs/core/src/agent/control.rs @@ -264,21 +264,37 @@ impl LocalAgentControl { context: AgentCommunicationContext, start_options: TurnStartOptions, ) -> CodexResult { - let observers = state.extensions.message_lifecycle_contributors(); - let delivery = (!observers.is_empty()) - .then(|| MessageDelivery::publish(context.sender_thread_id, observers)); + let terminal_result = context.kind == AgentCommunicationKind::Result; + let publish = || { + MessageDelivery::publish( + context.sender_thread_id, + state.extensions.message_lifecycle_contributors(), + terminal_result.then_some(&self.runtime.unconfirmed_results), + ) + }; let log_content = crate::agent_communication::send_log_content(&communication); - { - // Keep unloaded delivery atomic with publication of a reloaded session. + let (thread, guard) = { + // Choose the loaded submission queue or the eviction mailbox while reload is fenced. let threads = state.threads.read().await; - if !communication.trigger_turn && !threads.contains_key(&agent_id) { - let _membership = self.runtime.admit_start()?; - self.runtime.ensure_agent_known(agent_id)?; - let submission_id = uuid::Uuid::now_v7().to_string(); + let thread = threads.get(&agent_id); + let guard = thread + .filter(|_| terminal_result) + .and_then(|thread| Arc::clone(&thread.residency_gate).try_read_owned().ok()); + if !communication.trigger_turn + && (thread.is_none() || (terminal_result && guard.is_none())) + { + // A reporting child already holds membership until it publishes the result. + let _membership = (!terminal_result) + .then(|| self.runtime.admit_start()) + .transpose()?; + if thread.is_none() { + self.runtime.ensure_agent_known(agent_id)?; + } + let submission_id = Uuid::now_v7().to_string(); self.runtime.mailboxes.enqueue( agent_id, Some(submission_id.clone()), - vec![crate::PendingAgentMessage::new(communication, delivery)], + vec![crate::PendingAgentMessage::new(communication, publish())], )?; drop(threads); crate::agent_communication::emit_agent_communication_send( @@ -289,8 +305,19 @@ impl LocalAgentControl { ); return Ok(submission_id); } - } - // Loaded recipients retain submission ordering with follow-ups and interrupts. + ( + thread.cloned().ok_or(CodexErr::ThreadNotFound(agent_id))?, + guard, + ) + }; + let guard = match guard { + Some(guard) => { + self.runtime.touch_v2_recipient(&thread); + Some(guard) + } + None => self.runtime.pin_v2_residency(&thread).await?, + }; + let delivery = publish(); let result = submit_loaded_agent_communication( agent_id, communication, @@ -298,7 +325,7 @@ impl LocalAgentControl { log_content.as_deref(), start_options, |op, parent_turn_id, root_turn_id| { - state.send_op(agent_id, op, parent_turn_id, root_turn_id, delivery) + state.send_pinned_op(&thread, op, parent_turn_id, root_turn_id, guard, delivery) }, ) .await; @@ -519,8 +546,10 @@ impl LocalAgentControl { message, /*trigger_turn*/ false, ); - let context = - AgentCommunicationContext::new(AgentCommunicationKind::Result, child_thread_id); + let context = AgentCommunicationContext::new( + AgentCommunicationKind::CompletionWatcher, + child_thread_id, + ); let _ = control .send_inter_agent_communication( parent_thread_id, diff --git a/codex-rs/core/src/agent/control/api.rs b/codex-rs/core/src/agent/control/api.rs index 0426ca39d8..bdb03f1b5c 100644 --- a/codex-rs/core/src/agent/control/api.rs +++ b/codex-rs/core/src/agent/control/api.rs @@ -179,7 +179,7 @@ impl AgentControl for LocalAgentControl { } fn take_mailbox(&self, agent: ThreadId) -> Vec { - self.runtime.mailboxes.take(agent) + self.runtime.mailboxes.take(agent, |_| true) } fn watch_mailbox(&self, agent: ThreadId) -> tokio::sync::watch::Receiver { diff --git a/codex-rs/core/src/agent/control/completion.rs b/codex-rs/core/src/agent/control/completion.rs index 587f4cdaeb..ef64a6096d 100644 --- a/codex-rs/core/src/agent/control/completion.rs +++ b/codex-rs/core/src/agent/control/completion.rs @@ -1,7 +1,7 @@ //! Delivers terminal child results and completion activity to the agent tree. //! //! Sessions capture terminal state; the controller owns routing and queue-only delivery. -//! Delivery remains best effort; metrics distinguish queue acceptance from delivery failure. +//! Unconsumed results remain available to the source writer through graceful suspension. use super::LocalAgentControl; use crate::TurnStartOptions; @@ -11,6 +11,7 @@ use crate::agent_communication::AgentCommunicationKind; use crate::session_prefix::format_guardian_interruption_message; use crate::session_prefix::format_inter_agent_completion_message; use codex_protocol::AgentPath; +use codex_protocol::ResponseItemId; use codex_protocol::items::SubAgentActivityItem; use codex_protocol::protocol::AgentStatus; use codex_protocol::protocol::CodexErrorInfo; @@ -106,23 +107,20 @@ impl LocalAgentControl { let Some(message) = message else { return; }; - // `communication` owns the message. Keep a second copy only when the - // recorder will actually need it after parent delivery succeeds. let trace_message = trace.is_enabled().then(|| message.clone()); - let communication = InterAgentCommunication::new( + let mut communication = InterAgentCommunication::new( child_agent_path.clone(), parent_agent_path, Vec::new(), message, /*trigger_turn*/ false, ); - let context = - AgentCommunicationContext::new(AgentCommunicationKind::Result, outcome.thread_id); + communication.id = Some(ResponseItemId::with_suffix("amsg", uuid::Uuid::now_v7())); let delivery = self .send_inter_agent_communication( parent_thread_id, communication, - context, + AgentCommunicationContext::new(AgentCommunicationKind::Result, outcome.thread_id), TurnStartOptions::default(), ) .await; diff --git a/codex-rs/core/src/agent/control/mailbox.rs b/codex-rs/core/src/agent/control/mailbox.rs index c5e7e06ee8..4ca6ff93a9 100644 --- a/codex-rs/core/src/agent/control/mailbox.rs +++ b/codex-rs/core/src/agent/control/mailbox.rs @@ -45,18 +45,26 @@ impl Mailboxes { .ok_or_else(|| CodexErr::InvalidRequest("agent runtime is shutting down".into()))? .entry(agent) .or_default(); - mailbox - .pending - .extend(messages.into_iter().map(|communication| StoredMail { - id: id.clone(), - communication, - _diagnostics_guard: PENDING_MAILBOX_MESSAGES.track(), - })); + let mail = messages.into_iter().map(|communication| StoredMail { + id: id.clone(), + communication, + _diagnostics_guard: PENDING_MAILBOX_MESSAGES.track(), + }); + if id.is_some() { + mailbox.pending.extend(mail); + } else { + // Mail returned by eviction predates results published while its writer closed. + mailbox.pending.splice(..0, mail); + } mailbox.state.send_replace(!mailbox.pending.is_empty()); Ok(()) } - pub(super) fn take(&self, agent: ThreadId) -> Vec { + pub(super) fn take( + &self, + agent: ThreadId, + accept: impl Fn(&PendingAgentMessage) -> bool, + ) -> Vec { let mut mailboxes = self.0.lock().unwrap_or_else(PoisonError::into_inner); let Some(mailbox) = mailboxes .as_mut() @@ -66,8 +74,10 @@ impl Mailboxes { }; let pending = std::mem::take(&mut mailbox.pending); mailbox.state.send_if_modified(std::mem::take); + drop(mailboxes); pending .into_iter() + .filter(|mail| accept(&mail.communication)) .map(|mail| { // Mail transferred during eviction was already received by the old session. if let Some(id) = mail.id { diff --git a/codex-rs/core/src/agent/control/residency.rs b/codex-rs/core/src/agent/control/residency.rs index 8383722075..e0c9860ecc 100644 --- a/codex-rs/core/src/agent/control/residency.rs +++ b/codex-rs/core/src/agent/control/residency.rs @@ -66,8 +66,8 @@ impl ThreadManagerState { teardown.complete(); return Err(err); } - // Dispatch has stopped and the residency guard excludes senders. A tree - // shutdown deliberately discards its mailbox instead of retaining it. + // Dispatch has stopped. Put previously received mail ahead of results + // published directly to the mailbox while eviction owned the residency guard. let mail = thread .session .input_queue @@ -80,10 +80,11 @@ impl ThreadManagerState { let _ = runtime.mailboxes.enqueue(thread_id, /*id*/ None, mail); // Reuse the same effective configuration snapshot as a fork. This includes live // permissions and network restrictions that a profile name cannot reconstruct. - let (config, ..) = thread.session.fork_config().await; + let (config, snapshot, ..) = thread.session.fork_config().await; // Preserve configured bindings across eviction without retaining turn-only attachments. // Raw bindings keep the original pending owner configuration alive. let environments = EvictedEnvironments { + history_mode: snapshot.history_mode, config: Arc::new(config), has_independent_service_tier: matches!( &thread.session.agent_runtime, @@ -174,6 +175,12 @@ impl LocalAgentControl { impl LocalAgentRuntime { /// Pins and touches the registered runtime without waiting for unrelated eviction. + pub(super) fn touch_v2_recipient(&self, thread: &CodexThread) { + if is_resident_candidate(thread) { + self.residency.touch(thread.session.thread_id); + } + } + pub(crate) async fn pin_v2_residency( &self, thread: &CodexThread, @@ -187,7 +194,7 @@ impl LocalAgentRuntime { if !std::ptr::eq(thread, registered.as_ref()) { return Err(CodexErr::ThreadNotFound(thread_id)); } - self.residency.touch(thread_id); + self.touch_v2_recipient(thread); Ok(Some(guard)) } } diff --git a/codex-rs/core/src/agent/control/residency_tests.rs b/codex-rs/core/src/agent/control/residency_tests.rs index e098af6440..bb052e9d24 100644 --- a/codex-rs/core/src/agent/control/residency_tests.rs +++ b/codex-rs/core/src/agent/control/residency_tests.rs @@ -107,7 +107,7 @@ async fn residency_slot_reservation_unloads_oldest_idle_v2_agent() { .services .local_agent_runtime .mailboxes - .take(first.thread_id); + .take(first.thread_id, |_| true); assert_eq!(retained, vec![communication.into()]); let second = spawn_v2_subagent(&control, &state, config, root.thread_id, "worker-2").await; second_slot.commit(second.thread_id); diff --git a/codex-rs/core/src/agent/control/resume.rs b/codex-rs/core/src/agent/control/resume.rs index c311ff2e90..f69d8f07a0 100644 --- a/codex-rs/core/src/agent/control/resume.rs +++ b/codex-rs/core/src/agent/control/resume.rs @@ -1,16 +1,102 @@ //! Reopens an agent only when its runtime is absent, then returns its effective state. -//! Rollout and descendant restoration remain in the existing local startup path. +//! Restores recorded ancestor ownership before loading a suspended descendant. use super::LocalAgentControl; +use crate::CodexThread; use crate::agent::types::LiveAgent; use crate::codex_thread::ThreadConfigSnapshot; use crate::config::Config; +use crate::thread_manager::LoadedSuspendedAgent; use codex_protocol::ThreadId; +use codex_protocol::error::CodexErr; use codex_protocol::error::CodexErrorDetails; use codex_protocol::error::Result as CodexResult; +use codex_protocol::protocol::MultiAgentVersion; use codex_protocol::protocol::SessionSource; +use codex_protocol::protocol::SubAgentSource; +use codex_thread_store::ReadThreadParams; +use std::collections::HashSet; +use std::sync::Arc; impl LocalAgentControl { + /// Load ancestors first so each child inherits its actual parent's runtime and permissions. + pub(crate) async fn load_suspended_child( + &self, + root: &Arc, + child_id: ThreadId, + ) -> CodexResult { + let root_id = root.startup_metadata().thread_id; + if child_id == root_id || root.multi_agent_version() != Some(MultiAgentVersion::V2) { + return Err(CodexErr::InvalidRequest( + "expected a named V2 descendant".to_owned(), + )); + } + let state = self.runtime.upgrade()?; + let mut chain = Vec::new(); + let mut visited = HashSet::new(); + let mut id = child_id; + while id != root_id { + if !visited.insert(id) { + return Err(CodexErr::InvalidRequest( + "cycle in stored agent ancestry".into(), + )); + } + let stored = state + .read_stored_thread(ReadThreadParams { + thread_id: id, + include_archived: false, + include_history: false, + }) + .await?; + let SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id, + agent_path: Some(path), + .. + }) = &stored.source + else { + return Err(CodexErr::InvalidRequest(format!( + "thread {id} is not a named child of {root_id}" + ))); + }; + if stored + .parent_thread_id + .is_some_and(|parent| parent != *parent_thread_id) + || stored + .agent_path + .as_deref() + .is_some_and(|name| name != path.as_str()) + { + return Err(CodexErr::InvalidRequest(format!( + "inconsistent stored ancestry for {id}" + ))); + } + chain.push(id); + id = *parent_thread_id; + } + chain.reverse(); + let mut thread = Arc::clone(root); + let mut residency = None; + for id in chain { + let config = thread.session.get_config().await; + // Usually root resume registered the name from the graph. This also handles + // stores without a graph, using the same identity restoration path. + self.restore_agent_identities(&config, &state, vec![id]) + .await; + let pin = self + .ensure_v2_agent_loaded(config.as_ref().clone(), id, Some(thread)) + .await?; + thread = state.get_thread(id).await?; + // Allow an idle ancestor to give up its slot while its child reloads at capacity. + if id == child_id { + residency = pin; + } + } + Ok(LoadedSuspendedAgent { + thread, + _residency: residency, + }) + } + pub(crate) async fn resume_agent( &self, config: Config, diff --git a/codex-rs/core/src/agent/control/runtime.rs b/codex-rs/core/src/agent/control/runtime.rs index 1672d2aee7..7b92b5c055 100644 --- a/codex-rs/core/src/agent/control/runtime.rs +++ b/codex-rs/core/src/agent/control/runtime.rs @@ -11,6 +11,7 @@ use crate::rollout_budget::RolloutBudget; use crate::thread_manager::AgentTreeShutdownFailure; use crate::thread_manager::AgentTreeShutdownFailureReason; use crate::thread_manager::AgentTreeShutdownReport; +use crate::thread_manager::SuspendedAgentTree; use crate::thread_manager::ThreadIdGenerator; use crate::thread_manager::ThreadManagerState; use arc_swap::ArcSwapOption; @@ -24,11 +25,14 @@ use codex_protocol::error::Result as CodexResult; use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::SubAgentSource; use codex_rollout_trace::ThreadTraceContext; +use std::collections::HashMap; use std::sync::Arc; use std::sync::Mutex; use std::sync::OnceLock; use std::sync::PoisonError; use std::sync::Weak; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; use tokio_util::sync::CancellationToken; use tokio_util::task::TaskTracker; use tokio_util::task::task_tracker::TaskTrackerToken; @@ -41,12 +45,27 @@ const MAX_RETAINED_SHUTDOWN_FAILURES: usize = 64; pub(crate) struct AgentTreeShutdownState { members: TaskTracker, report: Mutex, + disposition: OnceLock, + suspended: Mutex, + children: Mutex>, +} + +#[derive(Debug)] +enum TreeDisposition { + Shutdown, + Suspend { + root: ThreadId, + evicted_writers: TaskTracker, + }, } impl Default for AgentTreeShutdownState { fn default() -> Self { Self { members: TaskTracker::new(), + disposition: OnceLock::new(), + suspended: Mutex::default(), + children: Mutex::default(), report: Mutex::new(AgentTreeShutdownReport { tree_id: Uuid::now_v7(), failures: Vec::new(), @@ -65,8 +84,14 @@ impl AgentTreeShutdownState { 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. + // The source may need to reopen an evicted parent's writer after all live + // sessions stop. Both shutdown waiters and suspension callers join that write. + if let Some(TreeDisposition::Suspend { + evicted_writers, .. + }) = self.disposition.get() + { + evicted_writers.wait().await; + } let report = self .report .lock() @@ -79,6 +104,15 @@ impl AgentTreeShutdownState { } } + pub(crate) async fn suspended_tree(&self) -> CodexResult { + self.wait().await?; + Ok(self + .suspended + .lock() + .unwrap_or_else(PoisonError::into_inner) + .clone()) + } + fn record_failure(&self, failure: AgentTreeShutdownFailure) { let tree_id = { let mut report = self.report.lock().unwrap_or_else(PoisonError::into_inner); @@ -112,6 +146,7 @@ impl AgentTreeShutdownState { pub(crate) struct AgentTreeMembership { state: Arc, _member: TaskTrackerToken, + _child: Option, } impl AgentTreeMembership { @@ -202,6 +237,7 @@ pub(crate) struct LocalAgentRuntime { Arc>>, pub(super) registry: Arc, pub(super) residency: Arc, + pub(super) unconfirmed_results: Arc, pub(super) mailboxes: Arc, /// Shared by every session in this tree, including private delegates. pub(crate) shutdown: CancellationToken, @@ -219,6 +255,7 @@ impl LocalAgentRuntime { thread_id_generator, registry: Arc::default(), residency: Arc::default(), + unconfirmed_results: Arc::default(), mailboxes: Arc::default(), shutdown: CancellationToken::new(), shutdown_state: Arc::default(), @@ -233,6 +270,35 @@ impl LocalAgentRuntime { runtime } + pub(crate) async fn wait_for_children(&self, parent: ThreadId) { + let children = self + .shutdown_state + .children + .lock() + .unwrap_or_else(PoisonError::into_inner) + .get(&parent) + .cloned(); + if let Some(children) = children { + children.close(); + children.wait().await; + } + } + + pub(crate) fn admit_session( + &self, + parent: Option, + ) -> CodexResult { + let mut children = self + .shutdown_state + .children + .lock() + .unwrap_or_else(PoisonError::into_inner); + let mut membership = self.admit_start()?; + children.retain(|_, tracker| !tracker.is_empty()); + membership._child = parent.map(|id| children.entry(id).or_default().token()); + Ok(membership) + } + pub(crate) fn tree_id(&self) -> Uuid { self.shutdown_state .report @@ -294,6 +360,7 @@ impl LocalAgentRuntime { let membership = AgentTreeMembership { state: Arc::clone(&self.shutdown_state), _member: self.shutdown_state.members.token(), + _child: None, }; // 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. @@ -307,12 +374,145 @@ impl LocalAgentRuntime { } pub(crate) fn request_shutdown(&self) -> Arc { + let disposition = self + .shutdown_state + .disposition + .get_or_init(|| TreeDisposition::Shutdown); self.shutdown_state.members.close(); self.shutdown.cancel(); - self.mailboxes.close(); + if matches!(disposition, TreeDisposition::Shutdown) { + self.mailboxes.close(); + } Arc::clone(&self.shutdown_state) } + pub(crate) fn request_suspension( + &self, + root: ThreadId, + ) -> CodexResult> { + let disposition = self.shutdown_state.disposition.get_or_init(|| { + let evicted_writers = TaskTracker::new(); + let guard = AgentTreeMembership { + state: Arc::clone(&self.shutdown_state), + _member: evicted_writers.token(), + _child: None, + } + .into_teardown_guard("save_evicted_parent_results", Some(root)); + evicted_writers.close(); + let runtime = self.clone(); + tokio::spawn(async move { + if let Err(failure) = runtime.save_results_after_eviction().await { + runtime.record_shutdown_failure(failure); + } + runtime.mailboxes.close(); + guard.complete(); + }); + TreeDisposition::Suspend { + root, + evicted_writers, + } + }); + if !matches!(disposition, TreeDisposition::Suspend { root: owner, .. } if *owner == root) { + return Err(CodexErr::InvalidRequest( + "agent tree shutdown already has a different disposition".into(), + )); + } + Ok(self.request_shutdown()) + } + + async fn save_results_after_eviction(&self) -> Result<(), AgentTreeShutdownFailure> { + self.shutdown_state.members.wait().await; + if self.unconfirmed_results.load(Ordering::Acquire) == 0 { + return Ok(()); + } + // A failed teardown may still own a writer. Never reopen it. + if !self + .shutdown_state + .report + .lock() + .unwrap_or_else(PoisonError::into_inner) + .failures + .is_empty() + { + return Ok(()); + } + for parent in self + .registry + .live_agents() + .into_iter() + .filter_map(|agent| agent.agent_id) + { + let results = self + .mailboxes + .take(parent, crate::PendingAgentMessage::is_terminal_result); + if results.is_empty() { + continue; + } + let failed = |phase, kind| { + AgentTreeShutdownFailure::operation_failed( + "save_evicted_parent_results", + phase, + Some(parent), + kind, + ) + }; + let owner = self + .registry + .evicted_environments(parent) + .ok_or_else(|| failed("locate_parent", "parent_not_evicted"))?; + self.upgrade() + .map_err(|_| failed("open_writer", "manager_unavailable"))? + .save_evicted_parent_results(parent, &owner.config, owner.history_mode, &results) + .await + .map_err(|err| { + failed( + "persist", + crate::thread_manager::thread_store_error_kind(&err), + ) + })?; + } + // A task that claimed a result and stopped before its append returned leaves + // an unconfirmed write. Only its writer knows whether the append took effect. + if self.unconfirmed_results.load(Ordering::Acquire) != 0 { + return Err(AgentTreeShutdownFailure::operation_failed( + "tree_suspension", + "persist_child_result", + /*thread_id*/ None, + "unconfirmed_write", + )); + } + Ok(()) + } + + pub(crate) fn suspension_root(&self) -> Option { + match self.shutdown_state.disposition.get() { + Some(TreeDisposition::Suspend { root, .. }) => Some(*root), + _ => None, + } + } + + pub(crate) fn record_root_wake_after(&self, previous_turn_id: String) { + self.shutdown_state + .suspended + .lock() + .unwrap_or_else(PoisonError::into_inner) + .wake_root_after = Some(previous_turn_id); + } + + pub(crate) fn record_suspended_turn(&self, thread_id: ThreadId, turn_id: String) { + let mut tree = self + .shutdown_state + .suspended + .lock() + .unwrap_or_else(PoisonError::into_inner); + if self.suspension_root() == Some(thread_id) { + tree.root_turn_id = Some(turn_id); + tree.wake_root_after = None; + } else { + tree.child_turns.insert(thread_id, turn_id); + } + } + /// 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) { diff --git a/codex-rs/core/src/agent/control/spawn.rs b/codex-rs/core/src/agent/control/spawn.rs index c7e97c62b4..f93bf91b5a 100644 --- a/codex-rs/core/src/agent/control/spawn.rs +++ b/codex-rs/core/src/agent/control/spawn.rs @@ -139,23 +139,31 @@ impl LocalAgentControl { } }; + self.restore_agent_identities(config, &state, descendant_ids) + .await; + } + + pub(super) async fn restore_agent_identities( + &self, + config: &Config, + state: &ThreadManagerState, + descendant_ids: Vec, + ) { + let registry = &self.runtime.registry; // Overlap storage reads, but reserve paths and nicknames in graph order. let mut stored_threads = stream::iter( descendant_ids .into_iter() .filter(|thread_id| registry.agent_metadata_for_thread(*thread_id).is_none()) - .map(|thread_id| { - let state = &state; - async move { - let stored_thread = state - .read_stored_thread(ReadThreadParams { - thread_id, - include_archived: true, - include_history: false, - }) - .await; - (thread_id, stored_thread) - } + .map(|thread_id| async move { + let stored_thread = state + .read_stored_thread(ReadThreadParams { + thread_id, + include_archived: true, + include_history: false, + }) + .await; + (thread_id, stored_thread) }), ) .buffered(/*n*/ 8); @@ -306,9 +314,8 @@ impl LocalAgentControl { let stored_reasoning_effort = stored_thread.reasoning_effort.clone(); let stored_source = stored_thread.source.clone(); let stored_parent_thread_id = stored_thread.parent_thread_id; - let history = load_agent_model_context(&state, thread_id, stored_thread.history_mode) - .await? - .ok_or(CodexErr::ThreadNotFound(thread_id))?; + let history = + load_agent_model_context(&state, thread_id, stored_thread.history_mode).await?; let initial_history = InitialHistory::Resumed(ResumedHistory { history_revision: history.revision, conversation_id: thread_id, @@ -568,9 +575,30 @@ impl LocalAgentControl { }; // Reserving a slot can evict an idle nested parent. Capture its instructions and // analytics metadata alongside its authority before the live parent can disappear. - let residency_slot = self - .reserve_v2_residency_slot(&state, &config, &membership, Some(thread_id)) - .await?; + let reserve = + || self.reserve_v2_residency_slot(&state, &config, &membership, Some(thread_id)); + let residency_slot = match reserve().await { + Err(limit) + if matches!(limit.details(), CodexErrorDetails::AgentLimitReached { .. }) => + { + // A cold parent can report PendingInit despite having spawned this stored child. + // Try that exact parent only after the ordinary eviction policy found no slot. + if let Some((parent, _)) = &parent + && matches!(parent.agent_status().await, AgentStatus::PendingInit) + && matches!( + state + .try_evict_v2_thread(Arc::clone(parent), &membership) + .await, + Ok(ThreadEvictionOutcome::Evicted | ThreadEvictionOutcome::NotCurrent) + ) + { + reserve().await? + } else { + return Err(limit); + } + } + result => result?, + }; match state .resume_thread_with_history_with_source(ResumeThreadWithHistoryOptions { @@ -1137,9 +1165,8 @@ impl LocalAgentControl { .map_err(|err| CodexErr::InvalidRequest(format!("invalid stored agent path: {err}")))?; let resumed_agent_nickname = stored_thread.agent_nickname.clone(); let resumed_agent_role = stored_thread.agent_role.clone(); - let history = load_agent_model_context(&state, thread_id, stored_thread.history_mode) - .await? - .ok_or(CodexErr::ThreadNotFound(thread_id))?; + let history = + load_agent_model_context(&state, thread_id, stored_thread.history_mode).await?; let initial_history = InitialHistory::Resumed(ResumedHistory { history_revision: history.revision, conversation_id: thread_id, diff --git a/codex-rs/core/src/agent/control_tests.rs b/codex-rs/core/src/agent/control_tests.rs index 0c26ab4482..8d4e5c3663 100644 --- a/codex-rs/core/src/agent/control_tests.rs +++ b/codex-rs/core/src/agent/control_tests.rs @@ -4588,6 +4588,8 @@ async fn multi_agent_v2_completion_ignores_dead_direct_parent() { .await .expect("tester thread should exist"); let tester_turn = tester_thread.session.new_default_turn().await; + let writes = &harness.control.runtime.unconfirmed_results; + let before = writes.load(std::sync::atomic::Ordering::Acquire); tester_thread .session .send_event( @@ -4605,6 +4607,7 @@ async fn multi_agent_v2_completion_ignores_dead_direct_parent() { ) .await; + assert_eq!(writes.load(std::sync::atomic::Ordering::Acquire), before); sleep(Duration::from_millis(100)).await; assert!( @@ -4638,6 +4641,18 @@ async fn multi_agent_v2_completion_ignores_dead_direct_parent() { assert!(!has_subagent_notification(root_history.raw_items())); } +#[test] +fn terminal_result_writes_require_confirmation_after_the_parent_claims_mail() { + let count = Arc::new(std::sync::atomic::AtomicUsize::default()); + let publish = || MessageDelivery::publish(ThreadId::new(), &[], Some(&count)).unwrap(); + let parent = ThreadId::new(); + drop(publish().consumed(parent)); + let saved = publish(); + saved.consumed(parent).unwrap().confirm(); + assert!(saved.consumed(parent).is_none()); + assert_eq!(count.load(std::sync::atomic::Ordering::Acquire), 1); +} + #[tokio::test] async fn multi_agent_v2_completion_queues_message_for_direct_parent() { let harness = AgentControlHarness::new().await; @@ -5122,6 +5137,11 @@ async fn resume_thread_subagent_restores_stored_metadata() { Vec::new() ); + assert_eq!( + original_snapshot.history_mode, + codex_protocol::protocol::ThreadHistoryMode::Legacy + ); + let before_resume = thread_store.calls().await; let resumed_thread_id = harness .control .resume_agent_from_rollout( @@ -5137,6 +5157,13 @@ async fn resume_thread_subagent_restores_stored_metadata() { ) .await .expect("resume should succeed"); + // Remote stores reject unbounded reads on their metadata API. + let after_resume = thread_store.calls().await; + assert_eq!( + after_resume.read_thread_with_history, + before_resume.read_thread_with_history + ); + assert!(after_resume.load_history > before_resume.load_history); assert_eq!(resumed_thread_id, child_thread_id); let resumed_thread = harness diff --git a/codex-rs/core/src/agent/history.rs b/codex-rs/core/src/agent/history.rs index f82446b371..d985b54344 100644 --- a/codex-rs/core/src/agent/history.rs +++ b/codex-rs/core/src/agent/history.rs @@ -38,7 +38,6 @@ use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::SubAgentSource; use codex_protocol::protocol::ThreadHistoryMode; use codex_thread_store::LoadThreadHistoryParams; -use codex_thread_store::ReadThreadParams; fn keep_forked_rollout_item(item: &RolloutItem, preserve_context_baselines: bool) -> bool { match item { @@ -132,29 +131,14 @@ pub(super) async fn load_agent_model_context( state: &ThreadManagerState, thread_id: ThreadId, history_mode: ThreadHistoryMode, -) -> CodexResult> { +) -> CodexResult { + let params = LoadThreadHistoryParams { + thread_id, + include_archived: true, + }; match history_mode { - ThreadHistoryMode::Legacy => Ok(state - .read_stored_thread(ReadThreadParams { - thread_id, - include_archived: true, - include_history: true, - }) - .await? - .history - .map(|history| codex_thread_store::StoredModelContext { - thread_id, - revision: history.revision, - items: history.items, - })), - ThreadHistoryMode::Paginated => Ok(Some( - state - .load_latest_model_context(LoadThreadHistoryParams { - thread_id, - include_archived: true, - }) - .await?, - )), + ThreadHistoryMode::Legacy => state.load_legacy_model_context(params).await, + ThreadHistoryMode::Paginated => state.load_latest_model_context(params).await, } } @@ -227,11 +211,6 @@ pub(super) async fn prepare_forked_history( let mut forked_rollout_items = load_agent_model_context(state, parent_thread_id, parent_history_mode) .await? - .ok_or_else(|| { - CodexErr::Fatal(format!( - "parent thread history unavailable for fork: {parent_thread_id}" - )) - })? .items; let multi_agent_v2_usage_hint_texts_to_filter: Vec = diff --git a/codex-rs/core/src/agent/registry.rs b/codex-rs/core/src/agent/registry.rs index eeed8f450d..2385a98748 100644 --- a/codex-rs/core/src/agent/registry.rs +++ b/codex-rs/core/src/agent/registry.rs @@ -42,6 +42,7 @@ struct ActiveAgents { /// Effective configuration and executor bindings retained while a child's runtime is unloaded. #[derive(Clone)] pub(crate) struct EvictedEnvironments { + pub(crate) history_mode: codex_protocol::protocol::ThreadHistoryMode, pub(crate) config: Arc, pub(crate) has_independent_service_tier: bool, pub(crate) selections: Vec, diff --git a/codex-rs/core/src/agent_communication.rs b/codex-rs/core/src/agent_communication.rs index 141737c50d..ebcd23cfe5 100644 --- a/codex-rs/core/src/agent_communication.rs +++ b/codex-rs/core/src/agent_communication.rs @@ -14,6 +14,8 @@ use std::ops::Deref; use std::sync::Arc; use std::sync::Mutex; use std::sync::PoisonError; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; pub(crate) static PENDING_MAILBOX_MESSAGES: codex_diagnostics::Gauge = codex_diagnostics::Gauge::new("core.mailbox.pending"); @@ -26,6 +28,7 @@ pub(crate) enum AgentCommunicationKind { Message, Followup, Result, + CompletionWatcher, } impl AgentCommunicationKind { @@ -34,14 +37,14 @@ impl AgentCommunicationKind { Self::Spawn => "spawn", Self::Message => "message", Self::Followup => "followup", - Self::Result => "result", + Self::Result | Self::CompletionWatcher => "result", } } } #[derive(Clone, Debug, PartialEq, Eq)] pub(crate) struct AgentCommunicationContext { - kind: AgentCommunicationKind, + pub(crate) kind: AgentCommunicationKind, pub(crate) sender_thread_id: ThreadId, } @@ -140,12 +143,33 @@ impl PendingAgentMessage { delivery, } } + + pub(crate) fn is_terminal_result(&self) -> bool { + self.delivery + .as_ref() + .is_some_and(MessageDelivery::awaits_result_write) + } + + pub(crate) fn consumed(&self, recipient: ThreadId) -> Option { + self.delivery.as_ref()?.consumed(recipient) + } +} + +/// Only a confirmed append releases a terminal result. Dropping this receipt +/// leaves the write outstanding, so an ambiguous handoff fails. +pub(crate) struct ResultWrite(Arc); + +impl ResultWrite { + pub(crate) fn confirm(self) { + self.0.fetch_sub(1, Ordering::Release); + } } /// The publication state retained until this delivery is consumed or abandoned. struct MessagePublication { input: MessageLifecycleInput, observers: Box<[Arc]>, + result_write: Option, } /// Clones follow one logical delivery through queues and share its consumption state. @@ -156,7 +180,11 @@ impl MessageDelivery { pub(crate) fn publish( sender: ThreadId, observers: &[Arc], - ) -> Self { + unconfirmed_results: Option<&Arc>, + ) -> Option { + if observers.is_empty() && unconfirmed_results.is_none() { + return None; + } let input = MessageLifecycleInput { sender, data: ExtensionData::new(uuid::Uuid::now_v7().to_string()), @@ -164,21 +192,37 @@ impl MessageDelivery { for observer in observers { observer.on_publish(&input); } + let result_write = unconfirmed_results.map(|count| { + count.fetch_add(1, Ordering::Relaxed); + ResultWrite(Arc::clone(count)) + }); let publication = MessagePublication { input, observers: observers.into(), + result_write, }; - Self(Arc::new(Mutex::new(Some(publication)))) + Some(Self(Arc::new(Mutex::new(Some(publication))))) } - /// Commit this delivery and release its publication state, even if clones remain. - pub(crate) fn consumed(&self, recipient: ThreadId) { - let publication = self.0.lock().unwrap_or_else(PoisonError::into_inner).take(); - if let Some(publication) = publication { - for observer in &publication.observers { - observer.on_consumed(&publication.input, recipient); - } + pub(crate) fn awaits_result_write(&self) -> bool { + self.0 + .lock() + .unwrap_or_else(PoisonError::into_inner) + .as_ref() + .is_some_and(|publication| publication.result_write.is_some()) + } + + /// The first consumer owns both the publication hooks and the write receipt. + pub(crate) fn consumed(&self, recipient: ThreadId) -> Option { + let publication = self + .0 + .lock() + .unwrap_or_else(PoisonError::into_inner) + .take()?; + for observer in &publication.observers { + observer.on_consumed(&publication.input, recipient); } + publication.result_write } } diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index 6e375dca71..1ad0ce3d68 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -587,7 +587,9 @@ impl CodexThread { /// Starts a new internal continuation turn when idle, including in Plan mode. /// Rejects if a newer task has started, even if it has already finished. - /// The input must be a response item; it is never treated as user authorization. + /// Supply an internal response item, or empty input to wake on a result already + /// saved in history. An empty wake inherits the preceding turn's attribution + /// and adds no message or user authorization. pub async fn continue_turn_if_idle( &self, request: impl Into>, diff --git a/codex-rs/core/src/config/mod.rs b/codex-rs/core/src/config/mod.rs index 5c732876f6..f958572c8a 100644 --- a/codex-rs/core/src/config/mod.rs +++ b/codex-rs/core/src/config/mod.rs @@ -1680,6 +1680,21 @@ async fn config_toml_from_layers(layers: &ConfigLayerStack) -> std::io::Result codex_thread_store::ThreadPersistenceMetadata { + use codex_protocol::protocol::ThreadMemoryMode; + codex_thread_store::ThreadPersistenceMetadata { + cwd: Some(self.cwd.to_path_buf()), + model_provider: self.model_provider_id.clone(), + memory_mode: if self.memories.generate_memories { + ThreadMemoryMode::Enabled + } else { + ThreadMemoryMode::Disabled + }, + } + } + /// Capture storage and sync preparation use one session-static rollout choice. pub(crate) fn guardian_context_budgeting(&self) -> codex_guardian_context::ContextBudgeting { if self.guardian_sync_elastic_budget.is_some() { diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index 12d65a0b78..109cb6e4dd 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -188,11 +188,13 @@ pub use thread_manager::AgentTreeShutdownFailureReason; pub use thread_manager::AgentTreeShutdownReport; pub use thread_manager::ForkSnapshot; pub use thread_manager::InternalSessionParent; +pub use thread_manager::LoadedSuspendedAgent; pub use thread_manager::McpRefreshReason; pub use thread_manager::NewThread; pub use thread_manager::PromptCacheKey; pub use thread_manager::StartThreadOptions; pub use thread_manager::StartupModels; +pub use thread_manager::SuspendedAgentTree; pub use thread_manager::ThreadManager; pub use thread_manager::ThreadRegistration; pub use thread_manager::ThreadShutdownReport; diff --git a/codex-rs/core/src/session/handlers.rs b/codex-rs/core/src/session/handlers.rs index 746ff4f220..670d1f484b 100644 --- a/codex-rs/core/src/session/handlers.rs +++ b/codex-rs/core/src/session/handlers.rs @@ -316,6 +316,9 @@ pub(super) async fn shutdown_session_runtime(sess: &Arc) { if let Some(reporting) = &sess.outcome_reporting { reporting.handoffs.close(); } + let runtime = &sess.services.local_agent_runtime; + let wakes_root = runtime.suspension_root() == Some(sess.thread_id) + && (sess.has_outstanding_durable_sleep() || sess.should_wake_on_agent_messages().await); let startup_prewarm = { let mut state = sess.state.lock().await; // Stop admission and take the current warmup together so resume cannot replace it. @@ -330,6 +333,49 @@ pub(super) async fn shutdown_session_runtime(sess: &Arc) { if let Some(reporting) = &sess.outcome_reporting { reporting.handoffs.wait().await; } + if runtime.suspension_root().is_some() { + let pending: Vec<_> = sess + .input_queue + .drain_mailbox() + .await + .into_iter() + .map(|mail| mail.communication) + .filter(crate::PendingAgentMessage::is_terminal_result) + .collect(); + if !pending.is_empty() && sess.live_thread().is_some() { + let turn_id = { + let state = sess.state.lock().await; + state.last_started_turn_id.clone().or_else(|| { + state + .turn_attribution + .as_ref() + .map(|turn| turn.turn_id.clone()) + }) + }; + if wakes_root { + if let Some(turn_id) = turn_id.as_ref() { + runtime.record_root_wake_after(turn_id.clone()); + } else { + runtime.record_shutdown_failure(AgentTreeShutdownFailure::operation_failed( + "tree_suspension", + "record_root_wake", + Some(sess.thread_id), + "missing_previous_turn", + )); + } + } + let context = sess + .new_turn_with_default_settings( + turn_id.unwrap_or_else(super::new_submission_id), + Default::default(), + ) + .await; + for result in pending { + sess.record_inter_agent_communication(&context, context.model_info(), result) + .await; + } + } + } super::async_compaction::shutdown(sess).await; let shell_snapshot_prewarm = sess.state.lock().await.shell_snapshot_prewarm.take(); if let Some(shell_snapshot_prewarm) = shell_snapshot_prewarm { @@ -491,7 +537,28 @@ pub(super) async fn submission_loop( let sub = tokio::select! { biased; _ = sess.services.local_agent_runtime.shutdown.cancelled() => { - shutdown_received = shutdown(&sess, super::new_submission_id()).await; + if sess.services.local_agent_runtime.suspension_root().is_some() { + // Keep accepting results until all direct children stop, then drain + // buffered submissions before deciding which parent work to hand off. + let runtime = &sess.services.local_agent_runtime; + tokio::join!( + async { + runtime.wait_for_children(sess.thread_id).await; + rx_sub.close(); + }, + async { + while let Ok(sub) = rx_sub.recv().await { + if let Op::InterAgentCommunication { communication, start_options } = sub.op + && sub.message_delivery.as_ref().is_some_and(crate::agent_communication::MessageDelivery::awaits_result_write) + { + let mail = crate::PendingAgentMessage::new(communication, sub.message_delivery); + inter_agent_communication(&sess, sub.id, mail, start_options).await; + } + } + }, + ); + } + shutdown_received = super::turn_suspension::stop_tree_member(&sess).await; break; } update = async { diff --git a/codex-rs/core/src/session/input_queue.rs b/codex-rs/core/src/session/input_queue.rs index 955491e708..f5aa2f69ef 100644 --- a/codex-rs/core/src/session/input_queue.rs +++ b/codex-rs/core/src/session/input_queue.rs @@ -269,15 +269,25 @@ impl InputQueue { pending.drain(..).collect() } - pub(crate) async fn drain_mailbox_input_items(&self) -> (Vec, TurnStartOptions) { - let pending_mails = self.drain_mailbox().await; + pub(crate) async fn pending_wake_options(&self) -> Option { + let mut mails = self.mailbox_pending_mails.lock().await; + self.read_mailbox(&mut mails); + let mut options = Self::trigger_turn_options(&mails)?; + if let Some(mail) = mails.iter().find(|mail| mail.communication.trigger_turn) { + options.initiating_agent_path = Some(mail.communication.author.clone()); + } + Some(options) + } + + fn trigger_turn_options( + pending_mails: &VecDeque, + ) -> Option { // A later follow-up supersedes the earlier choice, including an omitted choice. let mut start_options = pending_mails .iter() .rev() .find(|mail| mail.communication.trigger_turn) - .map(|mail| mail.start_options.clone()) - .unwrap_or_default(); + .map(|mail| mail.start_options.clone())?; start_options.parent_turn_id = pending_mails .iter() .filter(|mail| mail.communication.trigger_turn) @@ -296,8 +306,15 @@ impl InputQueue { .filter(|id| !id.trim().is_empty()) }) .map(str::to_string); + Some(start_options) + } + + pub(crate) async fn drain_mailbox_input_items(&self) -> (Vec, TurnStartOptions) { + let mut pending_mails = self.mailbox_pending_mails.lock().await; + self.read_mailbox(&mut pending_mails); + let start_options = Self::trigger_turn_options(&pending_mails).unwrap_or_default(); let items = pending_mails - .into_iter() + .drain(..) .map(|mail| TurnInput::InterAgentCommunication(mail.communication)) .collect(); (items, start_options) @@ -343,12 +360,34 @@ impl InputQueue { /// Clear any pending waiters and input buffered for the current turn. pub(crate) async fn clear_pending(&self, active_turn: &ActiveTurn) { - let mut turn_state = active_turn.turn_state.lock().await; - turn_state.clear_pending_waiters(); // Pending mail can own extension attachments whose destructors need the lock released. - let pending_input = std::mem::take(&mut turn_state.pending_input); - drop(turn_state); - drop(pending_input); + let pending_input = { + let mut turn_state = active_turn.turn_state.lock().await; + turn_state.clear_pending_waiters(); + std::mem::take(&mut turn_state.pending_input) + }; + // A terminal result still in the queue has no writer yet. Keep it available + // to the next turn or to the source writer during a graceful drain. + let results: Vec<_> = pending_input + .items + .into_iter() + .filter_map(|input| match input { + TurnInput::InterAgentCommunication(message) if message.is_terminal_result() => { + Some(message) + } + _ => None, + }) + .collect(); + if !results.is_empty() { + let mut pending = self.mailbox_pending_mails.lock().await; + for communication in results.into_iter().rev() { + pending.push_front(PendingMailboxCommunication { + communication, + start_options: TurnStartOptions::default(), + _diagnostics_guard: PENDING_MAILBOX_MESSAGES.track(), + }); + } + } } pub(crate) async fn defer_mailbox_delivery_to_next_turn( diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index b932d39dca..5d02f2b6d5 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -195,7 +195,6 @@ use codex_thread_store::LocalThreadStore; use codex_thread_store::PersistContext; use codex_thread_store::ReadThreadParams; use codex_thread_store::ResumeThreadParams; -use codex_thread_store::ThreadPersistenceMetadata; use codex_thread_store::ThreadStore; use codex_thread_store::ThreadStoreResumePreparation; use codex_utils_audio::prepare_response_items as prepare_audio_response_items; @@ -396,6 +395,7 @@ use codex_protocol::protocol::DeprecationNoticeEvent; use codex_protocol::protocol::Event; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::ExecApprovalRequestEvent; +use codex_protocol::protocol::InterAgentCommunication; use codex_protocol::protocol::ModelRerouteEvent; use codex_protocol::protocol::ModelRerouteReason; use codex_protocol::protocol::ModelVerification; @@ -2832,7 +2832,7 @@ impl Session { let reporter = self .outcome_reporting .as_ref() - .map(|reporting| &reporting.reporter); + .and_then(|reporting| reporting.reporter.as_ref()); if !notify_native_parent && reporter.is_none() { return; } @@ -4656,13 +4656,11 @@ impl Session { })) } - pub(crate) async fn record_inter_agent_communication( - &self, - turn_context: &TurnContext, - model_info: &ModelInfo, - communication: impl Into, - ) { - let communication = communication.into(); + /// Builds the same history item whether the recipient is live or has been evicted. + pub(crate) fn inter_agent_history_item( + communication: &InterAgentCommunication, + config: &Config, + ) -> ResponseItem { let mut response_item = communication.to_model_input_item(); if let ResponseItem::AgentMessage { content, @@ -4676,8 +4674,7 @@ impl Session { .get_or_insert_with(|| { vec![ ContentItemMetadata::tool(Some( - turn_context - .config + config .multi_agent_v2 .tool_namespace .clone() @@ -4688,6 +4685,17 @@ impl Session { ] }); } + response_item + } + + pub(crate) async fn record_inter_agent_communication( + &self, + turn_context: &TurnContext, + model_info: &ModelInfo, + communication: impl Into, + ) { + let communication = communication.into(); + let response_item = Self::inter_agent_history_item(&communication, &turn_context.config); let (items, _) = self .prepare_conversation_items_for_history( turn_context, @@ -4728,16 +4736,21 @@ impl Session { model_info.truncation_policy.into(), ); } - if let Some(delivery) = &communication.delivery { - delivery.consumed(self.thread_id); + let result_write = communication.consumed(self.thread_id); + let saved = self + .persist_rollout_items(&[ + RolloutItem::InterAgentCommunicationMetadata { + trigger_turn: communication.trigger_turn, + }, + RolloutItem::ResponseItem(response_item), + ]) + .await; + if saved + && self.live_thread().is_some() + && let Some(write) = result_write + { + write.confirm(); } - self.persist_rollout_items(&[ - RolloutItem::InterAgentCommunicationMetadata { - trigger_turn: communication.trigger_turn, - }, - RolloutItem::ResponseItem(response_item), - ]) - .await; drop(boundary); self.send_raw_response_items(turn_context, &items).await; } @@ -5394,6 +5407,21 @@ impl Session { && let Err(e) = live_thread.append_items(items).await { error!("failed to record rollout items: {e:#}"); + if self + .services + .local_agent_runtime + .suspension_root() + .is_some() + { + self.services.local_agent_runtime.record_shutdown_failure( + crate::thread_manager::AgentTreeShutdownFailure::operation_failed( + "tree_suspension", + "append_history", + Some(self.thread_id), + crate::thread_manager::thread_store_error_kind(&e), + ), + ); + } return false; } true diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 56dd0214cf..db260bfd21 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -60,9 +60,10 @@ use tokio_util::task::TaskTracker; type McpToolApprovalMetadataMap = HashMap<(codex_mcp::McpServerKey, String), std::sync::Weak>; -/// A thread's installed reporter and the tasks that may still hand it an outcome. +/// Joins terminal publishers before their session writer closes. Native child results +/// and optional host reporters both use this barrier. pub(crate) struct AgentOutcomeReporting { - pub(super) reporter: Arc, + pub(crate) reporter: Option>, pub(crate) handoffs: TaskTracker, } @@ -634,10 +635,15 @@ async fn warm_plugins_and_skills_for_session_init( } impl Session { - pub(crate) fn outcome_handoffs_closed(&self) -> bool { - self.outcome_reporting - .as_ref() - .is_some_and(|reporting| reporting.handoffs.is_closed()) + pub(crate) fn turn_starts_closed(&self) -> bool { + self.services + .local_agent_runtime + .suspension_root() + .is_some() + || self + .outcome_reporting + .as_ref() + .is_some_and(|reporting| reporting.handoffs.is_closed()) } pub(crate) fn new_item_id_generator() -> ItemIdGenerator { @@ -1143,15 +1149,7 @@ impl Session { .window_id .clone(), runtime_workspace_roots: Some(config.workspace_roots.clone()), - metadata: ThreadPersistenceMetadata { - cwd: Some(config.cwd.to_path_buf()), - model_provider: config.model_provider_id.clone(), - memory_mode: if config.memories.generate_memories { - ThreadMemoryMode::Enabled - } else { - ThreadMemoryMode::Disabled - }, - }, + metadata: config.thread_persistence_metadata(), }; if is_paginated_subagent && matches!( @@ -1180,15 +1178,7 @@ impl Session { rollout_path: resumed_history.rollout_path.clone(), history: Some(resumed_history.history.clone()), include_archived: true, - metadata: ThreadPersistenceMetadata { - cwd: Some(config.cwd.to_path_buf()), - model_provider: config.model_provider_id.clone(), - memory_mode: if config.memories.generate_memories { - ThreadMemoryMode::Enabled - } else { - ThreadMemoryMode::Disabled - }, - }, + metadata: config.thread_persistence_metadata(), }; let store = Arc::clone(&thread_store); let history_mode = session_configuration.history_mode; @@ -1797,10 +1787,14 @@ impl Session { .await; } - let outcome_reporting = thread_extension_data + let reporter = thread_extension_data .get::>() - .map(|reporter| AgentOutcomeReporting { - reporter: Arc::clone(&reporter), + .map(|reporter| Arc::clone(reporter.as_ref())); + // The turn's resolved version determines whether it publishes a native result. + let local_durable = live_thread.is_some() && !has_provided_agent_control; + let outcome_reporting = + (reporter.is_some() || local_durable).then(|| AgentOutcomeReporting { + reporter, handoffs: TaskTracker::new(), }); let (agent_runtime, mut startup_error) = diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 6c330fcbb0..8e6bdfb309 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -2,6 +2,7 @@ use crate::McpRefreshReason; use crate::agent::LocalAgentControl; use codex_mcp::McpServerKey; use codex_protocol::ContextWindowId; +use codex_thread_store::ThreadPersistenceMetadata; #[path = "compaction_replay_tests.rs"] mod compaction_replay_tests; @@ -6254,12 +6255,10 @@ async fn standalone_settings_invalidate_continuation_before_delivering_acceptanc .tx_event .send(Event { id: "occupied".into(), - msg: EventMsg::ThreadSettingsApplied( - codex_protocol::protocol::ThreadSettingsAppliedEvent { - thread_id: Some(session.thread_id()), - thread_settings: session.thread_settings_snapshot().await, - }, - ), + msg: EventMsg::ThreadSettingsApplied(ThreadSettingsAppliedEvent { + thread_id: Some(session.thread_id()), + thread_settings: session.thread_settings_snapshot().await, + }), }) .await .expect("fill event channel"); @@ -6271,7 +6270,7 @@ async fn standalone_settings_invalidate_continuation_before_delivering_acceptanc id: "settings".into(), turn_extension_init: None, op: Op::ThreadSettings { - thread_settings: codex_protocol::protocol::ThreadSettingsOverrides::default(), + thread_settings: ThreadSettingsOverrides::default(), reply: Some(reply), }, trace: None, @@ -6356,7 +6355,7 @@ async fn compaction_persists_resume_metadata_and_companion_records() { }), previous_turn_settings: Some(previous_turn_settings), }; - let expected_settings = codex_protocol::protocol::ThreadSettingsAppliedEvent { + let expected_settings = ThreadSettingsAppliedEvent { thread_id: Some(session.thread_id()), thread_settings: session.thread_settings_snapshot().await, }; diff --git a/codex-rs/core/src/session/turn_input.rs b/codex-rs/core/src/session/turn_input.rs index fb7575936d..3c9e91b6fa 100644 --- a/codex-rs/core/src/session/turn_input.rs +++ b/codex-rs/core/src/session/turn_input.rs @@ -306,13 +306,20 @@ pub(super) async fn handle( TurnInputMode::ContinueIfIdle { expected_previous_turn_id, } => { - if !matches!( - &request.request.input, - SubmittedTurnInput::ResponseItem(_) | SubmittedTurnInput::AnnotatedResponseItem(_) - ) { - return Err(CodexErr::InvalidRequest( - "continuation requires internal response input".to_string(), - )); + let mut request = request; + match &request.request.input { + SubmittedTurnInput::ResponseItem(_) + | SubmittedTurnInput::AnnotatedResponseItem(_) => {} + SubmittedTurnInput::UserInput { content, .. } if content.is_empty() => { + request.request.start = session.agent_message_wake_options().await; + } + SubmittedTurnInput::UserInput { .. } + | SubmittedTurnInput::InterAgentCommunication(_) => { + return Err(CodexErr::InvalidRequest( + "continuation requires internal response input or an empty wake" + .to_string(), + )); + } } start_if_idle( session, @@ -501,13 +508,9 @@ async fn start_or_steer( .await, ); } - session - .spawn_task(Arc::clone(&turn_context), task_input, RegularTask::new()) - .await; - Ok(TurnInputSubmission::Started { - turn_id: submission_id, - root_turn_id: turn_context.root_turn_id(), - }) + Ok(session + .spawn_task(turn_context, task_input, RegularTask::new()) + .await) } Err(reason) => Ok(TurnInputSubmission::NotSubmitted { reason }), } @@ -537,7 +540,9 @@ async fn start_if_idle( responsesapi_client_metadata, .. } = request; - if session.input_queue.has_trigger_turn_mailbox_items().await { + if expected_previous_turn_id.is_none() + && session.input_queue.has_trigger_turn_mailbox_items().await + { return Ok(TurnInputSubmission::NotSubmitted { reason: NotSubmittedReason::PendingTriggerTurn, }); @@ -575,18 +580,23 @@ async fn start_if_idle( let turn_state = { let mut active_turn = session.active_turn.lock().await; + if let Some(expected) = expected_previous_turn_id { + let state = session.state.lock().await; + // A registered successor already owns this work, even if it is still running. + // An unregistered reservation can still fail to start and must be retried. + if state.last_started_turn_id.as_ref() != Some(&expected) + && (state.last_started_turn_id.is_some() || active_turn.is_none()) + { + return Ok(TurnInputSubmission::NotSubmitted { + reason: NotSubmittedReason::Superseded, + }); + } + } if active_turn.is_some() { return Ok(TurnInputSubmission::NotSubmitted { reason: NotSubmittedReason::NotIdle, }); } - if let Some(expected) = expected_previous_turn_id - && session.state.lock().await.last_started_turn_id.as_ref() != Some(&expected) - { - return Ok(TurnInputSubmission::NotSubmitted { - reason: NotSubmittedReason::Superseded, - }); - } let active_turn = active_turn.get_or_insert_with(ActiveTurn::default); Arc::clone(&active_turn.turn_state) }; @@ -680,13 +690,9 @@ async fn start_if_idle( } } } - session - .start_task(Arc::clone(&turn_context), task_input, RegularTask::new()) - .await; - Ok(TurnInputSubmission::Started { - turn_id: submission_id, - root_turn_id: turn_context.root_turn_id(), - }) + Ok(session + .start_task(turn_context, task_input, RegularTask::new()) + .await) } async fn steer( @@ -814,7 +820,7 @@ impl Session { &self, ) -> Option<(Arc>, TurnStartOptions)> { let mut active_turn = self.active_turn.lock().await; - if active_turn.is_some() || self.outcome_handoffs_closed() { + if active_turn.is_some() || self.turn_starts_closed() { return None; } let has_trigger_turn = self.input_queue.has_trigger_turn_mailbox_items().await; @@ -828,17 +834,24 @@ impl Session { let previous_options = if has_trigger_turn { Default::default() } else { - self.state - .lock() - .await - .turn_attribution - .as_ref() - .map(codex_history::TurnAttribution::start_options) - .unwrap_or_default() + self.agent_message_wake_options().await }; Some((Arc::clone(&active_turn.turn_state), previous_options)) } + async fn agent_message_wake_options(&self) -> TurnStartOptions { + let state = self.state.lock().await; + let mut options = state + .turn_attribution + .as_ref() + .map(codex_history::TurnAttribution::start_options) + .unwrap_or_default(); + options.cyber_access_program = state + .reference_context_item() + .and_then(|context| context.cyber_access_program); + options + } + /// Revalidates a reserved wake and selects its settings before any mail is consumed. #[expect( clippy::await_holding_invalid_type, @@ -871,11 +884,8 @@ impl Session { Some(configuration.clone()) } - async fn should_wake_on_agent_messages(&self) -> bool { - if !self.features.enabled(Feature::WakeOnAgentMessages) - || self.is_interrupted() - || !self.input_queue.has_pending_mailbox_items().await - { + pub(crate) async fn should_wake_on_agent_messages(&self) -> bool { + if !self.features.enabled(Feature::WakeOnAgentMessages) || self.is_interrupted() { return false; } diff --git a/codex-rs/core/src/session/turn_input_tests.rs b/codex-rs/core/src/session/turn_input_tests.rs index 219a4229a4..2d49feb256 100644 --- a/codex-rs/core/src/session/turn_input_tests.rs +++ b/codex-rs/core/src/session/turn_input_tests.rs @@ -250,6 +250,29 @@ async fn start_only_rejects_active_turn_without_injecting() { session.abort_all_tasks(TurnAbortReason::Interrupted).await; } +#[tokio::test] +async fn saved_message_wake_checks_the_last_started_turn_during_admission() { + let (session, _, _rx) = make_session_and_context_with_rx().await; + *session.active_turn.lock().await = Some(ActiveTurn::default()); + for (last_started, reason) in [ + ("original", NotSubmittedReason::NotIdle), + ("successor", NotSubmittedReason::Superseded), + ] { + session.state.lock().await.last_started_turn_id = Some(last_started.into()); + let result = handle( + &session, + TurnInputRequest::user_input(vec![]), + TurnInputMode::ContinueIfIdle { + expected_previous_turn_id: "original".into(), + }, + "wake".into(), + ) + .await + .unwrap(); + assert_eq!(result, TurnInputSubmission::NotSubmitted { reason }); + } +} + #[tokio::test] async fn recovery_rejects_active_turn_without_injecting_or_applying_settings() { let (session, turn_context, _rx) = make_session_and_context_with_rx().await; diff --git a/codex-rs/core/src/session/turn_suspension.rs b/codex-rs/core/src/session/turn_suspension.rs index 2a49328c9b..b3318efb0d 100644 --- a/codex-rs/core/src/session/turn_suspension.rs +++ b/codex-rs/core/src/session/turn_suspension.rs @@ -1,10 +1,12 @@ use super::handlers; use super::session::Session; use crate::state::TaskKind; +use crate::thread_manager::AgentTreeShutdownFailure; use codex_protocol::error::CodexErr; use codex_protocol::error::Result as CodexResult; use codex_protocol::protocol::Event; use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::MultiAgentVersion; use codex_protocol::turn_input::SuspendTurnOutcome; use std::sync::Arc; use std::time::Duration; @@ -120,3 +122,39 @@ pub(super) async fn suspend_turn_and_shutdown( .await; Ok(SuspendTurnOutcome::Suspended { turn_id }) } + +pub(super) async fn stop_tree_member(session: &Arc) -> bool { + let runtime = &session.services.local_agent_runtime; + let submission_id = super::new_submission_id(); + if let Some(root) = runtime.suspension_root() + && session.live_thread().is_some() + && (root == session.thread_id + || session.multi_agent_version() == Some(MultiAgentVersion::V2) + && session + .state + .lock() + .await + .session_configuration + .session_source + .get_agent_path() + .is_some()) + { + match suspend_turn_and_shutdown(session, submission_id.clone()).await { + Ok(SuspendTurnOutcome::Suspended { turn_id }) => { + runtime.record_suspended_turn(session.thread_id, turn_id); + return true; + } + Ok(SuspendTurnOutcome::NotActive | SuspendTurnOutcome::UnsupportedTask) => {} + Err(error) => { + runtime.record_shutdown_failure(AgentTreeShutdownFailure::operation_failed( + "tree_suspension", + "suspend_regular_turn", + Some(session.thread_id), + "suspension_failed", + )); + warn!(thread_id = %session.thread_id, %error, "failed to suspend agent turn"); + } + } + } + handlers::shutdown(session, submission_id).await +} diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index 506b17877f..44f195ef57 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -61,6 +61,8 @@ use codex_protocol::protocol::TurnAbortReason; use codex_protocol::protocol::TurnAbortedEvent; use codex_protocol::protocol::TurnCompleteEvent; use codex_protocol::protocol::WarningEvent; +use codex_protocol::turn_input::NotSubmittedReason; +use codex_protocol::turn_input::TurnInputSubmission; use codex_thread_store::PersistContext; pub(crate) use admission::HostTurnStarted; @@ -280,22 +282,22 @@ impl Session { turn_context: Arc, input: Vec, task: T, - ) { + ) -> TurnInputSubmission { self.abort_all_tasks(TurnAbortReason::Replaced).await; self.clear_connector_selection().await; - self.start_task(turn_context, input, task).await; + self.start_task(turn_context, input, task).await } #[expect( clippy::await_holding_invalid_type, - reason = "record the started turn atomically with task registration" + reason = "register the turn and its pending input atomically with the suspension fence" )] pub(crate) async fn start_task( self: &Arc, turn_context: Arc, input: Vec, task: T, - ) { + ) -> TurnInputSubmission { // Replace the cache; in-flight acquisitions retain their previous Arc. self.services .thread_extension_data @@ -335,11 +337,12 @@ impl Session { let cancellation_token = CancellationToken::new(); let done = Arc::new(Notify::new()); - let (pending_items, _) = self.input_queue.drain_mailbox_input_items().await; let turn_state = { let mut active = self.active_turn.lock().await; - if self.outcome_handoffs_closed() { - return; + if self.turn_starts_closed() { + return TurnInputSubmission::NotSubmitted { + reason: NotSubmittedReason::ServerDraining, + }; } let turn = active.get_or_insert_with(ActiveTurn::default); debug_assert!(turn.task.is_none()); @@ -347,9 +350,6 @@ impl Session { Arc::clone(&turn.turn_state) }; turn_state.lock().await.token_usage_at_turn_start = token_usage_at_turn_start.clone(); - self.input_queue - .extend_pending_input_for_turn_state(turn_state.as_ref(), pending_items) - .await; self.emit_turn_start_lifecycle( turn_context.as_ref(), Some(&token_usage_at_turn_start), @@ -359,16 +359,35 @@ impl Session { let mut active = self.active_turn.lock().await; // Shutdown closes registration before taking this same lock to abort active work. - let outcome_handoff = if let Some(reporting) = &self.outcome_reporting { - if reporting.handoffs.is_closed() { - return; - } - Some(reporting.handoffs.token()) - } else { - None + if self.turn_starts_closed() { + return TurnInputSubmission::NotSubmitted { + reason: NotSubmittedReason::ServerDraining, + }; + } + let outcome_handoff = self + .outcome_reporting + .as_ref() + .filter(|reporting| { + reporting.reporter.is_some() + || turn_context.multi_agent_version == MultiAgentVersion::V2 + && (!turn_context.session_source.is_non_root_agent() + || turn_context.session_source.get_agent_path().is_some()) + }) + .map(|reporting| reporting.handoffs.token()); + let Some(turn) = active + .as_mut() + .filter(|turn| Arc::ptr_eq(&turn.turn_state, &turn_state)) + else { + return TurnInputSubmission::NotSubmitted { + reason: NotSubmittedReason::Superseded, + }; }; - let turn = active.get_or_insert_with(ActiveTurn::default); debug_assert!(turn.task.is_none()); + // Keep mail in its queue until a task owns it. Suspension can save it if it wins this lock. + let (pending_items, _) = self.input_queue.drain_mailbox_input_items().await; + self.input_queue + .extend_pending_input_for_turn_state(&turn_state, pending_items) + .await; // Record lineage only when the task starts: settings changes can invalidate a reservation. self.record_started_turn( &turn_context.sub_id, @@ -494,6 +513,10 @@ impl Session { _timer: timer, }; turn.task = Some(running_task); + TurnInputSubmission::Started { + turn_id: turn_context.sub_id.clone(), + root_turn_id: turn_context.root_turn_id(), + } } /// Returns whether an extension has marked this thread as durably asleep. @@ -559,19 +582,10 @@ impl Session { } None }; - let (input, mut start_options) = - self.input_queue.get_pending_input(&self.active_turn).await; - let triggered = input.iter().any( - |item| matches!(item, TurnInput::InterAgentCommunication(mail) if mail.trigger_turn), - ); - if !triggered { - // Queue-only wakes retain the previous turn context without selecting a new task. - start_options = previous_options; - start_options.cyber_access_program = self - .reference_context_item() - .await - .and_then(|context| context.cyber_access_program); - } + // Read the routing options now and leave the mail queued until task registration. + let wake_options = self.input_queue.pending_wake_options().await; + let triggered = wake_options.is_some(); + let start_options = wake_options.unwrap_or(previous_options); let options = NewTurnContextOptions { direct_only_tool_namespaces: start_options.direct_only_tool_namespaces, final_output_json_schema: start_options.final_output_json_schema, @@ -590,18 +604,7 @@ impl Session { turn_context.turn_metadata_state.set_turn_trigger(trigger); } if let Some(id) = start_options.parent_turn_id { - if let Some(initiating_agent_path) = input - .iter() - .find_map(|item| { - let TurnInput::InterAgentCommunication(communication) = item else { - return None; - }; - communication - .trigger_turn - .then(|| communication.author.clone()) - }) - .or(start_options.initiating_agent_path) - { + if let Some(initiating_agent_path) = start_options.initiating_agent_path { turn_context .turn_metadata_state .set_initiating_agent_path(initiating_agent_path); @@ -613,10 +616,6 @@ impl Session { } self.maybe_emit_model_warnings_for_turn(turn_context.as_ref()) .await; - // Task completion must still save this mail if pre-turn compaction fails. - self.input_queue - .extend_pending_input_for_turn_state(turn_state.as_ref(), input) - .await; // Context construction can outlive an interrupt or replacement. Do not let // start_task recreate a turn whose reservation is no longer ours. if self diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index b7d841c399..043661be2e 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -7,6 +7,8 @@ pub use shutdown::AgentTreeShutdown; pub use shutdown::AgentTreeShutdownFailure; pub use shutdown::AgentTreeShutdownFailureReason; pub use shutdown::AgentTreeShutdownReport; +pub use shutdown::LoadedSuspendedAgent; +pub use shutdown::SuspendedAgentTree; pub(crate) use shutdown::thread_store_error_kind; use crate::CodexAppsToolsCache; @@ -2083,6 +2085,22 @@ impl ThreadManagerState { }) } + pub(crate) async fn load_legacy_model_context( + &self, + params: LoadThreadHistoryParams, + ) -> CodexResult { + let history = self + .thread_store + .load_history(params) + .await + .map_err(thread_store_rollout_read_error)?; + Ok(StoredModelContext { + thread_id: history.thread_id, + revision: history.revision, + items: history.items, + }) + } + pub(crate) async fn load_latest_model_context( &self, params: LoadThreadHistoryParams, @@ -2121,11 +2139,31 @@ impl ThreadManagerState { } else { None }; + self.send_pinned_op( + &thread, + op, + parent_turn_id, + root_turn_id, + residency_guard, + message_delivery, + ) + .await + } + + pub(crate) async fn send_pinned_op( + &self, + thread: &CodexThread, + op: Op, + parent_turn_id: Option, + root_turn_id: Option, + residency_guard: Option>, + message_delivery: Option, + ) -> CodexResult { if let Some(ops_log) = &self.ops_log && let Ok(mut log) = ops_log.lock() && let Some(captured_op) = capture_test_op(&op) { - log.push((thread_id, captured_op)); + log.push((thread.session.thread_id, captured_op)); } thread .io @@ -2461,7 +2499,16 @@ impl ThreadManagerState { .config_layer_stack .with_mcp_ema_cancellation(cancellation); let runtime = request.agent_control.runtime().clone(); - let membership = runtime.admit_start()?; + let options = &request.options; + let source = options + .session_source + .as_ref() + .unwrap_or(&self.session_source); + request.parent_thread_id = request + .parent_thread_id + .or_else(|| options.initial_history.get_resumed_parent_thread_id()) + .or_else(|| source.parent_thread_id()); + let membership = runtime.admit_session(request.parent_thread_id)?; let owns_startup = request.startup.is_none(); let startup_state = Arc::clone( request @@ -2658,14 +2705,6 @@ impl ThreadManagerState { } (_, control) => control, }; - let parent_thread_id = parent_thread_id - .or_else(|| initial_history.get_resumed_parent_thread_id()) - .or(match &session_source { - SessionSource::SubAgent(SubAgentSource::ThreadSpawn { - parent_thread_id, .. - }) => Some(*parent_thread_id), - _ => None, - }); // Both resume entry points must restore identities before children can be loaded lazily. if let InitialHistory::Resumed(resumed) = &initial_history && initial_history.get_multi_agent_version() == Some(MultiAgentVersion::V2) diff --git a/codex-rs/core/src/thread_manager/shutdown.rs b/codex-rs/core/src/thread_manager/shutdown.rs index e622f41024..147583e650 100644 --- a/codex-rs/core/src/thread_manager/shutdown.rs +++ b/codex-rs/core/src/thread_manager/shutdown.rs @@ -1,11 +1,19 @@ //! Unloads locally controlled idle MAv2 children and exposes tracked agent-tree teardown. use super::ThreadManager; +use super::ThreadManagerState; +use crate::CodexThread; use crate::agent::control::AgentTreeShutdownState; +use crate::state::TaskKind; +use codex_history::ResponseItemEnvelope; +use codex_history::RolloutItem; use codex_protocol::ThreadId; use codex_protocol::error::CodexErr; use codex_protocol::error::Result as CodexResult; +use codex_thread_store::LiveThread; +use codex_thread_store::ResumeThreadParams; use codex_thread_store::ThreadStoreError; +use std::collections::HashMap; use std::fmt; use std::sync::Arc; use uuid::Uuid; @@ -128,7 +136,88 @@ impl AgentTreeShutdown { } } +/// Work to resume or start after a successful tree suspension. +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct SuspendedAgentTree { + /// Original ID of the root turn that was running when suspended. + pub root_turn_id: Option, + /// Last real root turn when a saved result should wake an idle root. The + /// replacement starts a fresh turn only if this predecessor is still current. + pub wake_root_after: Option, + pub child_turns: HashMap, +} + +/// Holds the restored child resident until the caller admits its turn. +pub struct LoadedSuspendedAgent { + pub thread: Arc, + pub(crate) _residency: Option>, +} + impl ThreadManager { + /// Stops this locally controlled tree and reports the original IDs of running regular + /// turns on the root and durable named V2 children, including turns without a recorded + /// model step. Returns None before installing the fence if the root lacks a durable + /// writer or has an unsupported active task. + /// Idle members and unsupported descendant tasks shut down normally. On success, + /// finished children have delivered their results to their parents' saved histories, + /// and every source writer has closed. Repeated calls on the same loaded root + /// return the same unfinished turns, even after caller cancellation. + pub async fn suspend_agent_tree( + &self, + root: &Arc, + ) -> CodexResult> { + self.check_local_agent_root(root).await?; + if root.session.live_thread().is_none() { + return Ok(None); + } + let id = root.startup_metadata().thread_id; + let shutdown = { + let active = root.session.active_turn.lock().await; + if active + .as_ref() + .and_then(|turn| turn.task.as_ref()) + .is_some_and(|task| task.kind != TaskKind::Regular) + { + return Ok(None); + } + root.session + .services + .local_agent_runtime + .request_suspension(id)? + }; + shutdown.suspended_tree().await.map(Some) + } + + /// Loads the named V2 child and its recorded ancestors under the replacement root. + /// This restores the original local identities without starting any model turns. + pub async fn load_suspended_agent( + &self, + root: &Arc, + child_id: ThreadId, + ) -> CodexResult { + self.check_local_agent_root(root).await?; + let runtime = &root.session.services.local_agent_runtime; + let _membership = runtime.admit_start()?; + runtime + .control(root.session.session_id()) + .load_suspended_child(root, child_id) + .await + } + + async fn check_local_agent_root(&self, root: &Arc) -> CodexResult<()> { + if self.state.agent_control_factory.is_some() + || root.session_source().is_non_root_agent() + || !Arc::ptr_eq( + root, + &self.get_thread(root.startup_metadata().thread_id).await?, + ) + { + return Err(CodexErr::InvalidRequest( + "agent tree recovery requires its loaded locally controlled root".into(), + )); + } + Ok(()) + } /// Unloads this exact idle MAv2 child using the manager's local controller backend, /// preserving mail and environments for reloading. Host controllers own their eviction /// and recovery; managers configured with a host controller factory are rejected. @@ -169,3 +258,53 @@ impl ThreadManager { }) } } + +impl ThreadManagerState { + /// An evicted parent has no session to consume queue-only mail. After the live + /// writers stop, reopen its writer long enough to save its unread results. + pub(crate) async fn save_evicted_parent_results( + &self, + parent: ThreadId, + config: &crate::config::Config, + history_mode: codex_protocol::protocol::ThreadHistoryMode, + results: &[crate::PendingAgentMessage], + ) -> codex_thread_store::ThreadStoreResult<()> { + let params = ResumeThreadParams { + thread_id: parent, + rollout_path: None, + history: None, + history_revision: None, + include_archived: false, + metadata: config.thread_persistence_metadata(), + }; + let (writer, _) = LiveThread::resume( + Arc::clone(&self.thread_store), + history_mode, + params, + /*resume_preparation*/ None, + ) + .await?; + let items = results + .iter() + .flat_map(|message| { + [ + RolloutItem::InterAgentCommunicationMetadata { + trigger_turn: message.trigger_turn, + }, + RolloutItem::ResponseItem(ResponseItemEnvelope::new( + crate::session::session::Session::inter_agent_history_item(message, config), + )), + ] + }) + .collect::>(); + let saved = writer.append_items(&items).await; + // Shutdown flushes pending items and always releases the acquired writer. + saved.and(writer.shutdown().await)?; + for result in results { + if let Some(write) = result.consumed(parent) { + write.confirm(); + } + } + Ok(()) + } +} diff --git a/codex-rs/core/tests/suite/abort_tasks.rs b/codex-rs/core/tests/suite/abort_tasks.rs index 0e039ee77d..8c3419edc8 100644 --- a/codex-rs/core/tests/suite/abort_tasks.rs +++ b/codex-rs/core/tests/suite/abort_tasks.rs @@ -1,26 +1,41 @@ use assert_matches::assert_matches; +use codex_core::CodexThread; +use codex_core::RecoverTurnRequest; use codex_core::SuspendTurnOutcome; use codex_core::TurnInputRequest; +use codex_extension_api::ExtensionFuture; +use codex_extension_api::TurnLifecycleContributor; +use codex_extension_api::TurnStartInput; +use codex_extension_api::TurnStopInput; use codex_features::Feature; use codex_history::RolloutItem; +use std::collections::HashMap; use std::sync::Arc; use std::time::Duration; +use codex_protocol::models::ResponseItem; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::Op; +use codex_protocol::protocol::ThreadHistoryMode; use codex_protocol::user_input::UserInput; +use core_test_support::ThreadIdle; +use core_test_support::responses::ev_assistant_message; use core_test_support::responses::ev_completed; use core_test_support::responses::ev_function_call; use core_test_support::responses::ev_function_call_with_namespace; use core_test_support::responses::ev_response_created; -use core_test_support::responses::mount_response_once; +use core_test_support::responses::mount_response_handler; use core_test_support::responses::mount_sse_once; use core_test_support::responses::mount_sse_sequence; use core_test_support::responses::sse; +use core_test_support::responses::sse_completed; +use core_test_support::responses::sse_failed; use core_test_support::responses::sse_response; use core_test_support::responses::start_mock_server; +use core_test_support::test_codex::TestCodex; use core_test_support::test_codex::test_codex; use core_test_support::wait_for_event; +use core_test_support::wait_for_event_match; use pretty_assertions::assert_eq; use regex_lite::Regex; use serde_json::json; @@ -71,48 +86,92 @@ async fn interrupt_long_running_tool_emits_turn_aborted() { codex.submit(Op::CleanBackgroundTerminals).await.unwrap(); } -/// Suspending a root preserves its unfinished turn and initiates child shutdown. The host -/// separately waits for tree cleanup, including any child writer failure, before recovery. -#[test_case(false; "recovers root after stopping child")] -#[test_case(true; "reports child writer failure separately")] -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn root_turn_suspension_preserves_unfinished_turn_history(fail_child_writer: bool) { - let server = start_mock_server().await; - mount_sse_once( - &server, +const CHILD_RESULT: &str = "the original worker found the issue"; + +#[derive(Clone, Copy, Default, PartialEq, Eq)] +enum ChildHandoff { + #[default] + QueuedBeforeSuspend, + FinishesDuringSuspend, + ErrorsDuringSuspend, + WakesIdleRoot, + WakeReservedBeforeSuspend, + ParentWriterClosed, + GrandchildFinishes, + ParentEvicted, +} + +impl ChildHandoff { + fn has_grandchild(self) -> bool { + matches!(self, Self::GrandchildFinishes | Self::ParentEvicted) + } + fn wakes_root(self) -> bool { + matches!(self, Self::WakesIdleRoot | Self::WakeReservedBeforeSuspend) + } +} + +#[derive(Default)] +struct ChildCompletionGate { + scenario: ChildHandoff, + child: std::sync::OnceLock, + entered: tokio::sync::Notify, + release: tokio::sync::Notify, +} + +impl ChildCompletionGate { + fn pause_if(&self, selected: bool) -> ExtensionFuture<'_, ()> { + Box::pin(async move { + if selected { + self.entered.notify_one(); + self.release.notified().await; + } + }) + } +} + +impl TurnLifecycleContributor for ChildCompletionGate { + fn on_turn_start<'a>(&'a self, input: TurnStartInput<'a>) -> ExtensionFuture<'a, ()> { + let is_root = input.thread_store.level_id() == input.session_store.level_id(); + self.pause_if( + is_root + && self.scenario == ChildHandoff::WakeReservedBeforeSuspend + && self.child.get().is_some(), + ) + } + + fn on_turn_stop<'a>(&'a self, input: TurnStopInput<'a>) -> ExtensionFuture<'a, ()> { + self.pause_if(self.child.get().map(String::as_str) == Some(input.thread_store.level_id())) + } +} + +async fn start_active_tree( + server: &wiremock::MockServer, + finishing: Option<&Arc>, + history_mode: ThreadHistoryMode, +) -> anyhow::Result<(TestCodex, String, Arc, String)> { + let spawn = |name, message| { sse(vec![ - ev_response_created("spawn_response"), + ev_response_created(name), ev_function_call_with_namespace( - "spawn_worker", + name, "collaboration", "spawn_agent", - &json!({ - "message": "keep working until shutdown", - "task_name": "worker", - "fork_turns": "none", - }) - .to_string(), + &json!({"message": message, "task_name": name, "fork_turns": "none"}).to_string(), ), - ev_completed("spawn_response"), - ]), - ) - .await; - // Waiting on the mocked response keeps the turn active on local and remote executors - // without requiring an OS-specific command or a working sandboxed child process. - for response_id in ["pending_child", "pending_root"] { - mount_response_once( - &server, - sse_response(sse(vec![ - ev_response_created(response_id), - ev_completed(response_id), - ])) - .set_delay(Duration::from_secs(60)), - ) - .await; - } - let test = test_codex() + ev_completed(name), + ]) + }; + mount_sse_once(server, spawn("worker", "keep working until shutdown")).await; + let wake_root = finishing.is_some_and(|gate| gate.scenario.wakes_root()); + let mut builder = test_codex() + .with_history_mode(history_mode) .with_model("gpt-5.6-sol") - .with_config(|config| { + .with_config(move |config| { + config.multi_agent_v2.max_concurrent_threads_per_session = 3; + config + .features + .set_enabled(Feature::WakeOnAgentMessages, wake_root) + .expect("configure root wake"); config .features .enable(Feature::Collab) @@ -121,10 +180,46 @@ async fn root_turn_suspension_preserves_unfinished_turn_history(fail_child_write .features .enable(Feature::MultiAgentV2) .expect("enable multi-agent v2 for the suspension test"); - }) - .build_with_auto_env(&server) - .await - .expect("start persistent root thread"); + }); + if let Some(gate) = finishing { + let mut extensions = codex_extension_api::ExtensionRegistryBuilder::new(); + extensions.turn_lifecycle_contributor(gate.clone()); + extensions.thread_lifecycle_contributor(Arc::new(ThreadIdle)); + builder = builder.with_extensions(Arc::new(extensions.build())); + } + let test = builder.build_with_auto_env(server).await?; + let root_id = test.session_configured.thread_id.to_string(); + let gate = finishing.cloned(); + mount_response_handler(server, move |request| { + let done = sse_response(sse_completed("done")); + let Some(gate) = &gate else { + return done.set_delay(Duration::from_secs(60)); + }; + let metadata = &request.body_json()["client_metadata"]; + let Some(parent) = metadata["x-codex-parent-thread-id"].as_str() else { + return done; + }; + if gate.scenario.has_grandchild() && parent == root_id { + if request.inputs_of_type("function_call_output").is_empty() { + return sse_response(spawn("leaf", "investigate the leaf")); + } + return if gate.scenario == ChildHandoff::GrandchildFinishes { + done.set_delay(Duration::from_secs(60)) + } else { + done + }; + } + let id = metadata["thread_id"].as_str().expect("child ID"); + gate.child.set(id.to_owned()).expect("one finishing child"); + if gate.scenario == ChildHandoff::ErrorsDuringSuspend { + return sse_response(sse_failed("child", "insufficient_quota", "quota exceeded")); + } + sse_response(sse(vec![ + ev_assistant_message("child", CHILD_RESULT), + ev_completed("child"), + ])) + }) + .await; let codex = Arc::clone(&test.codex); let mut created_threads = test.thread_manager.subscribe_thread_created(); let submitted = codex @@ -132,25 +227,37 @@ async fn root_turn_suspension_preserves_unfinished_turn_history(fail_child_write text: "preserve this exact unfinished turn".into(), text_elements: Vec::new(), }])) - .await - .expect("start root turn"); + .await?; let codex_core::TurnInputSubmission::Started { turn_id, .. } = submitted else { panic!("expected a started root turn"); }; - let child_id = tokio::time::timeout(Duration::from_secs(10), created_threads.recv()) - .await - .expect("child creation notification must arrive") - .expect("spawned child"); - let descendant = test - .thread_manager - .get_thread(child_id) - .await - .expect("loaded child"); - wait_for_event(&descendant, |event| { - matches!(event, EventMsg::TurnStarted(_)) + let child_id = tokio::time::timeout(Duration::from_secs(10), created_threads.recv()).await??; + let descendant = test.thread_manager.get_thread(child_id).await?; + let child_turn = wait_for_event_match(&descendant, |event| match event { + EventMsg::TurnStarted(turn) => Some(turn.turn_id.clone()), + _ => None, }) .await; + if let Some(gate) = finishing { + ThreadIdle::wait(&test.codex).await; + tokio::time::timeout(Duration::from_secs(10), gate.entered.notified()).await?; + } + Ok((test, turn_id, descendant, child_turn)) +} +/// Suspending a root preserves its unfinished turn and initiates child shutdown. The host +/// separately waits for tree cleanup, including any child writer failure, before recovery. +#[test_case(false; "recovers root after stopping child")] +#[test_case(true; "reports child writer failure separately")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn root_turn_suspension_preserves_unfinished_turn_history(fail_child_writer: bool) { + let server = start_mock_server().await; + let (test, turn_id, descendant, _) = + start_active_tree(&server, /*finishing*/ None, ThreadHistoryMode::Legacy) + .await + .expect("start the root and worker"); + let codex = Arc::clone(&test.codex); + let child_id = descendant.startup_metadata().thread_id; if fail_child_writer { test.thread_store .shutdown_thread(child_id) @@ -269,6 +376,236 @@ async fn root_turn_suspension_preserves_unfinished_turn_history(fail_child_write ); } +#[test_case(ChildHandoff::QueuedBeforeSuspend; "result already queued")] +#[test_case(ChildHandoff::FinishesDuringSuspend; "result arrives during suspend")] +#[test_case(ChildHandoff::ErrorsDuringSuspend; "error arrives during suspend")] +#[test_case(ChildHandoff::WakesIdleRoot; "result wakes idle root")] +#[test_case(ChildHandoff::WakeReservedBeforeSuspend; "fence wins before wake registration")] +#[test_case(ChildHandoff::ParentWriterClosed; "parent cannot save result")] +#[test_case(ChildHandoff::GrandchildFinishes; "unfinished parent")] +#[test_case(ChildHandoff::ParentEvicted; "evicted idle parent")] +#[tokio::test] +async fn graceful_suspend_preserves_a_finishing_child_result( + scenario: ChildHandoff, +) -> anyhow::Result<()> { + let server = start_mock_server().await; + let gate = Arc::new(ChildCompletionGate { + scenario, + ..Default::default() + }); + let mode = if scenario.wakes_root() || scenario.has_grandchild() { + ThreadHistoryMode::Legacy + } else { + ThreadHistoryMode::Paginated + }; + let (source, root_turn, worker, worker_turn) = + start_active_tree(&server, Some(&gate), mode).await?; + let worker_id = worker.startup_metadata().thread_id; + match scenario { + ChildHandoff::QueuedBeforeSuspend => { + gate.release.notify_one(); + ThreadIdle::wait(&worker).await; + } + ChildHandoff::WakeReservedBeforeSuspend => { + gate.release.notify_one(); + gate.entered.notified().await; + } + ChildHandoff::ParentEvicted => { + ThreadIdle::wait(&worker).await; + assert_eq!( + source + .thread_manager + .try_evict_v2_thread(Arc::clone(&worker)) + .await?, + codex_core::ThreadEvictionOutcome::Evicted + ); + } + ChildHandoff::ParentWriterClosed => { + source + .thread_store + .shutdown_thread(source.session_configured.thread_id) + .await?; + } + _ => {} + } + let drain = source.thread_manager.suspend_agent_tree(&source.codex); + tokio::pin!(drain); + if scenario != ChildHandoff::QueuedBeforeSuspend { + assert!(futures::poll!(drain.as_mut()).is_pending()); + gate.release.notify_one(); + } + let outcome = tokio::time::timeout(Duration::from_secs(10), drain).await?; + if scenario == ChildHandoff::ParentWriterClosed { + assert!(outcome.is_err(), "the result could not be saved"); + return Ok(()); + } + let tree = outcome?.expect("durable root"); + let child_turns = if scenario == ChildHandoff::GrandchildFinishes { + HashMap::from([(worker_id, worker_turn.clone())]) + } else { + HashMap::new() + }; + assert_eq!( + tree, + codex_core::SuspendedAgentTree { + root_turn_id: None, + wake_root_after: scenario.wakes_root().then(|| root_turn.clone()), + child_turns, + } + ); + let (parent, sender) = if scenario.has_grandchild() { + (&worker, "/root/worker/leaf") + } else { + (&source.codex, "/root/worker") + }; + let (history, _, _) = codex_rollout::RolloutRecorder::load_rollout_items( + &parent.rollout_path().expect("closed parent writer"), + ) + .await?; + let messages: Vec<_> = history + .iter() + .filter_map(|entry| match entry { + RolloutItem::ResponseItem(item) => match &item.item { + ResponseItem::AgentMessage { + author, content, .. + } if author == sender => { + codex_protocol::models::plaintext_agent_message_content(content) + } + _ => None, + }, + _ => None, + }) + .collect(); + let [result] = messages.as_slice() else { + panic!("expected one result, got {messages:?}") + }; + let payload = if scenario == ChildHandoff::ErrorsDuringSuspend { + "Agent errored:" + } else { + CHILD_RESULT + }; + assert!(result.contains(payload), "{result}"); + + if !scenario.wakes_root() && !scenario.has_grandchild() { + return Ok(()); + } + let server = start_mock_server().await; + let response = mount_sse_once(&server, sse(vec![ev_completed("received")])).await; + let replacement = test_codex() + .with_model("gpt-5.6-sol") + .with_config(|config| { + config.multi_agent_v2.max_concurrent_threads_per_session = 2; + config + .features + .enable(Feature::MultiAgentV2) + .expect("enable V2"); + }) + .restart_with_auto_env(&server, &source) + .await?; + let root = replacement.codex.clone(); + if scenario == ChildHandoff::ParentEvicted { + // The cold worker has one child slot. Load the grandchild through its real + // parent and release the parent's slot before admitting the grandchild. + let leaf_id = + codex_protocol::ThreadId::from_string(gate.child.get().expect("finishing leaf ID"))?; + let unrelated = replacement + .thread_manager + .start_thread(replacement.start_thread_options().await) + .await?; + assert!( + replacement + .thread_manager + .load_suspended_agent(&root, unrelated.thread_id) + .await + .is_err() + ); + let leaf = replacement + .thread_manager + .load_suspended_agent(&root, leaf_id) + .await?; + assert_eq!( + leaf.thread.session_source().parent_thread_id(), + Some(worker_id) + ); + assert!( + replacement + .thread_manager + .get_thread(worker_id) + .await + .is_err() + ); + return Ok(()); + } + let parent = if scenario == ChildHandoff::GrandchildFinishes { + let loaded = replacement + .thread_manager + .load_suspended_agent(&root, worker_id) + .await?; + assert_matches!(loaded.thread.recover_turn_if_idle(RecoverTurnRequest { turn_id: worker_turn.clone(), thread_settings: Default::default(), trace: None, prompt_cache_ttl: None, cyber_access_program: None }).await?, + codex_core::StartIfIdleSubmission::Started { turn_id, .. } if turn_id == worker_turn); + Arc::clone(&loaded.thread) + } else { + assert_eq!( + history + .iter() + .filter(|item| matches!(item, RolloutItem::EventMsg(EventMsg::TurnStarted(_)))) + .count(), + 1 + ); + assert_matches!(root.continue_turn_if_idle(TurnInputRequest::user_input(vec![]), root_turn.clone()).await?, + codex_core::TurnInputSubmission::Started { turn_id, .. } if turn_id != root_turn); + root + }; + wait_for_event(&parent, |event| matches!(event, EventMsg::TurnComplete(_))).await; + assert_eq!( + response + .single_request() + .inputs_of_type("agent_message") + .into_iter() + .filter(|item| item["author"] == sender) + .map(|item| item["content"].clone()) + .collect::>(), + [json!([{"type":"input_text", "text":result}])] + ); + Ok(()) +} + +#[tokio::test] +async fn suspended_tree_records_the_original_unfinished_turns() -> anyhow::Result<()> { + use codex_core::SuspendedAgentTree; + + let source_server = start_mock_server().await; + let (source, root_turn, child, child_turn) = start_active_tree( + &source_server, + /*finishing*/ None, + ThreadHistoryMode::Paginated, + ) + .await?; + // A single poll establishes the writer fence. Dropping the caller on the + // current-thread runtime leaves all the actual teardown work to the sessions. + { + let pending = source.thread_manager.suspend_agent_tree(&source.codex); + tokio::pin!(pending); + assert!(futures::poll!(pending.as_mut()).is_pending()); + } + let expected = SuspendedAgentTree { + root_turn_id: Some(root_turn.clone()), + wake_root_after: None, + child_turns: HashMap::from([(child.startup_metadata().thread_id, child_turn)]), + }; + for _ in 0..2 { + assert_eq!( + tokio::time::timeout( + Duration::from_secs(10), + source.thread_manager.suspend_agent_tree(&source.codex) + ) + .await??, + Some(expected.clone()) + ); + } + Ok(()) +} + /// After an interrupt we expect the next request to the model to include both /// the original tool call and an `"aborted"` `function_call_output`. This test /// exercises the follow-up flow: it sends another user turn, inspects the mock diff --git a/codex-rs/core/tests/suite/host_turn_admission.rs b/codex-rs/core/tests/suite/host_turn_admission.rs index 99e190ee77..dc5060fff4 100644 --- a/codex-rs/core/tests/suite/host_turn_admission.rs +++ b/codex-rs/core/tests/suite/host_turn_admission.rs @@ -306,6 +306,25 @@ async fn suspending_root_during_admission_preserves_the_recoverable_prompt() -> Ok(()) } +#[tokio::test] +async fn tree_suspension_leaves_an_ephemeral_turn_running() -> Result<()> { + let server = responses::start_mock_server().await; + let reply = responses::mount_sse_once(&server, responses::sse_completed("done")).await; + let builder = test_codex().with_config(|config| config.ephemeral = true); + let (test, mut attempts) = + root_with_admission(&server, builder, /*session_reader*/ None).await?; + let prompt = "Keep working after the suspension attempt."; + let (_, grant) = start(&test, &mut attempts, prompt).await?; + assert_eq!( + test.thread_manager.suspend_agent_tree(&test.codex).await?, + None + ); + assert!(grant.send(Ok(AgentExecutionGuard::new(()))).is_ok()); + wait_for_successful_turn(&test.codex).await; + assert!(reply.single_request().body_contains_text(prompt)); + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn suspending_root_during_roster_read_preserves_the_recoverable_prompt() -> Result<()> { skip_if_no_network!(Ok(()));