From 4e7ef7d19d59c51ac04e62d062664bfd012d76da Mon Sep 17 00:00:00 2001 From: jif Date: Tue, 13 Oct 2026 12:33:37 +0000 Subject: [PATCH] Keep spawned children provisional until the spawn result returns (#53654) ## Why Cancellation after initial input admission could leave a child registered even though the spawn caller had not received its result. ## What changed Resolve the child's status and configuration before committing reservations and releasing provisional ownership. This keeps cancellation cleanup active through the final awaits and lets the spawn result return without yielding after commitment. Cancellation after the result returns does not retract the accepted child. ## Testing Add a regression test that cancels spawn during the final status lookup after input admission and verifies that the child is unregistered, terminated, and unloaded, and that shutdown completes. Extend the provisional-child cancellation test to cover both native V1 and V2 callers. GitOrigin-RevId: 028b873ad4c1fac2c4b5d1dc823b4561ff19e6bf --- codex-rs/core/src/agent/api.rs | 3 + codex-rs/core/src/agent/control/spawn.rs | 13 ++-- .../core/src/agent/control/spawn_guard.rs | 2 +- codex-rs/core/src/agent/control_tests.rs | 63 +++++++++++++++++++ .../tests/suite/pending_input_persistence.rs | 33 ++++++++-- 5 files changed, 103 insertions(+), 11 deletions(-) diff --git a/codex-rs/core/src/agent/api.rs b/codex-rs/core/src/agent/api.rs index 1a498228b6..622c451343 100644 --- a/codex-rs/core/src/agent/api.rs +++ b/codex-rs/core/src/agent/api.rs @@ -53,6 +53,9 @@ pub trait AgentControl: Send + Sync { ) -> BoxFuture<'a, Result>; /// Start a child and accept its initial input, returning its effective settings. + /// Retain provisional-child ownership until this future returns its acceptance receipt. + /// Once returned, caller cancellation does not retract the accepted child. Subsequent + /// tool-result or transport publication belongs to the caller, outside this boundary. /// V2 snapshots resolve effort against captured startup metadata without changing the /// child's configured effort (for example, Ultra still enables proactive behavior). fn spawn( diff --git a/codex-rs/core/src/agent/control/spawn.rs b/codex-rs/core/src/agent/control/spawn.rs index cca217e8d0..70cf8b7ef0 100644 --- a/codex-rs/core/src/agent/control/spawn.rs +++ b/codex-rs/core/src/agent/control/spawn.rs @@ -954,6 +954,13 @@ impl LocalAgentControl { } } let input_admission = input_admission_started_at.elapsed(); + // Keep provisional ownership through the final awaits. Commit and return to the + // backend caller in one poll; later tool-result publication is outside this boundary. + let status = self.get_status(new_thread.thread_id).await; + let config = match spawn_config { + Some(config) => config, + None => new_thread.thread.config_snapshot().await, + }; reservation.commit(agent_metadata.clone()); if let Some(residency_slot) = residency_slot { residency_slot.commit(new_thread.thread_id); @@ -981,11 +988,7 @@ impl LocalAgentControl { let agent = LiveAgent { thread_id: new_thread.thread_id, metadata: agent_metadata, - status: self.get_status(new_thread.thread_id).await, - }; - let config = match spawn_config { - Some(config) => config, - None => new_thread.thread.config_snapshot().await, + status, }; let session_telemetry = new_thread .thread diff --git a/codex-rs/core/src/agent/control/spawn_guard.rs b/codex-rs/core/src/agent/control/spawn_guard.rs index d8befedd40..c9095893b1 100644 --- a/codex-rs/core/src/agent/control/spawn_guard.rs +++ b/codex-rs/core/src/agent/control/spawn_guard.rs @@ -1,4 +1,4 @@ -//! Owns a spawned child until its initial input is accepted. +//! Owns a provisional child until the spawn result can be returned without yielding. use super::AgentTreeMembership; use crate::thread_manager::ThreadManagerState; diff --git a/codex-rs/core/src/agent/control_tests.rs b/codex-rs/core/src/agent/control_tests.rs index d22fb97d29..6d97eceb2e 100644 --- a/codex-rs/core/src/agent/control_tests.rs +++ b/codex-rs/core/src/agent/control_tests.rs @@ -820,6 +820,69 @@ async fn spawn_failure_context_preserves_shutdown_during_child_startup() { .expect("spawn releases its membership"); } +#[tokio::test] +async fn cancelled_spawn_after_input_admission_discards_child() { + let harness = AgentControlHarness::new().await; + let state = harness.control.runtime.upgrade().expect("live manager"); + let parent_id = ThreadId::new(); + let mut spawn = harness.control.spawn(SpawnRequest { + caller: parent_id, + config: harness.config.clone(), + input: AgentInput::UserInput(text_input("child task")), + source: SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id: parent_id, + depth: 1, + agent_path: None, + agent_nickname: None, + agent_role: None, + }), + options: SpawnAgentOptions::default(), + }); + let child = poll_fn(|cx| { + assert!(spawn.as_mut().poll(cx).is_pending()); + let threads = state.threads.try_read().expect("uncontended manager"); + // On this single-thread runtime the session loop cannot consume the first + // input until this poll returns. Leave spawn waiting for its admission reply. + match threads.values().find(|thread| !thread.io.tx_sub.is_empty()) { + Some(child) => Poll::Ready(Arc::clone(child)), + None => Poll::Pending, + } + }) + .await; + let child_id = child.session.thread_id(); + // This no-op goes through the same FIFO submission loop. Its receipt proves the + // earlier input has been admitted, without polling the spawn future again. + child + .update_thread_settings(ThreadSettingsOverrides::default()) + .await + .expect("input admission completed before the settings barrier"); + { + // Only the final status lookup now needs the manager lock. Cancel there, + // after input admission but before the backend can return its receipt. + let _threads = state.threads.write().await; + assert!( + spawn + .as_mut() + .poll(&mut Context::from_waker(Waker::noop())) + .is_pending() + ); + drop(spawn); + } + assert!( + harness.control.get_agent_metadata(child_id).is_none(), + "cancelled spawn must not leave an accepted child registered" + ); + child.wait_until_terminated().await; + assert_thread_not_loaded(&harness.manager, child_id).await; + harness + .control + .runtime + .request_shutdown() + .wait() + .await + .expect("provisional cleanup releases tree membership"); +} + #[tokio::test] async fn spawn_agent_errors_when_manager_dropped() { let control = LocalAgentControl::default(); diff --git a/codex-rs/core/tests/suite/pending_input_persistence.rs b/codex-rs/core/tests/suite/pending_input_persistence.rs index f12d89372c..0ebb340663 100644 --- a/codex-rs/core/tests/suite/pending_input_persistence.rs +++ b/codex-rs/core/tests/suite/pending_input_persistence.rs @@ -595,8 +595,16 @@ async fn local_preparation_is_durable_before_first_input_and_survives_restart( Ok(()) } +#[derive(Clone, Copy)] +enum SpawnCaller { + NativeV1, + NativeV2, +} + +#[test_case(SpawnCaller::NativeV1; "native_v1")] +#[test_case(SpawnCaller::NativeV2; "native_v2")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn cancelled_spawn_discards_provisional_child() -> anyhow::Result<()> { +async fn cancelled_spawn_discards_provisional_child(caller: SpawnCaller) -> anyhow::Result<()> { let server = responses::start_mock_server().await; let (checkpoints, mut checkpoint_requests) = mpsc::unbounded_channel(); let store = Arc::new(GatedCheckpointStore { @@ -610,8 +618,14 @@ async fn cancelled_spawn_discards_provisional_child() -> anyhow::Result<()> { let test = test_codex() .with_thread_store(store) .with_history_mode(ThreadHistoryMode::Legacy) - .with_config(|config| { + .with_config(move |config| { config.agent_max_depth = 1; + if matches!(caller, SpawnCaller::NativeV2) { + config + .features + .enable(Feature::MultiAgentV2) + .expect("enable V2"); + } config .features .enable(Feature::Collab) @@ -619,15 +633,24 @@ async fn cancelled_spawn_discards_provisional_child() -> anyhow::Result<()> { }) .build_with_auto_env(&server) .await?; + let (namespace, arguments) = match caller { + SpawnCaller::NativeV1 => ( + "multi_agent_v1", + json!({"message":"child startup task", "task_name":"worker", "fork_context":true}), + ), + SpawnCaller::NativeV2 => ( + "collaboration", + json!({"message":"child startup task", "task_name":"worker", "fork_turns":"all"}), + ), + }; responses::mount_sse_once( &server, responses::sse(vec![ responses::ev_function_call_with_namespace( "spawn-checkpoint", - "multi_agent_v1", + namespace, "spawn_agent", - &json!({"message":"child startup task", "task_name":"worker", "fork_context":true}) - .to_string(), + &arguments.to_string(), ), responses::ev_completed("parent-spawn"), ]),