Overlap subagent spawn persistence and clean up cancelled children (#47348)

## Why

Forked subagent startup serializes history persistence and spawn-edge writes. Cancellation during startup can also leave a provisional child behind.

## What changed

- Defer copied-history durability so it can run alongside spawn-edge persistence, while still awaiting both before submitting the child's initial input.
- Commit agent and residency reservations and notify clients only after initial input is accepted.
- Guard provisional children so cancellation or startup failure removes and shuts them down, discards their stored state, and closes their spawn edge after any pending edge write finishes.

## Testing

Add coverage for cancellation at the child-history durability barrier, checking child termination, removal, and edge closure. Extend fork-context tests in legacy and paginated history modes to verify inherited context can be read from storage after spawn acknowledgment without flushing the live child first.

GitOrigin-RevId: e02dda57efc32b7237bb7af0ce5e5844922890c3

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;

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

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<ThreadManagerState>,
+    child: Option<ThreadId>,
+    edge_write: Option<JoinHandle<()>>,
+}
+
+impl PendingSpawn {
+    pub(super) fn new(state: Arc<ThreadManagerState>, 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}");
+            }
+        }));
+    }
+}

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<HistoryPosition>,
         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 {

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(

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;

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<PersistContext>,
+    checkpoint_context: Option<PersistContext>,
     checkpoints: mpsc::UnboundedSender<PendingCheckpoint>,
 }
 
@@ -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(())
+}

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(_)) {

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

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,
         }
     }
 }