openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-11-13 21:30:47 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add graceful agent-tree suspension and recovery (#61198)

Browse Source 
## Why

Suspending a locally controlled agent tree must preserve unfinished turns and child results that arrive during teardown so they remain available after recovery.

## What changed

- Add `ThreadManager::suspend_agent_tree` to stop durable trees, close their writers, and return unfinished root and named V2 child turn IDs plus any pending root wake.
- Persist terminal child results before suspension succeeds, including results queued for evicted parents. Report failed or unconfirmed writes.
- Add `load_suspended_agent` to restore named V2 descendants through their recorded ancestors without starting model turns.
- Allow empty `continue_turn_if_idle` input to wake on saved results, guarded by the preceding turn ID. Keep mail queued until task registration so suspension can preserve it.

## Testing

Add coverage for completion and wake races, child errors, parent writer failures, evicted parents, descendant restoration, repeated suspension after caller cancellation, and ephemeral roots remaining active.

GitOrigin-RevId: 3041791f51bff2555210f6b2d2443f9f05f68330
This commit is contained in:
Bryan Ashley
2026-11-12 22:43:14 +00:00
committed by copyberry
parent cdfcc18dba
commit fae87361f4
30 changed files with 1480 additions and 320 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core-api/src/lib.rs
+2
View File
@@ -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;
codex-rs/core/src/agent/control.rs
+44 -15
View File
@@ -264,21 +264,37 @@ impl LocalAgentControl {
context: AgentCommunicationContext,
start_options: TurnStartOptions,
) -> CodexResult<String> {
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,
codex-rs/core/src/agent/control/api.rs
+1 -1
View File
@@ -179,7 +179,7 @@ impl AgentControl for LocalAgentControl {
}
fn take_mailbox(&self, agent: ThreadId) -> Vec<crate::PendingAgentMessage> {
self.runtime.mailboxes.take(agent)
self.runtime.mailboxes.take(agent, |_| true)
}
fn watch_mailbox(&self, agent: ThreadId) -> tokio::sync::watch::Receiver<bool> {
codex-rs/core/src/agent/control/completion.rs
+5 -7
View File
@@ -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;
codex-rs/core/src/agent/control/mailbox.rs
+18 -8
View File
@@ -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<PendingAgentMessage> {
pub(super) fn take(
&self,
agent: ThreadId,
accept: impl Fn(&PendingAgentMessage) -> bool,
) -> Vec<PendingAgentMessage> {
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 {
codex-rs/core/src/agent/control/residency.rs
+11 -4
View File
@@ -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))
}
}
codex-rs/core/src/agent/control/residency_tests.rs
+1 -1
View File
@@ -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);
codex-rs/core/src/agent/control/resume.rs
+87 -1
View File
@@ -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<CodexThread>,
child_id: ThreadId,
) -> CodexResult<LoadedSuspendedAgent> {
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,
codex-rs/core/src/agent/control/runtime.rs
+203 -3
View File
@@ -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<AgentTreeShutdownReport>,
disposition: OnceLock<TreeDisposition>,
suspended: Mutex<SuspendedAgentTree>,
children: Mutex<HashMap<ThreadId, TaskTracker>>,
}
#[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<SuspendedAgentTree> {
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<AgentTreeShutdownState>,
_member: TaskTrackerToken,
_child: Option<TaskTrackerToken>,
}
impl AgentTreeMembership {
@@ -202,6 +237,7 @@ pub(crate) struct LocalAgentRuntime {
Arc<OnceLock<Arc<dyn ThreadInstructionsProvider>>>,
pub(super) registry: Arc<AgentRegistry>,
pub(super) residency: Arc<V2Residency>,
pub(super) unconfirmed_results: Arc<AtomicUsize>,
pub(super) mailboxes: Arc<super::mailbox::Mailboxes>,
/// 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<ThreadId>,
) -> CodexResult<AgentTreeMembership> {
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<AgentTreeShutdownState> {
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<Arc<AgentTreeShutdownState>> {
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<ThreadId> {
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) {
codex-rs/core/src/agent/control/spawn.rs
+48 -21
View File
@@ -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<ThreadId>,
) {
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,
codex-rs/core/src/agent/control_tests.rs
+27
View File
@@ -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
codex-rs/core/src/agent/history.rs
+7 -28
View File
@@ -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<Option<codex_thread_store::StoredModelContext>> {
) -> CodexResult<codex_thread_store::StoredModelContext> {
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<String> =
codex-rs/core/src/agent/registry.rs
+1
View File
@@ -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<Config>,
pub(crate) has_independent_service_tier: bool,
pub(crate) selections: Vec<TurnEnvironmentSelection>,
codex-rs/core/src/agent_communication.rs
+55 -11
View File
@@ -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<ResultWrite> {
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<AtomicUsize>);
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<dyn MessageLifecycleContributor>]>,
result_write: Option<ResultWrite>,
}
/// 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<dyn MessageLifecycleContributor>],
) -> Self {
unconfirmed_results: Option<&Arc<AtomicUsize>>,
) -> Option<Self> {
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<ResultWrite> {
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
}
}
codex-rs/core/src/codex_thread.rs
+3 -1
View File
@@ -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<WithTurnExtensionData<TurnInputRequest>>,
codex-rs/core/src/config/mod.rs
+15
View File
@@ -1680,6 +1680,21 @@ async fn config_toml_from_layers(layers: &ConfigLayerStack) -> std::io::Result<C
}
impl Config {
pub(crate) fn thread_persistence_metadata(
&self,
) -> 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() {
codex-rs/core/src/lib.rs
+2
View File
@@ -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;
codex-rs/core/src/session/handlers.rs
+68 -1
View File
@@ -316,6 +316,9 @@ pub(super) async fn shutdown_session_runtime(sess: &Arc<Session>) {
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<Session>) {
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 {
codex-rs/core/src/session/input_queue.rs
+49 -10
View File
@@ -269,15 +269,25 @@ impl InputQueue {
pending.drain(..).collect()
}
pub(crate) async fn drain_mailbox_input_items(&self) -> (Vec<TurnInput>, TurnStartOptions) {
let pending_mails = self.drain_mailbox().await;
pub(crate) async fn pending_wake_options(&self) -> Option<TurnStartOptions> {
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<PendingMailboxCommunication>,
) -> Option<TurnStartOptions> {
// 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<TurnInput>, 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(
codex-rs/core/src/session/mod.rs
+48 -20
View File
@@ -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<crate::PendingAgentMessage>,
) {
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<crate::PendingAgentMessage>,
) {
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
codex-rs/core/src/session/session.rs
+21 -27
View File
@@ -60,9 +60,10 @@ use tokio_util::task::TaskTracker;
type McpToolApprovalMetadataMap =
HashMap<(codex_mcp::McpServerKey, String), std::sync::Weak<McpToolApprovalContext>>;
/// 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<dyn AgentTurnOutcomeReporter>,
pub(crate) reporter: Option<Arc<dyn AgentTurnOutcomeReporter>>,
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::<Arc<dyn AgentTurnOutcomeReporter>>()
.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) =
codex-rs/core/src/session/tests.rs
+7 -8
View File
@@ -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,
};
codex-rs/core/src/session/turn_input.rs
+52 -42
View File
@@ -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<tokio::sync::Mutex<TurnState>>, 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;
}
codex-rs/core/src/session/turn_input_tests.rs
+23
View File
@@ -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;
codex-rs/core/src/session/turn_suspension.rs
+38
View File
@@ -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<Session>) -> 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
}
codex-rs/core/src/tasks/mod.rs
+46 -47
View File
@@ -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<TurnContext>,
input: Vec<TurnInput>,
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<T: SessionTask>(
self: &Arc<Self>,
turn_context: Arc<TurnContext>,
input: Vec<TurnInput>,
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
codex-rs/core/src/thread_manager.rs
+49 -10
View File
@@ -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<StoredModelContext> {
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<String>,
root_turn_id: Option<String>,
residency_guard: Option<tokio::sync::OwnedRwLockReadGuard<()>>,
message_delivery: Option<crate::agent_communication::MessageDelivery>,
) -> CodexResult<String> {
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)
codex-rs/core/src/thread_manager/shutdown.rs
+139
View File
@@ -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<String>,
/// 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<String>,
pub child_turns: HashMap<ThreadId, String>,
}
/// Holds the restored child resident until the caller admits its turn.
pub struct LoadedSuspendedAgent {
pub thread: Arc<CodexThread>,
pub(crate) _residency: Option<tokio::sync::OwnedRwLockReadGuard<()>>,
}
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<CodexThread>,
) -> CodexResult<Option<SuspendedAgentTree>> {
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<CodexThread>,
child_id: ThreadId,
) -> CodexResult<LoadedSuspendedAgent> {
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<CodexThread>) -> 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::<Vec<_>>();
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(())
}
}
codex-rs/core/tests/suite/abort_tasks.rs
+391 -54
View File
@@ -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<String>,
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<ChildCompletionGate>>,
history_mode: ThreadHistoryMode,
) -> anyhow::Result<(TestCodex, String, Arc<CodexThread>, 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::<Vec<_>>(),
[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
codex-rs/core/tests/suite/host_turn_admission.rs
+19
View File
@@ -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(()));