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"), ]),