diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index c9a4516494..fb048a5da5 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/codex-rs/core/src/agent/control.rs @@ -77,6 +77,7 @@ mod runtime; mod sender_context; mod service_tier; mod spawn; +mod spawn_guard; mod target; mod user_authorization; mod watch; diff --git a/codex-rs/core/src/agent/control/spawn.rs b/codex-rs/core/src/agent/control/spawn.rs index 760f96c539..e7ed85cb41 100644 --- a/codex-rs/core/src/agent/control/spawn.rs +++ b/codex-rs/core/src/agent/control/spawn.rs @@ -1,4 +1,5 @@ use super::residency::is_v2_resident_session_source; +use super::spawn_guard::PendingSpawn; use super::*; use crate::agent::child_config::build_agent_resume_config; use crate::agent::role::apply_role_to_config; @@ -27,6 +28,7 @@ use codex_history::ResponseItemEnvelope; use codex_prompts::ResolvedModelMessages; use codex_protocol::intersect_effective_permission_profiles; use codex_protocol::protocol::EnvironmentConfigState; +use codex_thread_store::PersistContext; use codex_utils_path_uri::PathUri; const AGENT_NAMES: &str = include_str!("../../../assets/agent/agent_names.txt"); @@ -732,10 +734,7 @@ impl LocalAgentControl { (None, _, _) => Box::pin(state.spawn_new_thread(config.clone(), self.clone())).await?, }; agent_metadata.agent_id = Some(new_thread.thread_id); - reservation.commit(agent_metadata.clone()); - if let Some(residency_slot) = residency_slot { - residency_slot.commit(new_thread.thread_id); - } + let mut pending_spawn = PendingSpawn::new(Arc::clone(&state), new_thread.thread_id); if let Some(SessionSource::SubAgent( subagent_source @ SubAgentSource::ThreadSpawn { @@ -770,17 +769,30 @@ impl LocalAgentControl { ); } - // Notify a new thread has been created. This notification will be processed by clients - // to subscribe or drain this newly created thread. - // TODO(jif) add helper for drain - state.notify_thread_created(new_thread.thread_id); - - self.persist_thread_spawn_edge_for_source( - new_thread.thread.as_ref(), - new_thread.thread_id, - notification_source.as_ref(), - ) - .await; + let control = self.clone(); + let child = Arc::clone(&new_thread.thread); + let child_thread_id = new_thread.thread_id; + let source = notification_source.clone(); + pending_spawn.set_edge_write(tokio::spawn(async move { + control + .persist_thread_spawn_edge_for_source( + child.as_ref(), + child_thread_id, + source.as_ref(), + ) + .await; + })); + if options.fork_mode.is_some() { + tokio::join!( + new_thread + .thread + .session + .ensure_rollout_materialized(PersistContext::Standard), + pending_spawn.wait_for_edge(), + ); + } else { + pending_spawn.wait_for_edge().await; + } let start_options = TurnStartOptions { parent_turn_id: options.parent_turn_id, @@ -805,6 +817,16 @@ impl LocalAgentControl { .await?; } } + reservation.commit(agent_metadata.clone()); + if let Some(residency_slot) = residency_slot { + residency_slot.commit(new_thread.thread_id); + } + pending_spawn.disarm(); + + // Notify a new thread has been created. This notification will be processed by clients + // to subscribe or drain this newly created thread. + // TODO(jif) add helper for drain + state.notify_thread_created(new_thread.thread_id); if multi_agent_version != MultiAgentVersion::V2 { let child_reference = agent_metadata .agent_path diff --git a/codex-rs/core/src/agent/control/spawn_guard.rs b/codex-rs/core/src/agent/control/spawn_guard.rs new file mode 100644 index 0000000000..7df61e64ac --- /dev/null +++ b/codex-rs/core/src/agent/control/spawn_guard.rs @@ -0,0 +1,75 @@ +//! Owns a spawned child until its initial input is accepted. + +use crate::thread_manager::ThreadManagerState; +use codex_agent_graph_store::ThreadSpawnEdgeStatus; +use codex_protocol::ThreadId; +use std::sync::Arc; +use tokio::task::JoinHandle; +use tracing::warn; + +pub(super) struct PendingSpawn { + state: Arc, + child: Option, + edge_write: Option>, +} + +impl PendingSpawn { + pub(super) fn new(state: Arc, child: ThreadId) -> Self { + Self { + state, + child: Some(child), + edge_write: None, + } + } + + pub(super) fn set_edge_write(&mut self, edge_write: JoinHandle<()>) { + self.edge_write = Some(edge_write); + } + + pub(super) async fn wait_for_edge(&mut self) { + if let Some(edge_write) = self.edge_write.as_mut() { + assert!( + edge_write.await.is_ok(), + "spawn edge write task should complete" + ); + } + self.edge_write = None; + } + + pub(super) fn disarm(mut self) { + self.child = None; + } +} + +impl Drop for PendingSpawn { + fn drop(&mut self) { + let Some(child) = self.child.take() else { + return; + }; + let state = Arc::clone(&self.state); + let edge_write = self.edge_write.take(); + drop(tokio::spawn(async move { + if let Some(thread) = state.remove_thread(&child).await { + if let Err(error) = thread.shutdown_and_wait().await { + warn!("failed to stop cancelled child spawn: {error}"); + } + if let Some(live_thread) = thread.session.live_thread() + && let Err(error) = live_thread.discard().await + { + warn!("failed to discard cancelled child spawn: {error}"); + } + } + // A pending Open write must finish before cleanup writes Closed. + if let Some(edge_write) = edge_write { + let _ = edge_write.await; + } + if let Some(store) = state.agent_graph_store() + && let Err(error) = store + .set_thread_spawn_edge_status(child, ThreadSpawnEdgeStatus::Closed) + .await + { + warn!("failed to close cancelled child spawn edge: {error}"); + } + })); + } +} diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index c8c1cef2fd..9c48261b7f 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -416,6 +416,8 @@ pub(crate) enum GitEnrichmentPolicy { /// Controls which fork history belongs in the newly created thread's own rollout. pub(crate) enum ForkPersistence { Copied, + /// The spawn caller owns the durability barrier before acknowledging the child. + CopiedDeferred, Referenced { history_base: Option, inherited_item_count: usize, @@ -1641,13 +1643,15 @@ impl Session { rollout_items.drain(..*inherited_item_count); rollout_items.insert(0, thread_settings_applied); } - ForkPersistence::Copied if is_paginated_subagent => { + ForkPersistence::Copied | ForkPersistence::CopiedDeferred + if is_paginated_subagent => + { // Paginated subagents already persist inherited context when their live // thread is created. rollout_items.clear(); rollout_items.push(thread_settings_applied); } - ForkPersistence::Copied => { + ForkPersistence::Copied | ForkPersistence::CopiedDeferred => { // Keep the copied prefix and effective child settings in one append so a // cold resume cannot observe inherited settings as the latest value. rollout_items.push(thread_settings_applied); @@ -1655,9 +1659,14 @@ impl Session { } self.persist_rollout_items(&rollout_items).await; - // Forked threads should remain file-backed immediately after startup. - self.ensure_rollout_materialized(PersistContext::Standard) - .await; + // Agent spawn owns its final durability barrier so it can overlap the edge write. + let persist_context = match self.fork_persistence { + ForkPersistence::CopiedDeferred => PersistContext::SubagentSpawn, + ForkPersistence::Copied | ForkPersistence::Referenced { .. } => { + PersistContext::Standard + } + }; + self.ensure_rollout_materialized(persist_context).await; // Flush after seeding history and any persisted rollout copy. if !is_subagent { diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index cc52c2bf93..21c68aad6f 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -808,7 +808,7 @@ impl Session { ForkPersistence::Referenced { history_base, .. } => { history_base.map(|position| position.end_ordinal_exclusive) } - ForkPersistence::Copied => match &initial_history { + ForkPersistence::Copied | ForkPersistence::CopiedDeferred => match &initial_history { InitialHistory::Resumed(resumed) => { // Both local and CCA thread stores place the resumed thread's // canonical SessionMeta first. Never inspect inherited metadata: @@ -1018,7 +1018,7 @@ impl Session { multi_agent_version: initial_multi_agent_version, history_mode: session_configuration.history_mode, history_base: match &fork_persistence { - ForkPersistence::Copied => None, + ForkPersistence::Copied | ForkPersistence::CopiedDeferred => None, ForkPersistence::Referenced { history_base, .. } => *history_base, }, subagent_history_start_ordinal: None, @@ -1037,7 +1037,10 @@ impl Session { }, }; if is_paginated_subagent - && matches!(&fork_persistence, ForkPersistence::Copied) + && matches!( + &fork_persistence, + ForkPersistence::Copied | ForkPersistence::CopiedDeferred + ) && let InitialHistory::Forked(items) = &initial_history { LiveThread::create_with_inherited_model_context( diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 6725939d2f..521388dbbe 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -2009,6 +2009,7 @@ impl ThreadManagerState { }; let mut request = ThreadSpawnRequest::new(options, Arc::clone(&self.auth_manager), agent_control); + request.fork_persistence = ForkPersistence::CopiedDeferred; request.parent_thread_id = parent_thread_id; request.forked_from_thread_id = forked_from_thread_id; request.inherited_environments = inherited_environments; diff --git a/codex-rs/core/tests/suite/pending_input_persistence.rs b/codex-rs/core/tests/suite/pending_input_persistence.rs index 84f71d94a4..09261abfa7 100644 --- a/codex-rs/core/tests/suite/pending_input_persistence.rs +++ b/codex-rs/core/tests/suite/pending_input_persistence.rs @@ -9,16 +9,19 @@ use std::time::Duration; use codex_core::TurnInput; use codex_core::TurnInputRequest; use codex_core::TurnInputSubmission; +use codex_features::Feature; use codex_history::InitialHistory; use codex_history::ResumedHistory; use codex_protocol::ThreadId; use codex_protocol::mcp::ClientMcpExtensions; 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 codex_rollout::RolloutItem; use codex_rollout::RolloutRecorder; +use codex_state::DirectionalThreadSpawnEdgeStatus; use codex_thread_store::AppendThreadItemsParams; use codex_thread_store::ArchiveThreadParams; use codex_thread_store::CreateThreadParams; @@ -63,6 +66,7 @@ enum InputKind { #[derive(Debug)] struct PendingCheckpoint { + thread_id: ThreadId, context: PersistContext, complete: oneshot::Sender<()>, } @@ -72,6 +76,8 @@ struct GatedCheckpointStore { inner: InMemoryThreadStore, policy: CheckpointPolicy, armed: AtomicBool, + arm_on: Option, + checkpoint_context: Option, checkpoints: mpsc::UnboundedSender, } @@ -129,15 +135,25 @@ impl ThreadStore for GatedCheckpointStore { context: PersistContext, ) -> ThreadStoreFuture<'_, ()> { Box::pin(async move { + if self.arm_on == Some(context) { + self.armed.store(true, Ordering::SeqCst); + } if self.policy == CheckpointPolicy::Preparation && context == PersistContext::ThreadPreparation { return Ok(()); } - if self.armed.swap(false, Ordering::SeqCst) { + let should_checkpoint = self + .checkpoint_context + .is_none_or(|checkpoint_context| checkpoint_context == context); + if should_checkpoint && self.armed.swap(false, Ordering::SeqCst) { let (complete, completed) = oneshot::channel(); self.checkpoints - .send(PendingCheckpoint { context, complete }) + .send(PendingCheckpoint { + thread_id, + context, + complete, + }) .expect("checkpoint receiver should stay alive"); if self.policy == CheckpointPolicy::Synchronous || !context.allows_background_persistence() @@ -192,6 +208,8 @@ async fn steered_input_checkpoint_controls_next_request( inner: InMemoryThreadStore::default(), policy, armed: AtomicBool::new(false), + arm_on: None, + checkpoint_context: None, checkpoints, }); let base_url = format!("{}/v1", server.uri()); @@ -297,6 +315,8 @@ async fn preparation_and_first_sampling_do_not_wait_for_durable_metadata() -> an inner: InMemoryThreadStore::default(), policy: CheckpointPolicy::Preparation, armed: AtomicBool::new(false), + arm_on: None, + checkpoint_context: None, checkpoints, }); let base_url = format!("{}/v1", server.uri()); @@ -469,3 +489,99 @@ async fn local_preparation_is_durable_before_first_input_and_survives_restart( resumed.thread.shutdown_and_wait().await?; Ok(()) } + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn cancelled_spawn_discards_provisional_child() -> anyhow::Result<()> { + let server = responses::start_mock_server().await; + let (checkpoints, mut checkpoint_requests) = mpsc::unbounded_channel(); + let store = Arc::new(GatedCheckpointStore { + inner: InMemoryThreadStore::default(), + policy: CheckpointPolicy::Synchronous, + armed: AtomicBool::new(false), + arm_on: Some(PersistContext::SubagentSpawn), + checkpoint_context: Some(PersistContext::Standard), + checkpoints, + }); + let test = test_codex() + .with_thread_store(store) + .with_history_mode(ThreadHistoryMode::Legacy) + .with_config(|config| { + config.agent_max_depth = 1; + config + .features + .enable(Feature::Collab) + .expect("enable collaboration"); + }) + .build_with_auto_env(&server) + .await?; + responses::mount_sse_once( + &server, + responses::sse(vec![ + responses::ev_function_call_with_namespace( + "spawn-checkpoint", + "multi_agent_v1", + "spawn_agent", + &json!({"message":"child startup task", "task_name":"worker", "fork_context":true}) + .to_string(), + ), + responses::ev_completed("parent-spawn"), + ]), + ) + .await; + + test.codex + .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text { + text: "parent context to inherit".to_string(), + text_elements: Vec::new(), + }])) + .await?; + let checkpoint = timeout(Duration::from_secs(10), checkpoint_requests.recv()) + .await? + .expect("child history should reach its durability barrier"); + let child = test.thread_manager.get_thread(checkpoint.thread_id).await?; + let state_db = codex_core::init_state_db(&test.config) + .await + .expect("state db should be enabled"); + + test.codex.submit(Op::Interrupt).await?; + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnAborted(_)) + }) + .await; + timeout(Duration::from_secs(10), child.wait_until_terminated()).await?; + timeout(Duration::from_secs(10), async { + while test + .thread_manager + .get_thread(checkpoint.thread_id) + .await + .is_ok() + { + tokio::task::yield_now().await; + } + }) + .await?; + timeout(Duration::from_secs(10), async { + loop { + let open_children = state_db + .list_thread_spawn_children_with_status( + test.session_configured.thread_id, + DirectionalThreadSpawnEdgeStatus::Open, + ) + .await?; + let closed_children = state_db + .list_thread_spawn_children_with_status( + test.session_configured.thread_id, + DirectionalThreadSpawnEdgeStatus::Closed, + ) + .await?; + if open_children.is_empty() && closed_children == vec![checkpoint.thread_id] { + return Ok::<(), anyhow::Error>(()); + } + tokio::task::yield_now().await; + } + }) + .await??; + + test.codex.shutdown_and_wait().await?; + Ok(()) +} diff --git a/codex-rs/core/tests/suite/subagent_notifications.rs b/codex-rs/core/tests/suite/subagent_notifications.rs index cdb28f4a19..ccd44bb92d 100644 --- a/codex-rs/core/tests/suite/subagent_notifications.rs +++ b/codex-rs/core/tests/suite/subagent_notifications.rs @@ -28,6 +28,7 @@ use codex_protocol::protocol::SubAgentSource; use codex_protocol::protocol::ThreadHistoryMode; use codex_protocol::protocol::ThreadSettingsOverrides; use codex_protocol::user_input::UserInput; +use codex_thread_store::LoadThreadHistoryParams; use core_test_support::hooks::trust_discovered_hooks; use core_test_support::responses::ResponsesRequest; use core_test_support::responses::assert_parent_turn; @@ -1105,6 +1106,15 @@ async fn spawned_child_receives_forked_parent_context( .as_str() .expect("legacy child thread id"), )?; + // Read the acknowledged fork from storage without flushing the live child first. + let reopened_store = codex_core::thread_store_from_config(&test.config, /*state_db*/ None); + let persisted = reopened_store + .load_latest_model_context(LoadThreadHistoryParams { + thread_id: child_thread_id, + include_archived: false, + }) + .await?; + assert!(serde_json::to_string(&persisted.items)?.contains(TURN_0_FORK_PROMPT)); let child_thread = test.thread_manager.get_thread(child_thread_id).await?; tokio::time::timeout(Duration::from_secs(2), async { while !matches!(child_thread.agent_status().await, AgentStatus::Completed(_)) { diff --git a/codex-rs/thread-store/src/local/mod.rs b/codex-rs/thread-store/src/local/mod.rs index 66af0429ed..330ec0b765 100644 --- a/codex-rs/thread-store/src/local/mod.rs +++ b/codex-rs/thread-store/src/local/mod.rs @@ -529,6 +529,9 @@ impl ThreadStore for LocalThreadStore { thread_id: ThreadId, context: PersistContext, ) -> ThreadStoreFuture<'_, ()> { + if context == PersistContext::SubagentSpawn { + return Box::pin(async { Ok(()) }); + } Box::pin(async move { if context == PersistContext::ThreadPreparation { live_writer::flush_thread(self, thread_id).await diff --git a/codex-rs/thread-store/src/store.rs b/codex-rs/thread-store/src/store.rs index 38d52ee72c..d67c3b766a 100644 --- a/codex-rs/thread-store/src/store.rs +++ b/codex-rs/thread-store/src/store.rs @@ -68,6 +68,9 @@ pub enum PersistContext { ThreadPreparation, /// Standard persistence makes the thread and all queued items durable and readable. Standard, + /// Copied subagent history may be enqueued so its durability fence can overlap relationship + /// persistence. The caller must await standard persistence before acknowledging the child. + SubagentSpawn, /// A turn is about to begin sampling after its input has been recorded. TurnStart, /// Accepted user input is being recorded before an active turn's next sampling request. @@ -81,7 +84,7 @@ impl PersistContext { pub fn allows_background_persistence(self) -> bool { match self { Self::ThreadPreparation | Self::Standard => false, - Self::TurnStart | Self::SteeredUserInput => true, + Self::SubagentSpawn | Self::TurnStart | Self::SteeredUserInput => true, } } }