diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index 3ce775cc7f..541eedbb35 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/codex-rs/core/src/agent/control.rs @@ -15,7 +15,6 @@ use crate::config::RolloutBudgetConfig; use crate::context::ContextualUserFragment; use crate::context::SubagentNotification; use crate::environment_selection::TurnEnvironmentSnapshot; -use crate::session::emit_subagent_session_started; use crate::session_prefix::format_inter_agent_completion_message; use crate::thread_manager::ResumeThreadWithHistoryOptions; use crate::thread_manager::ThreadIdGenerator; @@ -62,6 +61,10 @@ pub(crate) use self::runtime::AgentTreeShutdownState; pub(crate) use self::runtime::AgentTreeTeardownGuard; pub(crate) use self::runtime::LocalAgentRuntime; pub(crate) use self::watch::StatusSubscription; +pub(in crate::agent) use spawn::effective_spawn_config; +pub(in crate::agent) use spawn::fresh_child_history_and_tools; +pub(in crate::agent) use spawn_guard::PendingSpawn; +pub(in crate::agent) use spawn_guard::SpawnOwner; mod api; mod budget; @@ -626,14 +629,8 @@ impl LocalAgentControl { state: &Arc, session_source: Option<&SessionSource>, ) -> Option { - let Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { - parent_thread_id, .. - })) = session_source - else { - return None; - }; - - let parent_thread = state.get_thread(*parent_thread_id).await.ok()?; + let parent_thread_id = session_source?.parent_thread_id()?; + let parent_thread = state.get_thread(parent_thread_id).await.ok()?; Some(parent_thread.session.inherited_environment_snapshot().await) } @@ -674,14 +671,15 @@ async fn inherited_exec_policy_for_source( session_source: Option<&SessionSource>, child_config: &Config, ) -> Option> { - let Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { - parent_thread_id, .. - })) = session_source - else { - return None; - }; + let parent_thread_id = session_source?.parent_thread_id()?; + let parent_thread = state.get_thread(parent_thread_id).await.ok()?; + inherited_exec_policy_for_parent(&parent_thread, child_config).await +} - let parent_thread = state.get_thread(*parent_thread_id).await.ok()?; +pub(in crate::agent) async fn inherited_exec_policy_for_parent( + parent_thread: &crate::CodexThread, + child_config: &Config, +) -> Option> { let parent_config = parent_thread.session.get_config().await; if !crate::exec_policy::child_uses_parent_exec_policy(&parent_config, child_config) { return None; diff --git a/codex-rs/core/src/agent/control/runtime_context.rs b/codex-rs/core/src/agent/control/runtime_context.rs index d4f8f9b86e..133c2c1535 100644 --- a/codex-rs/core/src/agent/control/runtime_context.rs +++ b/codex-rs/core/src/agent/control/runtime_context.rs @@ -130,7 +130,7 @@ impl LocalAgentRuntime { Ok(descendants) } - pub(super) fn upgrade(&self) -> CodexResult> { + pub(in crate::agent) fn upgrade(&self) -> CodexResult> { self.manager.upgrade().ok_or_else(|| { CodexErr::UnsupportedOperation("thread manager dropped".to_string()) .with_agent_context(AgentErrorContext::ManagerUnavailable) diff --git a/codex-rs/core/src/agent/control/spawn.rs b/codex-rs/core/src/agent/control/spawn.rs index 16323358e1..9f1660cf61 100644 --- a/codex-rs/core/src/agent/control/spawn.rs +++ b/codex-rs/core/src/agent/control/spawn.rs @@ -1,5 +1,6 @@ use super::residency::is_v2_resident_session_source; use super::spawn_guard::PendingSpawn; +use super::spawn_guard::SpawnOwner; use super::spawn_telemetry::SpawnMeasurements; use super::spawn_telemetry::record_spawn_success; use super::*; @@ -16,10 +17,10 @@ use crate::codex_thread::ThreadConfigSnapshot; use crate::config::PermissionProfileSnapshot; use crate::thread_manager::StartThreadOptions; use codex_features::Feature; +use codex_protocol::dynamic_tools::DynamicToolSpec; use codex_protocol::error::AgentErrorContext; use codex_protocol::intersect_effective_permission_profiles; use codex_protocol::protocol::EnvironmentConfigState; -use codex_thread_store::PersistContext; use codex_utils_path_uri::PathUri; use futures::StreamExt; use futures::stream; @@ -39,7 +40,9 @@ pub(super) enum SpawnInitialInput { InterAgentCommunication(InterAgentCommunication, AgentCommunicationContext), } -async fn effective_spawn_config(thread: &CodexThread) -> CodexResult { +pub(in crate::agent) async fn effective_spawn_config( + thread: &CodexThread, +) -> CodexResult { let mut config = thread.config_snapshot().await; let model_info = thread .thread_extension_data() @@ -74,6 +77,26 @@ pub(super) fn agent_nickname_candidates(config: &Config, role_name: Option<&str> .collect() } +pub(in crate::agent) async fn fresh_child_history_and_tools( + parent: &CodexThread, + config: &Config, + version: MultiAgentVersion, +) -> (Option, Vec) { + let history_mode = matches!( + parent.config_snapshot().await.history_mode, + ThreadHistoryMode::Paginated + ) + .then_some(ThreadHistoryMode::Paginated); + let tools = if version == MultiAgentVersion::V2 + && config.features.enabled(Feature::MultiAgentV2DynamicTools) + { + parent.session.dynamic_tools().await.to_vec() + } else { + Vec::new() + }; + (history_mode, tools) +} + impl LocalAgentControl { /// Restore persisted V2 agent identities without reopening their runtimes. pub(crate) async fn restore_v2_agent_metadata( @@ -511,15 +534,11 @@ impl LocalAgentControl { }) }; let subagent_analytics = - if let SessionSource::SubAgent(source @ SubAgentSource::ThreadSpawn { .. }) = - &session_source + if let SessionSource::SubAgent(SubAgentSource::ThreadSpawn { .. }) = &session_source && let Some(parent_thread_id) = parent_thread_id && let Ok(parent) = state.get_thread(parent_thread_id).await { - Some(( - parent.session.app_server_client_metadata().await, - source.clone(), - )) + Some(parent.session.app_server_client_metadata().await) } else { None }; @@ -559,22 +578,14 @@ impl LocalAgentControl { runtime.registry.clear_evicted_environments(thread_id); residency_slot.commit(reloaded_thread.thread_id); // Register before listeners can forward events from the resumed thread. - if let Some((client_metadata, source)) = subagent_analytics { - let thread_config = reloaded_thread.thread.config_snapshot().await; - emit_subagent_session_started( - &reloaded_thread - .thread - .session - .services - .analytics_events_client, - client_metadata, - reloaded_thread.thread.session.session_id(), - reloaded_thread.thread_id, - thread_config.parent_thread_id, - thread_config, - source, - /*resumed_created_at*/ Some(stored_thread.created_at), - ); + if let Some(client_metadata) = subagent_analytics { + reloaded_thread + .thread + .emit_subagent_session_started( + client_metadata, + Some(stored_thread.created_at), + ) + .await; } state.notify_thread_created(reloaded_thread.thread_id); Ok(residency_pin) @@ -734,17 +745,9 @@ impl LocalAgentControl { if let Some(parent_thread_id) = parent_thread_id && let Ok(parent_thread) = state.get_thread(parent_thread_id).await { - start_options.history_mode = matches!( - parent_thread.config_snapshot().await.history_mode, - ThreadHistoryMode::Paginated - ) - .then_some(ThreadHistoryMode::Paginated); - if multi_agent_version == MultiAgentVersion::V2 - && config.features.enabled(Feature::MultiAgentV2DynamicTools) - { - start_options.dynamic_tools = - parent_thread.session.dynamic_tools().await.to_vec(); - } + (start_options.history_mode, start_options.dynamic_tools) = + fresh_child_history_and_tools(&parent_thread, &config, multi_agent_version) + .await; } None } @@ -762,14 +765,16 @@ impl LocalAgentControl { .map_err(|err| err.with_agent_context(AgentErrorContext::ChildStartup))?; let child_create = child_create_started_at.elapsed(); agent_metadata.agent_id = Some(new_thread.thread_id); - let mut pending_spawn = - PendingSpawn::new(Arc::clone(&state), new_thread.thread_id, membership); + let mut pending_spawn = PendingSpawn::new( + Arc::clone(&state), + new_thread.thread_id, + membership, + SpawnOwner::Local, + ); - if let Some(SessionSource::SubAgent( - subagent_source @ SubAgentSource::ThreadSpawn { - parent_thread_id, .. - }, - )) = notification_source.as_ref() + if let Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id, .. + })) = notification_source.as_ref() { let client_metadata = match state.get_thread(*parent_thread_id).await { Ok(parent_thread) => parent_thread.session.app_server_client_metadata().await, @@ -785,18 +790,10 @@ impl LocalAgentControl { } } }; - let thread_config = new_thread.thread.config_snapshot().await; - let parent_thread_id = thread_config.parent_thread_id; - emit_subagent_session_started( - &new_thread.thread.session.services.analytics_events_client, - client_metadata, - new_thread.thread.session.session_id(), - new_thread.thread_id, - parent_thread_id, - thread_config, - subagent_source.clone(), - /*resumed_created_at*/ None, - ); + new_thread + .thread + .emit_subagent_session_started(client_metadata, /*resumed_created_at*/ None) + .await; } let control = self.clone(); @@ -815,10 +812,7 @@ impl LocalAgentControl { let durability_wait_started_at = Instant::now(); if options.fork_mode.is_some() { tokio::join!( - new_thread - .thread - .session - .ensure_rollout_materialized(PersistContext::Standard), + new_thread.thread.ensure_rollout_materialized(), pending_spawn.wait_for_edge(), ); } else { diff --git a/codex-rs/core/src/agent/control/spawn_guard.rs b/codex-rs/core/src/agent/control/spawn_guard.rs index c9095893b1..1918af66dc 100644 --- a/codex-rs/core/src/agent/control/spawn_guard.rs +++ b/codex-rs/core/src/agent/control/spawn_guard.rs @@ -12,22 +12,30 @@ use std::sync::Arc; use tokio::task::JoinHandle; use tracing::warn; -pub(super) struct PendingSpawn { +pub(in crate::agent) enum SpawnOwner { + Local, + Host, +} + +pub(in crate::agent) struct PendingSpawn { state: Arc, pending: Option<(ThreadId, AgentTreeMembership)>, edge_write: Option>, + owner: SpawnOwner, } impl PendingSpawn { - pub(super) fn new( + pub(in crate::agent) fn new( state: Arc, child: ThreadId, membership: AgentTreeMembership, + owner: SpawnOwner, ) -> Self { Self { state, pending: Some((child, membership)), edge_write: None, + owner, } } @@ -45,7 +53,7 @@ impl PendingSpawn { self.edge_write = None; } - pub(super) fn disarm(mut self) -> AgentTreeMembership { + pub(in crate::agent) fn disarm(mut self) -> AgentTreeMembership { let Some((_child, membership)) = self.pending.take() else { unreachable!("pending spawn must own agent-tree membership"); }; @@ -61,6 +69,7 @@ impl Drop for PendingSpawn { let state = Arc::clone(&self.state); let teardown = membership.into_teardown_guard("child_spawn_cleanup", Some(child)); let edge_write = self.edge_write.take(); + let local_owner = matches!(self.owner, SpawnOwner::Local); drop(tokio::spawn(async move { if let Some(thread) = state.remove_thread(&child).await { if let Err(error) = thread.shutdown_and_wait().await { @@ -85,7 +94,8 @@ impl Drop for PendingSpawn { if let Some(edge_write) = edge_write { let _ = edge_write.await; } - if let Some(store) = state.agent_graph_store() + if local_owner + && let Some(store) = state.agent_graph_store() && let Err(error) = store .set_thread_spawn_edge_status(child, ThreadSpawnEdgeStatus::Closed) .await diff --git a/codex-rs/core/src/agent/host/invocation.rs b/codex-rs/core/src/agent/host/invocation.rs new file mode 100644 index 0000000000..2bbf1b58ca --- /dev/null +++ b/codex-rs/core/src/agent/host/invocation.rs @@ -0,0 +1,272 @@ +//! Bind host child preparation and startup to a live tool invocation. + +use super::spawn::SpawnDestination; +use super::spawn::prepare_spawn; +use super::spawn::start_prepared; +use super::types::HostThreadSpawnReservation; +use super::types::PreparedThreadSpawn; +use crate::codex_thread::CodexThread; +use crate::function_tool::FunctionCallError; +use crate::session::session::Session; +use crate::session::step_context::StepContext; +use crate::session::turn_context::TurnContext; +use crate::thread_spawn::ChildSpawnSettings; +use crate::thread_spawn::ThreadCreateRequest; +use crate::thread_spawn::ThreadSpawnError; +use crate::thread_spawn::ThreadSpawnResult; +use crate::tools::context::ToolInvocation; +use codex_protocol::AgentPath; +use codex_protocol::error::CodexErr; +use codex_protocol::error::Result as CodexResult; +use codex_protocol::protocol::MultiAgentVersion; +use futures::TryFutureExt; +use std::collections::HashMap; +use std::collections::hash_map::Entry; +use std::future::Future; +use std::sync::Arc; +use std::sync::Mutex; +use std::sync::MutexGuard; +use std::sync::PoisonError; +use std::sync::Weak; +use tokio_util::sync::CancellationToken; + +struct InvocationContext { + step: Weak, + finished: Mutex, +} + +impl InvocationContext { + fn finished(&self) -> MutexGuard<'_, CancellationToken> { + self.finished.lock().unwrap_or_else(PoisonError::into_inner) + } +} + +/// Install in thread extension data to enable live V2 tool lookup and require unique active +/// call IDs within each turn. Host-spawned children inherit this opt-in. +pub struct HostToolInvocationLookup; + +#[derive(Default)] +struct InvocationContexts(Mutex>>); + +/// A handle to one tool call, valid only on its issuing runtime until dispatch ends. +/// It holds the runtime and step weakly. Child preparation captures history and policy. +pub struct ToolInvocationContext { + session: Weak, + call_id: String, + context: Arc, +} + +impl ToolInvocationContext { + /// Order publication with revocation of the issuing tool call. + pub(super) fn lock_active( + &self, + ) -> Result, ThreadSpawnError> { + let finished = self.context.finished(); + if finished.is_cancelled() { + return Err(ThreadSpawnError::Preparation( + "spawning tool call finished".into(), + )); + } + Ok(finished) + } +} + +pub(crate) fn tool_invocation( + thread: &CodexThread, + turn_id: &str, + call_id: &str, +) -> CodexResult { + let context = thread + .session + .services + .thread_extension_data + .get::() + .and_then(|contexts| { + contexts + .0 + .lock() + .unwrap_or_else(PoisonError::into_inner) + .get(&(turn_id.to_owned(), call_id.to_owned())) + .and_then(Weak::upgrade) + }) + .filter(|context| !context.finished().is_cancelled()) + .ok_or_else(|| CodexErr::InvalidRequest("tool call is not active".into()))?; + Ok(ToolInvocationContext { + session: Arc::downgrade(&thread.session), + call_id: call_id.to_owned(), + context, + }) +} + +pub(crate) struct ReservedInvocation { + contexts: Arc, + key: (String, String), +} + +/// Keeps the dispatch's ID reserved until its terminal callback returns. +#[derive(Default)] +pub(crate) struct InvocationReservation(Mutex>); + +impl InvocationReservation { + pub(crate) fn reserve( + &self, + session: &Session, + turn: &TurnContext, + call_id: &str, + ) -> Result<(), FunctionCallError> { + if turn.multi_agent_version != MultiAgentVersion::V2 + || session + .services + .thread_extension_data + .get::() + .is_none() + { + return Ok(()); + } + + let contexts = session + .services + .thread_extension_data + .get_or_init(InvocationContexts::default); + let key = (turn.sub_id.clone(), call_id.to_owned()); + match contexts + .0 + .lock() + .unwrap_or_else(PoisonError::into_inner) + .entry(key.clone()) + { + Entry::Vacant(entry) => { + entry.insert(Weak::new()); + } + Entry::Occupied(_) => { + return Err(FunctionCallError::RespondToModel(format!( + "tool call ID {call_id} is already active in this turn" + ))); + } + } + *self.0.lock().unwrap_or_else(PoisonError::into_inner) = + Some(ReservedInvocation { contexts, key }); + Ok(()) + } + + pub(crate) fn activate(&self, invocation: &ToolInvocation) -> Option { + let reservation = self.0.lock().unwrap_or_else(PoisonError::into_inner); + let reservation = reservation.as_ref()?; + let context = Arc::new(InvocationContext { + step: Arc::downgrade(&invocation.step_context), + finished: Mutex::new(invocation.cancellation_token.child_token()), + }); + let mut contexts = reservation + .contexts + .0 + .lock() + .unwrap_or_else(PoisonError::into_inner); + *contexts.get_mut(&reservation.key)? = Arc::downgrade(&context); + Some(InvocationGuard(context)) + } + + pub(crate) fn take(&self) -> Option { + self.0.lock().unwrap_or_else(PoisonError::into_inner).take() + } +} + +/// Revokes host access when execution ends, including unwinding or cancellation. +pub(crate) struct InvocationGuard(Arc); + +impl Drop for InvocationGuard { + fn drop(&mut self) { + self.0.finished().cancel(); + } +} + +impl Drop for ReservedInvocation { + fn drop(&mut self) { + self.contexts + .0 + .lock() + .unwrap_or_else(PoisonError::into_inner) + .remove(&self.key); + } +} + +pub(crate) async fn create( + thread: &CodexThread, + invocation: &ToolInvocationContext, + request: ThreadCreateRequest, + reservation: HostThreadSpawnReservation, +) -> Result { + with_spawn_request( + thread, + invocation, + request, + |request, task_name| async move { + let prepared = Box::pin(prepare_spawn( + thread, + request, + task_name, + SpawnDestination::SameManager, + )) + .await?; + Box::pin(start_prepared( + &thread.session.services.local_agent_runtime, + prepared, + reservation, + Some(invocation), + )) + .await + }, + ) + .await +} + +pub(crate) async fn prepare( + thread: &CodexThread, + invocation: &ToolInvocationContext, + request: ThreadCreateRequest, +) -> Result { + with_spawn_request(thread, invocation, request, |request, task_name| { + prepare_spawn(thread, request, task_name, SpawnDestination::Detached) + .map_err(ThreadSpawnError::from) + }) + .await +} + +async fn with_spawn_request( + thread: &CodexThread, + invocation: &ToolInvocationContext, + request: ThreadCreateRequest, + operation: F, +) -> Result +where + F: FnOnce(ChildSpawnSettings, AgentPath) -> Fut, + Fut: Future>, +{ + let session = &thread.session; + if !invocation.session.ptr_eq(&Arc::downgrade(session)) { + return Err(CodexErr::InvalidRequest( + "tool invocation belongs to a different runtime".into(), + ) + .into()); + } + let context = &invocation.context; + let finished = context.finished().clone(); + let step = context.step.upgrade().ok_or_else(|| { + ThreadSpawnError::Preparation("spawning step is no longer available".into()) + })?; + tokio::select! { + biased; + _ = finished.cancelled() => Err(ThreadSpawnError::Preparation("spawning tool call finished".into())), + result = async { + let _caller_pin = Arc::clone(&thread.residency_gate).read_owned().await; + if !thread.is_running() { + return Err(CodexErr::ThreadNotFound(session.thread_id).into()); + } + let (request, task_name) = crate::thread_spawn::prepare(session, &step, &invocation.call_id, request).await?; + operation(request, task_name).await + } => result, + } +} + +#[cfg(test)] +#[path = "invocation_tests.rs"] +mod tests; diff --git a/codex-rs/core/src/agent/host/invocation_tests.rs b/codex-rs/core/src/agent/host/invocation_tests.rs new file mode 100644 index 0000000000..fbd13e9c18 --- /dev/null +++ b/codex-rs/core/src/agent/host/invocation_tests.rs @@ -0,0 +1,161 @@ +#![cfg(test)] + +use super::HostToolInvocationLookup; +use super::InvocationContexts; +use crate::session::step_context::StepContext; +use crate::tools::context::ToolCallSource; +use crate::tools::context::ToolPayload; +use crate::tools::handlers::ListMcpResourcesHandler; +use crate::tools::parallel::ToolCallRuntime; +use crate::tools::registry::ToolRegistry; +use crate::tools::router::ToolCall; +use crate::tools::router::ToolRouter; +use codex_extension_api::ToolFinishInput; +use codex_extension_api::ToolLifecycleContributor; +use codex_extension_api::ToolLifecycleFuture; +use codex_extension_api::ToolStartInput; +use codex_protocol::openai_models::ToolMode; +use codex_protocol::protocol::MultiAgentVersion; +use codex_tools::ToolName; +use std::sync::Arc; +use std::time::Duration; +use tokio::sync::Notify; +use tokio_util::sync::CancellationToken; + +#[derive(Default)] +struct CallbackGate { + entered: Notify, + release: Notify, +} + +impl ToolLifecycleContributor for CallbackGate { + fn on_tool_start<'a>(&'a self, input: ToolStartInput<'a>) -> ToolLifecycleFuture<'a> { + Box::pin(async move { + self.entered.notify_one(); + self.release.notified().await; + if matches!( + input.source, + codex_extension_api::ToolCallSource::CodeMode { .. } + ) { + panic!("dispatch unwound"); + } + }) + } + + fn on_tool_finish<'a>(&'a self, _input: ToolFinishInput<'a>) -> ToolLifecycleFuture<'a> { + Box::pin(async { + self.entered.notify_one(); + self.release.notified().await; + }) + } +} + +#[tokio::test] +async fn dispatch_revokes_access_and_holds_call_id_through_callbacks() -> anyhow::Result<()> { + let (mut session, mut turn, _events) = + crate::session::tests::make_session_and_context_with_rx().await; + Arc::get_mut(&mut turn).unwrap().multi_agent_version = MultiAgentVersion::V2; + session + .services + .thread_extension_data + .insert(HostToolInvocationLookup); + let gate = Arc::new(CallbackGate::default()); + let mut extensions = codex_extension_api::ExtensionRegistryBuilder::new(); + extensions.tool_lifecycle_contributor(gate.clone()); + Arc::get_mut(&mut session).unwrap().services.extensions = Arc::new(extensions.build()); + let mut registry = ToolRegistry::default(); + registry.add(ListMcpResourcesHandler::new(/*messages*/ None)); + let router = ToolRouter::from_parts( + registry, + vec![], + ToolMode::Direct, + Default::default(), + /*tool_namespaces_info*/ None, + &[], + ); + let step = StepContext::for_test(turn.clone()).with_tool_router_for_test(Arc::new(router)); + let runtime = ToolCallRuntime::new(session.clone(), step.clone(), Default::default()); + let dispatch = |source, token, args: &str| { + runtime.clone().handle_tool_call_with_source( + step.clone(), + ToolCall { + tool_name: ToolName::plain("list_mcp_resources"), + call_id: "shared".into(), + payload: ToolPayload::Function { + arguments: args.into(), + }, + encrypted_function_args: None, + }, + source, + token, + Default::default(), + ) + }; + let wait_limit = Duration::from_secs(5); + let duplicate = async |args: &str| { + let result = tokio::time::timeout( + wait_limit, + dispatch(ToolCallSource::Direct, CancellationToken::new(), args), + ) + .await + .unwrap(); + assert!( + matches!(result, Err(crate::function_tool::FunctionCallError::RespondToModel(msg)) if msg.contains("already active")) + ); + }; + let key = (turn.sub_id.clone(), String::from("shared")); + let at_callback = || tokio::time::timeout(wait_limit, gate.entered.notified()); + let contexts = session + .services + .thread_extension_data + .get_or_init(InvocationContexts::default); + for (source, abort) in [ + ( + ToolCallSource::CodeMode { + cell_id: "cell".into(), + runtime_tool_call_id: "nested".into(), + }, + false, + ), + (ToolCallSource::Direct, true), + (ToolCallSource::Direct, false), + ] { + let panics = matches!(source, ToolCallSource::CodeMode { .. }); + let token = CancellationToken::new(); + let owner = tokio::spawn(dispatch(source, token.clone(), "{}")); + at_callback() + .await + .expect("owner entered its start callback"); + let captured = contexts.0.lock().unwrap()[&key].upgrade().unwrap(); + duplicate("invalid json").await; + assert!(Arc::ptr_eq( + &captured, + &contexts.0.lock().unwrap()[&key].upgrade().unwrap() + )); + assert!(!captured.finished().is_cancelled()); + if abort { + token.cancel(); + } else { + gate.release.notify_one(); + } + if !panics { + at_callback() + .await + .expect("owner entered its terminal callback"); + assert!(captured.finished().is_cancelled()); + duplicate("{}").await; + gate.release.notify_one(); + } + let outcome = tokio::time::timeout(wait_limit, owner).await??; + if panics { + assert!( + matches!(outcome, Err(crate::function_tool::FunctionCallError::Fatal(message)) if message.contains("dispatch unwound")) + ); + } else { + outcome?; + } + assert!(captured.finished().is_cancelled()); + assert!(!contexts.0.lock().unwrap().contains_key(&key)); + } + Ok(()) +} diff --git a/codex-rs/core/src/agent/host/mod.rs b/codex-rs/core/src/agent/host/mod.rs new file mode 100644 index 0000000000..10ef102e14 --- /dev/null +++ b/codex-rs/core/src/agent/host/mod.rs @@ -0,0 +1,14 @@ +mod invocation; +mod spawn; +mod types; + +pub use invocation::HostToolInvocationLookup; +pub(crate) use invocation::InvocationReservation; +pub use invocation::ToolInvocationContext; +pub(crate) use invocation::create; +pub(crate) use invocation::prepare; +pub(crate) use invocation::tool_invocation; +pub(crate) use spawn::start_prepared; +pub use types::HostThreadSpawnReservation; +pub(crate) use types::PreparedChildInheritance; +pub use types::PreparedThreadSpawn; diff --git a/codex-rs/core/src/agent/host/spawn.rs b/codex-rs/core/src/agent/host/spawn.rs new file mode 100644 index 0000000000..9f539fee6d --- /dev/null +++ b/codex-rs/core/src/agent/host/spawn.rs @@ -0,0 +1,275 @@ +//! Capture child startup inputs on the source and start the reserved child at the destination. + +use super::invocation::HostToolInvocationLookup; +use super::invocation::ToolInvocationContext; +use super::types::HostThreadSpawnReservation; +use super::types::PreparedChildInheritance; +use super::types::PreparedThreadSpawn; +use crate::agent::control::AgentControlInit; +use crate::agent::control::LocalAgentRuntime; +use crate::agent::control::PendingSpawn; +use crate::agent::control::SpawnOwner; +use crate::agent::control::effective_spawn_config; +use crate::agent::control::fresh_child_history_and_tools; +use crate::agent::control::inherited_exec_policy_for_parent; +use crate::agent::history::prepare_forked_history; +use crate::agent::types::AgentMetadata; +use crate::agent::types::LiveAgent; +use crate::codex_thread::CodexThread; +use crate::exec_policy::ExecPolicyManager; +use crate::thread_manager::StartThreadOptions; +use crate::thread_spawn::ChildSpawnSettings; +use crate::thread_spawn::ThreadSpawnError; +use crate::thread_spawn::ThreadSpawnResult; +use codex_config::McpServerAuth; +use codex_features::Feature; +use codex_history::InitialHistory; +use codex_protocol::AgentPath; +use codex_protocol::SessionId; +use codex_protocol::error::AgentErrorContext; +use codex_protocol::error::CodexErr; +use codex_protocol::error::Result; +use codex_protocol::protocol::EnvironmentConfigState; +use codex_protocol::protocol::MultiAgentVersion; +use codex_protocol::protocol::ThreadSource; +use codex_thread_store::ListThreadAttachmentsParams; +use std::sync::Arc; + +pub(super) enum SpawnDestination { + SameManager, + Detached, +} + +pub(super) async fn prepare_spawn( + parent: &CodexThread, + request: ChildSpawnSettings, + agent_path: AgentPath, + destination: SpawnDestination, +) -> Result { + if matches!(destination, SpawnDestination::Detached) + && (request.config.application_network_policy.is_managed() + || request.config.application_auth_route_config.is_some() + || request + .config + .config_layer_stack + .cloud_config_binding() + .is_some()) + { + return Err(CodexErr::InvalidRequest( + "prepared spawn requires destination binding for host-managed application authority" + .into(), + )); + } + let runtime = &parent.session.services.local_agent_runtime; + let _membership = runtime.admit_start()?; + let state = runtime.upgrade()?; + let ChildSpawnSettings { + config, + source, + options, + } = request; + if matches!(destination, SpawnDestination::Detached) + && config.features.enabled(Feature::UseXaa) + && config.mcp_enterprise_managed_auth.is_some() + && config.mcp_servers.get().values().any(|server| { + server.enabled + && matches!(server.auth, McpServerAuth::EmaAuth) + && server.validate_ema_auth_transport().is_ok() + }) + { + return Err(CodexErr::InvalidRequest( + "prepared spawn requires destination binding for EMA-authenticated MCP servers".into(), + )); + } + let caller = parent.session.thread_id; + let instructions = parent.session.inherited_instructions().await; + if matches!(destination, SpawnDestination::Detached) + && (instructions.thread_provider.is_some() || instructions.user_provider.is_some()) + { + return Err(CodexErr::InvalidRequest( + "prepared spawn requires destination binding for inherited instruction providers" + .into(), + )); + } + let store = &parent.session.services.thread_store; + if matches!(destination, SpawnDestination::Detached) && store.supports_thread_attachments() { + let page = store + .list_thread_attachments(ListThreadAttachmentsParams { + thread_id: caller, + cursor: None, + limit: 1, + }) + .await + .map_err(|error| { + CodexErr::InvalidRequest(format!("cannot prepare child attachments: {error}")) + })?; + if !page.attachments.is_empty() { + return Err(CodexErr::InvalidRequest( + "prepared spawn requires a host handoff for thread attachments".into(), + )); + } + } + let snapshot = options.environments.as_ref().ok_or_else(|| { + CodexErr::InvalidRequest("prepared spawn requires captured environments".into()) + })?; + let environments = snapshot.inheritable_requests(); + if matches!(destination, SpawnDestination::Detached) + && environments + .iter() + .any(|environment| matches!(environment.config, EnvironmentConfigState::Pending)) + { + return Err(CodexErr::InvalidRequest( + "prepared spawn requires resolved environment configuration".into(), + )); + } + let (initial_history, history_mode, dynamic_tools) = if options.fork_mode.is_some() { + let (history, mode) = + prepare_forked_history(&state, &source, &options, &config, MultiAgentVersion::V2) + .await + .map_err(|error| error.with_agent_context(AgentErrorContext::ForkHistory))?; + (history, mode, Vec::new()) + } else { + let (mode, tools) = + fresh_child_history_and_tools(parent, &config, MultiAgentVersion::V2).await; + (InitialHistory::New, mode, tools) + }; + if !std::ptr::eq(state.get_thread(caller).await?.as_ref(), parent) { + return Err(CodexErr::ThreadNotFound(caller)); + } + let inherited_policy = inherited_exec_policy_for_parent(parent, &config).await; + let mut thread_options = StartThreadOptions { + session_source: Some(source), + thread_source: Some(ThreadSource::Subagent), + initial_history, + history_mode, + dynamic_tools, + disabled_plugin_ids: options.disabled_plugin_ids, + client_mcp_extensions: parent.session.services.client_mcp_extensions.clone(), + prompt_cache_key: parent + .session + .services + .model_client + .prompt_cache_key_policy(), + ..StartThreadOptions::new(config) + }; + thread_options + .thread_extension_init + .insert(snapshot.selected_capability_roots()); + thread_options + .thread_extension_init + .insert(HostToolInvocationLookup); + let exec_policy = match destination { + SpawnDestination::SameManager => { + thread_options.inherited_environments = options.environments; + inherited_policy + } + SpawnDestination::Detached => { + thread_options.environments = Some(environments); + Some(Arc::new(match inherited_policy { + Some(policy) => ExecPolicyManager::new(policy.current()), + None => ExecPolicyManager::load(&thread_options.config.config_layer_stack) + .await + .map_err(|error| CodexErr::InvalidRequest(error.to_string()))?, + })) + } + }; + Ok(PreparedThreadSpawn { + caller, + session_id: parent.session.session_id(), + agent_path, + client_metadata: parent.session.app_server_client_metadata().await, + thread_options, + inheritance: PreparedChildInheritance { + instructions, + exec_policy, + originator: parent.config_snapshot().await.originator, + attachment_source: (matches!(destination, SpawnDestination::SameManager) + && options.fork_mode.is_some()) + .then_some(caller), + }, + }) +} + +pub(crate) async fn start_prepared( + runtime: &LocalAgentRuntime, + prepared: PreparedThreadSpawn, + reservation: HostThreadSpawnReservation, + invocation: Option<&ToolInvocationContext>, +) -> std::result::Result { + if reservation.control.identity() != prepared.session_id + || reservation.thread_id == prepared.caller + || SessionId::from(reservation.thread_id) == prepared.session_id + { + return Err(ThreadSpawnError::Preparation( + "reservation does not match prepared child".into(), + )); + } + let membership = runtime.admit_start()?; + let state = runtime.upgrade()?; + if state + .threads + .read() + .await + .contains_key(&reservation.thread_id) + { + return Err(CodexErr::InvalidRequest(format!( + "thread {} is already running", + reservation.thread_id, + )) + .into()); + } + let task_name = prepared.agent_path; + let mut options = prepared.thread_options; + options.reserved_thread_id = Some(reservation.thread_id); + let child = state + .spawn_prepared_child_thread( + options, + AgentControlInit::Provided { + control: reservation.control, + runtime: runtime.clone(), + }, + prepared.caller, + prepared.inheritance, + ) + .await + .map_err(|error| error.with_agent_context(AgentErrorContext::ChildStartup))?; + let pending = PendingSpawn::new( + Arc::clone(&state), + child.thread_id, + membership, + SpawnOwner::Host, + ); + child.thread.emit_thread_ready_lifecycle().await; + child + .thread + .emit_subagent_session_started(prepared.client_metadata, /*resumed_created_at*/ None) + .await; + if child.session_configured.forked_from_id.is_some() { + child.thread.ensure_rollout_materialized().await; + } + let config = effective_spawn_config(&child.thread).await?; + let status = child.thread.io.agent_status.borrow().clone(); + let _active = invocation + .map(ToolInvocationContext::lock_active) + .transpose()?; + let _membership = pending.disarm(); + state.notify_thread_created(child.thread_id); + Ok(ThreadSpawnResult { + agent: LiveAgent { + thread_id: child.thread_id, + metadata: AgentMetadata { + agent_id: Some(child.thread_id), + agent_path: Some(task_name.clone()), + agent_nickname: child.thread.session_source.get_nickname(), + agent_role: child.thread.session_source.get_agent_role(), + }, + status, + }, + config, + task_name, + }) +} + +#[cfg(test)] +#[path = "spawn_tests.rs"] +mod tests; diff --git a/codex-rs/core/src/agent/host/spawn_tests.rs b/codex-rs/core/src/agent/host/spawn_tests.rs new file mode 100644 index 0000000000..c0cb8e793b --- /dev/null +++ b/codex-rs/core/src/agent/host/spawn_tests.rs @@ -0,0 +1,254 @@ +#![cfg(test)] + +//! Tests live child bindings and the capture boundary for detached child authority. + +use super::SpawnDestination; +use super::prepare_spawn; +use crate::HostThreadSpawnReservation; +use crate::HostToolInvocationLookup; +use crate::ThreadManager; +use crate::agent::host::InvocationReservation; +use crate::agent::types::SpawnAgentForkMode; +use crate::config::Config; +use crate::config::ConfigBuilder; +use crate::thread_manager::StartThreadOptions; +use crate::thread_spawn::ThreadCreateRequest; +use crate::thread_spawn::ThreadSpawnError; +use crate::tools::context::ToolCallSource; +use crate::tools::context::ToolInvocation; +use crate::tools::context::ToolPayload; +use codex_config::ScopedSkillsConfig; +use codex_exec_server::ExecServerError; +use codex_exec_server::RemoteEnvironmentOptions; +use codex_features::Feature; +use codex_login::CodexAuth; +use codex_protocol::ThreadId; +use codex_protocol::approvals::ExecPolicyAmendment; +use codex_protocol::protocol::EnvironmentConfigState; +use codex_tools::ToolName; +use pretty_assertions::assert_eq; +use std::future::poll_fn; +use std::net::TcpListener; +use std::sync::Arc; +use std::task::Poll; +use std::time::Duration; +use tokio_util::sync::CancellationToken; + +async fn spawn_fixture() -> anyhow::Result<(tempfile::TempDir, Config, ThreadManager)> { + let home = tempfile::tempdir()?; + let config = ConfigBuilder::without_managed_config_for_tests() + .codex_home(home.path().to_path_buf()) + .build() + .await?; + let manager = ThreadManager::with_models_provider_and_home_for_tests( + CodexAuth::from_api_key("dummy"), + config.model_provider.clone(), + home.path().to_path_buf(), + Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), + ); + Ok((home, config, manager)) +} + +fn child_request(fork_mode: Option) -> ThreadCreateRequest { + ThreadCreateRequest { + task_name: "worker".into(), + fork_mode, + agent_type: None, + model: None, + reasoning_effort: None, + } +} + +#[tokio::test] +async fn same_manager_creation_retains_step_executor_and_shares_exec_policy() -> anyhow::Result<()> +{ + let (_home, mut config, manager) = spawn_fixture().await?; + config.features.enable(Feature::MultiAgentV2)?; + config.features.enable(Feature::DeferredExecutor)?; + let environments = manager.environment_manager(); + // Leave the handshake pending so the registration stays attached throughout the spawn. + let listener = TcpListener::bind(("127.0.0.1", 0))?; + let url = format!("ws://{}", listener.local_addr()?); + environments.upsert_environment_with_options( + "workspace".into(), + RemoteEnvironmentOptions { + exec_server_url: url.clone(), + connect_timeout: Some(Duration::from_secs(300)), + http_headers: Default::default(), + }, + ScopedSkillsConfig { + required: vec!["original-skill".into()], + }, + )?; + let mut requests = manager.default_environment_requests(&config.cwd, &config.workspace_roots); + requests + .iter_mut() + .find(|request| request.environment_id == "workspace") + .expect("registered workspace") + .config = EnvironmentConfigState::Pending; + let parent = manager + .start_thread(StartThreadOptions { + environments: Some(requests), + ..StartThreadOptions::new(config) + }) + .await? + .thread; + parent + .thread_extension_data() + .insert(HostToolInvocationLookup); + let session = &parent.session; + let step = session + .capture_step_context(session.new_default_turn().await, &CancellationToken::new()) + .await?; + let call = ToolInvocation { + session: Arc::clone(session), + turn: Arc::clone(&step.turn), + step_context: step, + cancellation_token: CancellationToken::new(), + tracker: Default::default(), + call_id: "spawn-call".into(), + tool_name: ToolName::plain("host_spawn"), + source: ToolCallSource::Direct, + payload: ToolPayload::Function { + arguments: "{}".into(), + }, + environments: Default::default(), + }; + let reservation = InvocationReservation::default(); + reservation.reserve(session, &call.turn, &call.call_id)?; + assert!( + parent + .tool_invocation(&call.turn.sub_id, &call.call_id) + .is_err() + ); + let guard = reservation.activate(&call); + let invocation = parent.tool_invocation(&call.turn.sub_id, &call.call_id)?; + environments.upsert_environment("workspace".into(), url, /*connect_timeout*/ None)?; + assert!(matches!( + parent.prepare_child(&invocation, child_request(/*fork_mode*/ None)).await, + Err(ThreadSpawnError::Backend(error)) if error.to_string().contains("resolved environment configuration") + )); + let child_id = ThreadId::new(); + parent + .create_child( + &invocation, + child_request(/*fork_mode*/ None), + HostThreadSpawnReservation { + thread_id: child_id, + control: Arc::clone(&session.services.agent_control), + }, + ) + .await?; + let child = manager.get_thread(child_id).await?; + assert!( + child + .thread_extension_data() + .get::() + .is_some() + ); + let snapshot = child.session.services.turn_environments.snapshot_now(); + assert_eq!( + snapshot + .required_skills() + .find(|required| required.environment_id == "workspace") + .expect("child retained the original executor registration") + .skill_names, + ["original-skill"] + ); + let pending = snapshot + .starting() + .find(|starting| starting.selection.environment_id == "workspace") + .expect("child awaits the workspace"); + let failure = "owner failed to configure workspace"; + parent + .environment_failed(&pending.selection, failure.into()) + .await?; + let error = tokio::time::timeout(Duration::from_secs(5), pending.wait_until_ready()) + .await? + .expect_err("child should receive the owner's failure"); + assert!(matches!(error.as_ref(), ExecServerError::Protocol(message) if message == failure)); + assert!(Arc::ptr_eq( + &session.services.exec_policy, + &child.session.services.exec_policy, + )); + drop(guard); + assert!( + parent + .tool_invocation(&call.turn.sub_id, &call.call_id) + .is_err() + ); + drop(reservation); + let next = InvocationReservation::default(); + next.reserve(session, &call.turn, &call.call_id)?; + let _next_guard = next.activate(&call); + parent.tool_invocation(&call.turn.sub_id, &call.call_id)?; + assert!(matches!( + parent.prepare_child(&invocation, child_request(/*fork_mode*/ None)).await, + Err(ThreadSpawnError::Preparation(message)) if message == "spawning tool call finished" + )); + child.shutdown_and_wait().await?; + parent.shutdown_and_wait().await?; + Ok(()) +} + +#[tokio::test] +async fn prepared_spawn_captures_policy_after_fork_history() -> anyhow::Result<()> { + let (home, config, manager) = spawn_fixture().await?; + let parent = manager + .start_thread(StartThreadOptions::new(config)) + .await? + .thread; + let session = &parent.session; + let step = session + .capture_step_context(session.new_default_turn().await, &CancellationToken::new()) + .await?; + let (request, path) = crate::thread_spawn::prepare( + session, + &step, + "fork-call", + child_request(Some(SpawnAgentForkMode::FullHistory)), + ) + .await?; + let rollout = parent + .rollout_path() + .expect("local parent has a rollout path"); + assert!(!rollout.try_exists()?); + let mut preparation = std::pin::pin!(prepare_spawn( + &parent, + request, + path, + SpawnDestination::Detached, + )); + // Pause the fork as soon as it materializes the parent's rollout, then amend the policy. + let materialized = poll_fn(|cx| { + if rollout.exists() { + return Poll::Ready(()); + } + assert!( + preparation.as_mut().poll(cx).is_pending(), + "fork finished before materializing history" + ); + Poll::Pending + }); + tokio::time::timeout(Duration::from_secs(10), materialized).await?; + let command = vec!["late_fork_policy_test".to_owned()]; + let parent_policy = &session.services.exec_policy; + parent_policy + .append_amendment_and_update(home.path(), &ExecPolicyAmendment::new(command.clone())) + .await?; + let prepared = tokio::time::timeout(Duration::from_secs(10), preparation).await??; + let later = vec!["after_prepare_policy_test".to_owned()]; + parent_policy + .append_amendment_and_update(home.path(), &ExecPolicyAmendment::new(later.clone())) + .await?; + let prefixes = prepared + .inheritance + .exec_policy + .expect("detached policy") + .current() + .get_allowed_prefixes(); + assert!(prefixes.contains(&command)); + assert!(!prefixes.contains(&later)); + parent.shutdown_and_wait().await?; + Ok(()) +} diff --git a/codex-rs/core/src/agent/host/types.rs b/codex-rs/core/src/agent/host/types.rs new file mode 100644 index 0000000000..4707edf4ac --- /dev/null +++ b/codex-rs/core/src/agent/host/types.rs @@ -0,0 +1,50 @@ +//! Prepared inputs and host reservations for native child startup. + +use crate::agent::api::AgentControl; +use crate::agents_md_manager::SessionInstructions; +use crate::exec_policy::ExecPolicyManager; +use crate::session::session::AppServerClientMetadata; +use crate::thread_manager::StartThreadOptions; +use codex_protocol::AgentPath; +use codex_protocol::SessionId; +use codex_protocol::ThreadId; +use std::sync::Arc; + +/// A fresh child ID reserved in the host's tree. +/// The host commits membership, serializes startup attempts and deduplicates retries. +/// Core derives parent, session and path from [`PreparedThreadSpawn`]; the controller +/// supplies runtime coordination, and its `spawn` method is not called. +pub struct HostThreadSpawnReservation { + pub thread_id: ThreadId, + pub control: Arc, +} + +/// Opaque, in-process snapshot for starting a child without its source runtime. +/// Settings and authority reflect preparation time. The destination must supply compatible +/// services and register the captured environment IDs with matching working directories and +/// capability mappings. Detached preparation rejects thread attachments, live instruction providers, +/// enabled EMA-authenticated MCP servers, and unresolved environment configuration. +/// The child inherits user and thread instructions and discovers repository instructions +/// in its destination environment. +pub struct PreparedThreadSpawn { + pub(super) caller: ThreadId, + pub(super) session_id: SessionId, + pub(super) agent_path: AgentPath, + pub(super) client_metadata: AppServerClientMetadata, + pub(super) thread_options: StartThreadOptions, + pub(super) inheritance: PreparedChildInheritance, +} + +impl PreparedThreadSpawn { + /// Canonical path of the child in the agent tree. + pub fn agent_path(&self) -> &AgentPath { + &self.agent_path + } +} + +pub(crate) struct PreparedChildInheritance { + pub(crate) instructions: SessionInstructions, + pub(crate) exec_policy: Option>, + pub(crate) originator: String, + pub(crate) attachment_source: Option, +} diff --git a/codex-rs/core/src/agent/mod.rs b/codex-rs/core/src/agent/mod.rs index 7bf6b2480e..dd54f564bb 100644 --- a/codex-rs/core/src/agent/mod.rs +++ b/codex-rs/core/src/agent/mod.rs @@ -3,6 +3,7 @@ pub(crate) mod api; pub(crate) mod child_config; pub(crate) mod control; mod history; +pub(crate) mod host; mod registry; pub(crate) mod role; pub(crate) mod status; diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index 624b74aa76..c7e27be989 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -8,11 +8,15 @@ use crate::elicitation::ElicitationRegistration; use crate::session::SessionIo; use crate::session::SessionSettingsUpdate; use crate::session::Submission; +use crate::session::emit_subagent_session_started; use crate::session::new_submission_id; +use crate::session::session::AppServerClientMetadata; use crate::session::session::Session; use crate::session::startup_prewarm::PrewarmInput; use crate::session::step_settings::StepSettingsUpdate; use crate::thread_startup_metadata::ThreadStartupMetadata; +use chrono::DateTime; +use chrono::Utc; use codex_diagnostics::Gauge; use codex_diagnostics::GaugeGuard; use codex_exec_server::SelectedCapabilityRootsStatus; @@ -45,6 +49,7 @@ use codex_protocol::protocol::MultiAgentVersion; use codex_protocol::protocol::Op; use codex_protocol::protocol::SandboxPolicy; use codex_protocol::protocol::SessionSource; +use codex_protocol::protocol::SubAgentSource; use codex_protocol::protocol::ThreadHistoryMode; use codex_protocol::protocol::ThreadMemoryMode; use codex_protocol::protocol::ThreadSettingsOverrides; @@ -289,6 +294,39 @@ impl CodexThread { self.session.services.analytics_events_client.is_enabled() } + /// Looks up a live V2 tool invocation on a runtime with [`crate::HostToolInvocationLookup`]. + pub fn tool_invocation( + &self, + turn_id: &str, + call_id: &str, + ) -> CodexResult { + crate::agent::host::tool_invocation(self, turn_id, call_id) + } + + /// Creates an idle V2 child here with the invoking step's environments and live application + /// bindings. Compatible children share the parent's exec policy. + /// Ending the invocation or dropping this future aborts pending startup. + /// [`crate::ThreadManager::create_prepared_child`] defines acceptance and cleanup. + pub async fn create_child( + &self, + invocation: &crate::ToolInvocationContext, + request: crate::ThreadCreateRequest, + reservation: crate::HostThreadSpawnReservation, + ) -> Result { + crate::agent::host::create(self, invocation, request, reservation).await + } + + /// Captures inputs from a live invocation for [`crate::ThreadManager::create_prepared_child`]. + /// The host owns cancellation and snapshot freshness after this returns. Detached preparation + /// rejects managed application network, auth, and cloud config bindings. + pub async fn prepare_child( + &self, + invocation: &crate::ToolInvocationContext, + request: crate::ThreadCreateRequest, + ) -> Result { + crate::agent::host::prepare(self, invocation, request).await + } + /// Returns extension-owned data attached to this thread runtime. pub fn thread_extension_data(&self) -> &codex_extension_api::ExtensionData { &self.session.services.thread_extension_data @@ -338,6 +376,29 @@ impl CodexThread { self.io.session_loop_termination.clone().await; } + pub(crate) async fn emit_subagent_session_started( + &self, + client_metadata: AppServerClientMetadata, + resumed_created_at: Option>, + ) { + let SessionSource::SubAgent(source @ SubAgentSource::ThreadSpawn { .. }) = + &self.session_source + else { + return; + }; + let config = self.config_snapshot().await; + emit_subagent_session_started( + &self.session.services.analytics_events_client, + client_metadata, + self.session.session_id(), + self.session.thread_id(), + config.parent_thread_id, + config, + source.clone(), + resumed_created_at, + ); + } + pub(crate) async fn emit_thread_ready_lifecycle(&self) { let contributors = self .session diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index 2c91c8e487..7f3363ada3 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -152,6 +152,13 @@ mod event_mapping; pub use codex_prompts as review_prompts; mod thread_manager; mod thread_spawn; +pub use agent::host::HostThreadSpawnReservation; +pub use agent::host::HostToolInvocationLookup; +pub use agent::host::PreparedThreadSpawn; +pub use agent::host::ToolInvocationContext; +pub use thread_spawn::ThreadCreateRequest; +pub use thread_spawn::ThreadSpawnError; +pub use thread_spawn::ThreadSpawnResult; pub(crate) mod web_search; pub(crate) mod windows_sandbox_read_grants; pub use agent::control::ThreadEvictionOutcome; diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index c62eeb70d8..f946d5707d 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -356,7 +356,14 @@ impl StartThreadOptions { } } -/// Parent dependencies resolved before the shared session startup. +#[derive(Clone, Copy)] +enum ThreadReadyNotification { + Emit, + // Caller installs cleanup before running ready callbacks. + Defer, +} + +/// Parent dependencies resolved from live Local state or a prepared snapshot. struct ThreadStartupContext { history_initialization: HistoryInitialization, instructions: SessionInstructions, @@ -367,13 +374,21 @@ struct ThreadStartupContext { attachment_source: Option, } +struct PreparedChildOrigin { + history_initialization: HistoryInitialization, + originator: String, + attachment_source: Option, +} + struct ThreadSpawnRequest { + ready_notification: ThreadReadyNotification, startup: Option>, options: StartThreadOptions, auth_manager: Arc, agent_control: AgentControlInit, parent_thread_id: Option, parent_originator: Option, + prepared_child: Option, forked_from_thread_id: Option, fork_persistence: ForkPersistence, inherited_environments: Option, @@ -389,12 +404,14 @@ impl ThreadSpawnRequest { agent_control: impl Into, ) -> Self { Self { + ready_notification: ThreadReadyNotification::Emit, startup: None, options, auth_manager, agent_control: agent_control.into(), parent_thread_id: None, parent_originator: None, + prepared_child: None, forked_from_thread_id: None, fork_persistence: ForkPersistence::Copied, inherited_environments: None, @@ -1139,6 +1156,27 @@ impl ThreadManager { .await } + /// Creates a [prepared child](crate::PreparedThreadSpawn) using this manager's services. + /// + /// On `Ok`, the child is loaded and idle, its creation has been published, and the host + /// owns it. The host chooses whether to submit work and handles any later delivery failure. + /// + /// An error or a dropped future schedules asynchronous cleanup of any provisional child. + /// Already-durable history may remain. + pub async fn create_prepared_child( + &self, + prepared: crate::agent::host::PreparedThreadSpawn, + reservation: crate::agent::host::HostThreadSpawnReservation, + ) -> Result { + crate::agent::host::start_prepared( + &self.agent_control().runtime, + prepared, + reservation, + /*invocation*/ None, + ) + .await + } + /// Starts a fresh internal session associated with an existing parent thread. pub async fn spawn_internal_session( &self, @@ -2059,6 +2097,52 @@ impl ThreadManagerState { Box::pin(self.spawn_thread(request)).await } + /// Starts from captured inheritance without requiring a live parent runtime. + pub(crate) async fn spawn_prepared_child_thread( + &self, + options: StartThreadOptions, + agent_control: AgentControlInit, + parent_thread_id: ThreadId, + inheritance: crate::agent::host::PreparedChildInheritance, + ) -> CodexResult { + if let Some(environments) = &options.environments { + crate::environment_selection::validate_environment_ids_and_cwds( + &self.environment_manager, + environments + .iter() + .map(|environment| (environment.environment_id.as_str(), &environment.cwd)), + )?; + } + options.reserved_thread_id.ok_or_else(|| { + CodexErr::InvalidRequest("prepared child startup requires a reserved thread ID".into()) + })?; + let mut request = + ThreadSpawnRequest::new(options, Arc::clone(&self.auth_manager), agent_control); + request.ready_notification = ThreadReadyNotification::Defer; + request.parent_thread_id = Some(parent_thread_id); + let history_initialization = match &request.options.initial_history { + InitialHistory::Forked(_) => { + request.forked_from_thread_id = Some(parent_thread_id); + request.fork_persistence = ForkPersistence::CopiedDeferred; + HistoryInitialization::WarmFork + } + InitialHistory::New => HistoryInitialization::New, + InitialHistory::Resumed(_) | InitialHistory::Cleared => { + return Err(CodexErr::InvalidRequest( + "prepared spawn requires new or forked history".into(), + )); + } + }; + request.inherited_instructions = Some(inheritance.instructions); + request.inherited_exec_policy = inheritance.exec_policy; + request.prepared_child = Some(PreparedChildOrigin { + history_initialization, + originator: inheritance.originator, + attachment_source: inheritance.attachment_source, + }); + Box::pin(self.spawn_thread(request)).await + } + pub(crate) async fn resume_thread_with_history_with_source( &self, options: ResumeThreadWithHistoryOptions, @@ -2151,12 +2235,14 @@ impl ThreadManagerState { let startup_guard = owns_startup.then(|| SessionStartupGuard::new(Arc::clone(&startup_state))); let ThreadSpawnRequest { + ready_notification, startup, options, auth_manager, agent_control, parent_thread_id, parent_originator, + prepared_child, forked_from_thread_id, fork_persistence, inherited_environments, @@ -2164,7 +2250,9 @@ impl ThreadManagerState { inherited_exec_policy, user_shell_override, } = request; - let history_initialization = if matches!( + let history_initialization = if let Some(child) = &prepared_child { + child.history_initialization + } else if matches!( &options.initial_history, InitialHistory::New | InitialHistory::Forked(_) ) && let Some(source_thread_id) = forked_from_thread_id @@ -2345,6 +2433,9 @@ impl ThreadManagerState { } let (exec_policy, multi_agent_version) = match isolation { SessionIsolation::Isolated => (None, Some(MultiAgentVersion::Disabled)), + SessionIsolation::Inherit if prepared_child.is_some() => { + (inherited_exec_policy, Some(MultiAgentVersion::V2)) + } SessionIsolation::Inherit => ( inherited_exec_policy, self.initial_multi_agent_version_for_spawn( @@ -2356,16 +2447,10 @@ impl ThreadManagerState { .await, ), }; - let startup_context = ThreadStartupContext { - history_initialization, - instructions, - exec_policy, - multi_agent_version, - parent_rollout_thread_trace: self - .parent_rollout_thread_trace_for_source(&session_source, &initial_history) - .await, - originator: self - .effective_originator( + let (originator, attachment_source) = match prepared_child { + Some(child) => (child.originator, child.attachment_source), + None => ( + self.effective_originator( &initial_history, metrics_service_name.as_deref(), &session_source, @@ -2374,8 +2459,20 @@ impl ThreadManagerState { parent_originator, ) .await, - attachment_source: forked_from_thread_id - .filter(|_| matches!(&initial_history, InitialHistory::Forked(_))), + forked_from_thread_id + .filter(|_| matches!(&initial_history, InitialHistory::Forked(_))), + ), + }; + let startup_context = ThreadStartupContext { + history_initialization, + instructions, + exec_policy, + multi_agent_version, + parent_rollout_thread_trace: self + .parent_rollout_thread_trace_for_source(&session_source, &initial_history) + .await, + originator, + attachment_source, }; let (extensions, mcp_manager) = if isolation == SessionIsolation::Isolated { let extensions = thread_extension_init @@ -2526,7 +2623,9 @@ impl ThreadManagerState { return Err(error); } }; - new_thread.thread.emit_thread_ready_lifecycle().await; + if matches!(ready_notification, ThreadReadyNotification::Emit) { + new_thread.thread.emit_thread_ready_lifecycle().await; + } if source_changed_during_startup.load(Ordering::Acquire) { new_thread.thread.session.request_mcp_runtime_refresh(); } @@ -2622,13 +2721,12 @@ impl ThreadManagerState { if matches!(initial_history, InitialHistory::Resumed(_)) { return codex_rollout_trace::ThreadTraceContext::disabled(); } - // Parent lookup can fail if the parent was closed or released between - // spawn preparation and session construction. Tracing is diagnostic, so - // that race should not block child creation; the child simply starts - // without a parent rollout trace. + // A stopped parent or one on another manager cannot supply a local trace. + // Tracing is best-effort and never blocks child creation. self.get_thread(*parent_thread_id) .await .ok() + .filter(|thread| thread.is_running()) .map(|thread| thread.session.services.rollout_thread_trace.clone()) .unwrap_or_else(codex_rollout_trace::ThreadTraceContext::disabled) } diff --git a/codex-rs/core/src/thread_spawn.rs b/codex-rs/core/src/thread_spawn.rs index 63f2bb014a..4623a50c08 100644 --- a/codex-rs/core/src/thread_spawn.rs +++ b/codex-rs/core/src/thread_spawn.rs @@ -14,6 +14,7 @@ use crate::agent::types::MessageDeliveryMode; use crate::agent::types::SpawnAgentForkMode; use crate::agent::types::SpawnAgentOptions; use crate::codex_thread::ThreadConfigSnapshot; +use crate::config::Config; use crate::session::multi_agents::resolve_usage_hints; use crate::session::session::Session; use crate::session::step_context::StepContext; @@ -26,6 +27,15 @@ use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::SubAgentSource; /// Child creation options; omitted overrides retain native MAv2 inheritance and role rules. +pub struct ThreadCreateRequest { + pub task_name: String, + /// `None` starts fresh; `FullHistory` uses native fork preparation. + pub fork_mode: Option, + pub agent_type: Option, + pub model: Option, + pub reasoning_effort: Option, +} + pub(crate) struct ThreadSpawnRequest { pub task_name: String, /// Explicit encoding: hosts normally provide plaintext; native tools can forward ciphertext. @@ -37,8 +47,39 @@ pub(crate) struct ThreadSpawnRequest { pub reasoning_effort: Option, } -/// Accepted child and its effective configuration. Acceptance does not mean the turn completed. -pub(crate) struct ThreadSpawnResult { +impl ThreadSpawnRequest { + pub(crate) fn into_parts( + self, + ) -> Result<(ThreadCreateRequest, AgentMessage), ThreadSpawnError> { + let message_text = match &self.message { + AgentMessage::Plaintext(text) | AgentMessage::Encrypted(text) => text, + }; + if message_text.trim().is_empty() { + return Err(ThreadSpawnError::Preparation( + "Empty message can't be sent to an agent".to_string(), + )); + } + Ok(( + ThreadCreateRequest { + task_name: self.task_name, + fork_mode: self.fork_mode, + agent_type: self.agent_type, + model: self.model, + reasoning_effort: self.reasoning_effort, + }, + self.message, + )) + } +} + +pub(crate) struct ChildSpawnSettings { + pub(crate) config: Config, + pub(crate) source: SessionSource, + pub(crate) options: SpawnAgentOptions, +} + +/// Child runtime and its effective configuration. +pub struct ThreadSpawnResult { pub agent: LiveAgent, pub config: ThreadConfigSnapshot, pub task_name: AgentPath, @@ -46,7 +87,7 @@ pub(crate) struct ThreadSpawnResult { /// Separates preparation errors from backend failures to preserve native tool error reporting. #[derive(Debug, thiserror::Error)] -pub(crate) enum ThreadSpawnError { +pub enum ThreadSpawnError { #[error("{0}")] Preparation(String), #[error(transparent)] @@ -59,7 +100,18 @@ pub(crate) async fn spawn( call_id: &str, request: ThreadSpawnRequest, ) -> Result { - let (request, task_name) = prepare(session, step_context, call_id, request).await?; + let (request, message) = request.into_parts()?; + let (settings, task_name) = prepare(session, step_context, call_id, request).await?; + let request = SpawnRequest { + caller: session.thread_id, + config: settings.config, + input: AgentInput::Message { + message, + mode: MessageDeliveryMode::TriggerTurn, + }, + source: settings.source, + options: settings.options, + }; let (agent, config) = session.services.agent_control.spawn(request).await?; Ok(ThreadSpawnResult { agent, @@ -73,25 +125,16 @@ pub(crate) async fn prepare( session: &Session, step_context: &StepContext, call_id: &str, - request: ThreadSpawnRequest, -) -> Result<(SpawnRequest, AgentPath), ThreadSpawnError> { + request: ThreadCreateRequest, +) -> Result<(ChildSpawnSettings, AgentPath), ThreadSpawnError> { let turn = &step_context.turn; - let ThreadSpawnRequest { + let ThreadCreateRequest { task_name, - message, fork_mode, agent_type, model, reasoning_effort, } = request; - let message_text = match &message { - AgentMessage::Plaintext(text) | AgentMessage::Encrypted(text) => text, - }; - if message_text.trim().is_empty() { - return Err(ThreadSpawnError::Preparation( - "Empty message can't be sent to an agent".to_string(), - )); - } let role_name = agent_type .as_deref() .map(str::trim) @@ -149,13 +192,8 @@ pub(crate) async fn prepare( None }; Ok(( - SpawnRequest { - caller: session.thread_id, + ChildSpawnSettings { config, - input: AgentInput::Message { - message, - mode: MessageDeliveryMode::TriggerTurn, - }, source: spawn_source, options: SpawnAgentOptions { fork_parent_spawn_call_id: fork_mode.as_ref().map(|_| call_id.to_string()), diff --git a/codex-rs/core/src/tools/context.rs b/codex-rs/core/src/tools/context.rs index bb9fd88a77..fe7de2c90b 100644 --- a/codex-rs/core/src/tools/context.rs +++ b/codex-rs/core/src/tools/context.rs @@ -48,6 +48,7 @@ pub type SharedTurnDiffTracker = Arc>; pub(crate) struct ToolCallState { pub(crate) terminal_outcome_reached: AtomicBool, pub(crate) delivered_assistant_message: OnceLock, + pub(crate) invocation: crate::agent::host::InvocationReservation, } #[derive(Clone, Debug, Eq, PartialEq)] diff --git a/codex-rs/core/src/tools/parallel.rs b/codex-rs/core/src/tools/parallel.rs index 09fd9000df..ffb17eb8a7 100644 --- a/codex-rs/core/src/tools/parallel.rs +++ b/codex-rs/core/src/tools/parallel.rs @@ -160,6 +160,13 @@ impl ToolCallRuntime { cancellation_token: CancellationToken, call_state: Arc, ) -> impl std::future::Future> { + if let Err(err) = + call_state + .invocation + .reserve(&self.session, &step_context.turn, &call.call_id) + { + return Either::Left(std::future::ready(Err(err))); + } let message_admission = super::user_messaging::admit_code_mode_send(&self.session, &source, &call.tool_name); self.session @@ -320,7 +327,7 @@ impl ToolCallRuntime { tracker, dispatch_call, source, - dispatch_call_state, + Arc::clone(&dispatch_call_state), ) .await; drop(_admitted_timing); @@ -339,22 +346,34 @@ impl ToolCallRuntime { started.elapsed(), ); } + // Finished calls release their IDs even if their results have not been collected. + // Otherwise the outer task may still need the ID for its abort callback. + if dispatch_call_state + .terminal_outcome_reached + .load(Ordering::Acquire) + { + drop(dispatch_call_state.invocation.take()); + } result } .instrument(dispatch_span.clone().or_current()), )); - async move { + let completion = async move { let mut tool_call_timing_guard = tool_call_timing_guard; + let join_error = |err| { + drop(call_state.invocation.take()); + Self::tool_task_join_error(err) + }; let mut result = tokio::select! { - res = &mut dispatch_handle => res.map_err(Self::tool_task_join_error)?, + res = &mut dispatch_handle => res.map_err(join_error)?, _ = cancellation_token.cancelled() => { // Dynamic handlers observe the token and publish their terminal item themselves. if finishes_on_cancellation || call_state.terminal_outcome_reached.load(Ordering::Acquire) || dispatch_handle.is_finished() { - dispatch_handle.await.map_err(Self::tool_task_join_error)? + dispatch_handle.await.map_err(join_error)? } else { let secs = started.elapsed().as_secs_f32().max(0.1); abort_dispatch_span.record("aborted", true); @@ -362,6 +381,7 @@ impl ToolCallRuntime { match dispatch_handle.await { Ok(result) => result, Err(err) if err.is_cancelled() => { + let _reservation = call_state.invocation.take(); let response = Self::aborted_response(&call, Some(secs)); call_trace::result_ready( thread_id, @@ -387,6 +407,7 @@ impl ToolCallRuntime { } }, }; + let _reservation = call_state.invocation.take(); // Cancellation outside the handler can return before the registry reports an outcome. if cancellation_token.is_cancelled() && !call_state @@ -416,7 +437,8 @@ impl ToolCallRuntime { } result } - .instrument(dispatch_span.or_current()) + .instrument(dispatch_span.or_current()); + Either::Right(completion) } } diff --git a/codex-rs/core/src/tools/registry.rs b/codex-rs/core/src/tools/registry.rs index 6e3cce6d37..3dccc2eece 100644 --- a/codex-rs/core/src/tools/registry.rs +++ b/codex-rs/core/src/tools/registry.rs @@ -711,6 +711,15 @@ impl ToolRegistry { mut invocation: ToolInvocation, call_state: Option>, ) -> Result { + // Standalone dispatch has no outer task to reserve the ID for it. + let standalone; + let reservation = if let Some(state) = &call_state { + &state.invocation + } else { + standalone = crate::agent::host::InvocationReservation::default(); + standalone.reserve(&invocation.session, &invocation.turn, &invocation.call_id)?; + &standalone + }; let cancellation_token = invocation.cancellation_token.clone(); let tool_name = invocation.tool_name.clone(); let call_id_owned = invocation.call_id.clone(); @@ -879,6 +888,7 @@ impl ToolRegistry { return Err(fail_before_execution(&invocation, err).await); } + let invocation_guard = reservation.activate(&invocation); if tool.mcp_server_name().is_none() { notify_tool_start(&invocation, /*mcp_tool*/ None) .or_cancel_tool(&cancellation_token) @@ -978,6 +988,7 @@ impl ToolRegistry { handler_executed: true, }, }; + drop(invocation_guard); notify_tool_finish_if_unclaimed(&invocation, call_state.as_deref(), lifecycle_outcome) .await; diff --git a/codex-rs/core/tests/common/lib.rs b/codex-rs/core/tests/common/lib.rs index 77758d0408..91848e8b9a 100644 --- a/codex-rs/core/tests/common/lib.rs +++ b/codex-rs/core/tests/common/lib.rs @@ -273,6 +273,16 @@ where wait_for_event_with_timeout(codex, predicate, Duration::from_secs(1)).await } +pub async fn wait_for_successful_turn(codex: &CodexThread) { + use codex_protocol::protocol::EventMsg; + + let event = wait_for_event(codex, |event| matches!(event, EventMsg::TurnComplete(_))).await; + assert!( + matches!(&event, EventMsg::TurnComplete(completed) if completed.error.is_none()), + "turn failed: {event:?}" + ); +} + /// Waits for a configured MCP server to finish startup and requires it to be ready. pub async fn wait_for_mcp_server(codex: &CodexThread, server_name: &str) -> anyhow::Result<()> { use codex_protocol::protocol::EventMsg; diff --git a/codex-rs/core/tests/common/responses.rs b/codex-rs/core/tests/common/responses.rs index 25e0078f39..672d960660 100644 --- a/codex-rs/core/tests/common/responses.rs +++ b/codex-rs/core/tests/common/responses.rs @@ -105,6 +105,19 @@ fn assert_turn_id(body: &Value, key: &str, expected: Option<&str>) -> Result<()> #[derive(Debug, Clone)] pub struct ResponsesRequest(wiremock::Request); +pub async fn received_analytics_events(server: &MockServer, event_type: &str) -> Vec { + server + .received_requests() + .await + .unwrap_or_default() + .into_iter() + .filter(|request| request.url.path() == "/codex/analytics-events/events") + .filter_map(|request| serde_json::from_slice::(&request.body).ok()) + .flat_map(|payload| payload["events"].as_array().cloned().unwrap_or_default()) + .filter(|event| event["event_type"] == event_type) + .collect() +} + pub async fn received_responses_requests(server: &MockServer) -> Vec { server .received_requests() diff --git a/codex-rs/core/tests/suite/agent_control_tests.rs b/codex-rs/core/tests/suite/agent_control_tests.rs index eb37ca5448..a4869bd1bf 100644 --- a/codex-rs/core/tests/suite/agent_control_tests.rs +++ b/codex-rs/core/tests/suite/agent_control_tests.rs @@ -51,7 +51,7 @@ use serde_json::json; use test_case::test_case; use wiremock::MockServer; -struct TestAgentControl { +pub(super) struct TestAgentControl { thread_id: ThreadId, service_tier: Mutex>, mail: Mutex>, @@ -60,7 +60,7 @@ struct TestAgentControl { } impl TestAgentControl { - fn new(thread_id: ThreadId) -> Self { + pub(super) fn new(thread_id: ThreadId) -> Self { Self { thread_id, service_tier: Mutex::default(), diff --git a/codex-rs/core/tests/suite/pending_input.rs b/codex-rs/core/tests/suite/pending_input.rs index d7ce6b5af9..2aa2c35ec3 100644 --- a/codex-rs/core/tests/suite/pending_input.rs +++ b/codex-rs/core/tests/suite/pending_input.rs @@ -57,6 +57,7 @@ use core_test_support::test_codex::TestCodex; use core_test_support::test_codex::test_codex; use core_test_support::test_codex::turn_permission_fields; use core_test_support::wait_for_event; +use core_test_support::wait_for_successful_turn; use pretty_assertions::assert_eq; use serde_json::Value; use serde_json::from_slice; @@ -439,14 +440,6 @@ async fn wait_for_turn_complete(codex: &CodexThread) { wait_for_event(codex, |event| matches!(event, EventMsg::TurnComplete(_))).await; } -async fn wait_for_successful_turn(codex: &CodexThread) { - let event = wait_for_event(codex, |event| matches!(event, EventMsg::TurnComplete(_))).await; - assert!( - matches!(&event, EventMsg::TurnComplete(completed) if completed.error.is_none()), - "turn failed: {event:?}" - ); -} - async fn wait_for_sleep_item_started(codex: &CodexThread, call_id: &str, duration_ms: u64) { let event = wait_for_event(codex, |event| { matches!( diff --git a/codex-rs/core/tests/suite/scenarios.rs b/codex-rs/core/tests/suite/scenarios.rs index 90df49de18..ace972ff0f 100644 --- a/codex-rs/core/tests/suite/scenarios.rs +++ b/codex-rs/core/tests/suite/scenarios.rs @@ -141,6 +141,9 @@ mod mailbox_preemption; #[path = "scenarios_agent_eviction_tests.rs"] mod agent_eviction; +#[path = "scenarios_reserved_spawn.rs"] +mod reserved_spawn; + #[path = "scenarios_partial_answers.rs"] mod partial_answers; diff --git a/codex-rs/core/tests/suite/scenarios_reserved_spawn.rs b/codex-rs/core/tests/suite/scenarios_reserved_spawn.rs new file mode 100644 index 0000000000..8083fb22b5 --- /dev/null +++ b/codex-rs/core/tests/suite/scenarios_reserved_spawn.rs @@ -0,0 +1,642 @@ +//! Exercises host-reserved child startup with fresh and inherited history. + +use super::super::agent_control::TestAgentControl; +use super::super::agents_md::RecordingThreadInstructionsProvider; +use super::*; +use anyhow::Context; +use codex_analytics::AnalyticsEventsClient; +use codex_core::AgentControl; +use codex_core::CodexThread; +use codex_core::HostThreadSpawnReservation; +use codex_core::HostToolInvocationLookup; +use codex_core::PromptCacheKey; +use codex_core::SpawnAgentForkMode; +use codex_core::StartIfIdleSubmission; +use codex_core::ThreadCreateRequest; +use codex_core::ThreadSpawnError; +use codex_extension_api::ExtensionFuture; +use codex_extension_api::Instructions; +use codex_extension_api::ThreadInstructionsProvider; +use codex_extension_api::ThreadLifecycleContributor; +use codex_extension_api::ThreadReadyInput; +use codex_http_client::DestinationPolicy; +use codex_http_client::NetworkPolicy; +use codex_http_client::NetworkPolicyController; +use codex_protocol::ThreadId; +use codex_protocol::dynamic_tools::DynamicToolCallRequest; +use codex_protocol::error::CodexErrorDetails; +use codex_protocol::models::PermissionProfile; +use codex_protocol::models::input_text; +use codex_protocol::protocol::ThreadHistoryMode; +use codex_protocol::protocol::ThreadSettingsOverrides; +use codex_thread_store::AddThreadAttachmentParams; +use codex_thread_store::ListThreadAttachmentsParams; +use codex_thread_store::LoadThreadHistoryParams; +use core_test_support::responses; +use core_test_support::submit_thread_settings; +use core_test_support::test_codex::TestCodexBuilder; +use core_test_support::wait_for_successful_turn; +use pretty_assertions::assert_eq; +use tokio::sync::Notify; +use wiremock::MockServer; + +async fn host_spawn_fixture( + server: &MockServer, + builder: TestCodexBuilder, + application_network_policy: NetworkPolicy, + prompt_cache_key: PromptCacheKey, + thread_instructions_provider: Option>, +) -> Result<(TestCodex, Arc, Arc)> { + let root_id = ThreadId::new(); + // This controller rejects every AgentControl::spawn call. Successful startup therefore + // demonstrates the direct host hook, while runtime calls still use the host controller. + let controller = Arc::new(TestAgentControl::new(root_id)); + let mut test = builder + .with_model("gpt-6-astra") + .with_config(move |config| { + configure_scenario_catalog(config); + config.application_network_policy = application_network_policy; + config + .features + .enable(Feature::MultiAgentV2) + .expect("enable multi-agent V2"); + config + .features + .enable(Feature::MultiAgentV2DynamicTools) + .expect("enable dynamic tool inheritance"); + config.multi_agent_v2.disable_native_tools = true; + config + .permissions + .set_permission_profile(PermissionProfile::read_only()) + .expect("configure read-only permissions"); + }) + .with_thread_manager({ + let controller = Arc::clone(&controller); + move |manager| { + manager.with_agent_control_factory(move |thread_id| { + let control: Arc = if thread_id == root_id { + controller.clone() + } else { + Arc::new(TestAgentControl::new(thread_id)) + }; + async move { Ok(control) } + }) + } + }) + .build_with_auto_env(server) + .await?; + let mut options = test.start_thread_options().await; + options + .thread_extension_init + .insert(HostToolInvocationLookup); + options.history_mode = Some(ThreadHistoryMode::Legacy); + options.prompt_cache_key = prompt_cache_key; + options.thread_instructions_provider = thread_instructions_provider; + options.reserved_thread_id = Some(root_id); + options.disabled_plugin_ids = Some(vec!["captured-plugin@local".into()]); + options.dynamic_tools = vec![DynamicToolSpec::Function(DynamicToolFunctionSpec { + name: "host_spawn".into(), + description: "Ask the host to start a child from this invocation.".into(), + input_schema: json!({"type": "object", "properties": {}}), + defer_loading: false, + })]; + test.codex.shutdown_and_wait().await?; + test.thread_manager + .remove_thread(&test.session_configured.thread_id) + .await; + let root = test.thread_manager.start_thread(options).await?; + test.codex = Arc::clone(&root.thread); + test.session_configured = root.session_configured; + let parent = root.thread; + Ok((test, parent, controller)) +} + +async fn host_spawn_invocation( + server: &MockServer, + parent: &Arc, + call_id: &str, + message: &str, +) -> Result { + responses::mount_sse_once( + server, + sse(vec![ + responses::ev_function_call(call_id, "host_spawn", "{}"), + ev_completed(call_id), + ]), + ) + .await; + parent + .start_or_steer_turn(TurnInputRequest::user_input(vec![text(message)])) + .await?; + Ok(wait_for_event_match(parent, |event| match event { + EventMsg::DynamicToolCallRequest(call) => Some(call.clone()), + _ => None, + }) + .await) +} + +async fn finish_spawn_tool( + server: &MockServer, + parent: &CodexThread, + invocation: &DynamicToolCallRequest, +) -> Result<()> { + responses::mount_sse_once(server, sse(vec![ev_completed("parent-done")])).await; + parent + .submit(Op::DynamicToolResponse { + id: invocation.call_id.clone(), + response: DynamicToolResponse { + content_items: vec![DynamicToolCallOutputContentItem::InputText { + text: "Done.".into(), + }], + success: true, + }, + }) + .await?; + wait_for_event(parent, |event| matches!(event, EventMsg::TurnComplete(_))).await; + assert!( + parent + .tool_invocation(&invocation.turn_id, &invocation.call_id) + .is_err() + ); + Ok(()) +} + +fn create_request(name: &str, fork_history: bool) -> ThreadCreateRequest { + ThreadCreateRequest { + task_name: name.into(), + fork_mode: fork_history.then_some(SpawnAgentForkMode::FullHistory), + agent_type: None, + model: None, + reasoning_effort: None, + } +} + +async fn start_first_turn(child: &CodexThread, prompt: &str) -> Result<()> { + let submission = child + .start_turn_if_idle(TurnInputRequest::user_input(vec![text(prompt)])) + .await?; + assert!( + matches!(submission, StartIfIdleSubmission::Started { .. }), + "expected the new child to be idle, got {submission:?}" + ); + wait_for_successful_turn(child).await; + Ok(()) +} + +#[test_case::test_case(false; "fresh")] +#[test_case::test_case(true; "forked")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn same_manager_creation_preserves_live_bindings_and_attachments( + fork_history: bool, +) -> Result<()> { + const PR: &str = "openai/codex#123"; + const UPDATED_RULE: &str = "Updated shared worker instructions."; + skip_if_no_network!(Ok(())); + let server = start_mock_server().await; + let network_controller = NetworkPolicyController::default(); + let network_policy = network_controller.policy(); + assert!(network_controller.publish(network_policy.revision(), DestinationPolicy::Unrestricted)); + let provider = + Arc::new(RecordingThreadInstructionsProvider::with_text("Initial rules.").shared()); + let (test, parent, controller) = host_spawn_fixture( + &server, + test_codex(), + network_policy.clone(), + PromptCacheKey::SessionDerived, + Some(provider.clone()), + ) + .await?; + let root_id = test.session_configured.thread_id; + let captured = parent.config_snapshot().await; + let invocation = + host_spawn_invocation(&server, &parent, "spawn-worker", "Start the worker.").await?; + let context = parent.tool_invocation(&invocation.turn_id, &invocation.call_id)?; + test.thread_store + .add_thread_attachment(AddThreadAttachmentParams { + thread_id: root_id, + attachment_type: "pull_request".into(), + identity_key: PR.into(), + payload: json!({"number": 123}), + }) + .await?; + // Changing current thread settings must not change the invocation's captured settings. + submit_thread_settings( + &parent, + ThreadSettingsOverrides { + disabled_plugin_ids: Some(vec!["later-plugin@local".into()]), + ..Default::default() + }, + ) + .await?; + let child_id = ThreadId::new(); + let request = || create_request("worker", fork_history); + let reservation = || HostThreadSpawnReservation { + thread_id: child_id, + control: controller.clone(), + }; + // A detached snapshot cannot carry the source's live application policy. + assert!(matches!( + parent + .prepare_child(&context, request()) + .await, + Err(ThreadSpawnError::Backend(error)) + if error.to_string().contains("host-managed application authority") + )); + let mut mismatched = reservation(); + mismatched.control = Arc::new(TestAgentControl::new(ThreadId::new())); + assert!(matches!( + parent.create_child(&context, request(), mismatched).await, + Err(ThreadSpawnError::Preparation(_)) + )); + assert!(test.thread_manager.get_thread(child_id).await.is_err()); + + provider.set_instructions(Some(Instructions { + text: UPDATED_RULE.into(), + source: None, + })); + let spawned = parent + .create_child(&context, request(), reservation()) + .await?; + assert_eq!(spawned.config.approval_policy, captured.approval_policy); + assert_eq!(spawned.config.environments, captured.environments); + assert_eq!( + spawned.config.disabled_plugin_ids, + captured.disabled_plugin_ids + ); + let worker = test.thread_manager.get_thread(child_id).await?; + // A handle cannot be used through a different source runtime. + assert!(matches!( + worker.prepare_child(&context, request()).await, + Err(ThreadSpawnError::Backend(error)) + if error.to_string().contains("tool invocation belongs to a different runtime") + )); + let child_network_policy = worker.config().await.application_network_policy.clone(); + assert_eq!(child_network_policy, network_policy); + assert_eq!( + worker.rollout_path().context("child rollout")?.is_file(), + fork_history, + "only copied history should be materialized at creation" + ); + // Once creation returns, ending the spawning invocation leaves the idle child available. + finish_spawn_tool(&server, &parent, &invocation).await?; + let child_request = + responses::mount_sse_once(&server, responses::sse_completed("worker-done")).await; + start_first_turn(&worker, "Report your status.").await?; + assert!( + child_request + .single_request() + .body_contains_text(UPDATED_RULE) + ); + let page = test + .thread_store + .list_thread_attachments(ListThreadAttachmentsParams { + thread_id: child_id, + cursor: None, + limit: 1, + }) + .await?; + assert_eq!( + page.attachments + .first() + .map(|attachment| attachment.identity_key.as_str()), + fork_history.then_some(PR) + ); + + worker.shutdown_and_wait().await?; + parent.shutdown_and_wait().await?; + Ok(()) +} + +#[test_case::test_case(false, "inherited instruction providers"; "thread instructions")] +#[test_case::test_case(true, "EMA-authenticated MCP servers"; "enterprise MCP")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn prepare_child_rejects_live_bindings(ema_enabled: bool, expected: &str) -> Result<()> { + skip_if_no_network!(Ok(())); + let server = start_mock_server().await; + let home = Arc::new(TempDir::new()?); + std::fs::write( + home.path().join("config.toml"), + format!( + r#"features.use_xaa = true +mcp_enterprise_managed_auth.idp = {{ issuer = "https://idp.example", client_id = "test" }} +mcp_servers.enterprise = {{ url = "https://mcp.example/mcp", auth = "ema_auth", enabled = {ema_enabled} }}"# + ), + )?; + let (test, parent, controller) = host_spawn_fixture( + &server, + test_codex().with_home(home), + NetworkPolicy::unmanaged(), + PromptCacheKey::SessionDerived, + (!ema_enabled).then(|| { + Arc::new(RecordingThreadInstructionsProvider::with_text("Shared rules").shared()) + as Arc + }), + ) + .await?; + let call = host_spawn_invocation(&server, &parent, "prepare-worker", "Start a worker.").await?; + let context = parent.tool_invocation(&call.turn_id, &call.call_id)?; + let request = || create_request("worker", /*fork_history*/ false); + let Err(ThreadSpawnError::Backend(error)) = parent.prepare_child(&context, request()).await + else { + anyhow::bail!("expected detached preparation to reject {expected}"); + }; + assert!(matches!( + error.details(), + CodexErrorDetails::InvalidRequest(message) if message.contains(expected) + )); + let child_id = ThreadId::new(); + parent + .create_child( + &context, + request(), + HostThreadSpawnReservation { + thread_id: child_id, + control: controller, + }, + ) + .await?; + test.thread_manager + .get_thread(child_id) + .await? + .shutdown_and_wait() + .await?; + parent.shutdown_and_wait().await?; + Ok(()) +} + +#[test_case::test_case(true, None; "full_history")] +#[test_case::test_case(false, Some("host-cache-key"); "fresh_history")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn prepared_spawn_starts_on_an_independent_destination_after_parent_shutdown( + fork_history: bool, + cache_key: Option<&str>, +) -> Result<()> { + const PROJECT_RULE: &str = "Run the destination repository checks."; + const EARLIER_INPUT: &str = "Remember PRIOR_CONTEXT_5241."; + const SPAWN_INPUT: &str = + "Remember TRANSFERRED_CONTEXT_7391 and ask the host to start a worker."; + const TASK: &str = "Report TRANSFERRED_TASK_8812 and the remembered context."; + skip_if_no_network!(Ok(())); + let server = start_mock_server().await; + let (source, parent, source_control) = host_spawn_fixture( + &server, + test_codex(), + NetworkPolicy::unmanaged(), + cache_key.map_or(PromptCacheKey::Omit, |key| { + PromptCacheKey::Custom(key.into()) + }), + /*thread_instructions_provider*/ None, + ) + .await?; + let root_id = source.session_configured.thread_id; + let session_id = source.session_configured.session_id; + parent + .set_app_server_client_info( + Some("source-client".into()), + Some("1.0".into()), + /*mcp_elicitations_auto_deny*/ false, + ) + .await?; + let captured = parent.config_snapshot().await; + responses::mount_sse_once(&server, sse(vec![ev_completed("remembered")])).await; + source.submit_text_turn(EARLIER_INPUT).await?; + let invocation = host_spawn_invocation(&server, &parent, "prepare-worker", SPAWN_INPUT).await?; + let context = parent.tool_invocation(&invocation.turn_id, &invocation.call_id)?; + let prepared = parent + .prepare_child(&context, create_request("remote_worker", fork_history)) + .await?; + let child_path = prepared.agent_path().clone(); + assert_eq!(child_path.as_str(), "/root/remote_worker"); + let agents_md = source.workspace_path_uri("AGENTS.md")?; + + // Keep only durable storage and referenced paths; startup cannot consult A's runtime, + // manager, controller, or environment services after this point. + let source_store = Arc::clone(&source.thread_store); + let _source_home = Arc::clone(&source.home); + let _source_cwd = Arc::clone(&source.cwd); + let source_manager = Arc::downgrade(&source.thread_manager); + let source_parent = Arc::downgrade(&parent); + parent.shutdown_and_wait().await?; + source.thread_manager.remove_thread(&root_id).await; + drop(parent); + drop(source); + drop(source_control); + assert!(source_manager.upgrade().is_none()); + assert!(source_parent.upgrade().is_none()); + + // Independent runtimes may share durable storage. + wiremock::Mock::given(wiremock::matchers::path("/codex/analytics-events/events")) + .respond_with(wiremock::ResponseTemplate::new(/*status*/ 200)) + .mount(&server) + .await; + let analytics = AnalyticsEventsClient::new( + codex_core::test_support::auth_manager_from_auth( + CodexAuth::create_dummy_chatgpt_auth_for_testing(), + ), + server.uri(), + /*analytics_enabled*/ Some(true), + ); + let destination = test_codex() + .with_analytics_events_client(analytics.clone()) + .with_thread_store(source_store) + .with_model("gpt-6-astra") + .with_config(configure_scenario_catalog) + .build_with_auto_env(&server) + .await?; + let manager = &destination.thread_manager; + destination.codex.shutdown_and_wait().await?; + manager + .remove_thread(&destination.session_configured.thread_id) + .await; + destination + .fs() + .write_file( + &agents_md, + PROJECT_RULE.as_bytes().to_vec(), + Default::default(), + /*sandbox*/ None, + ) + .await?; + let child_id = ThreadId::new(); + // The destination controller shares the tree identity but has no entry for the child path. + let destination_control = Arc::new(TestAgentControl::new(root_id)); + let child_request = + responses::mount_sse_once(&server, responses::sse_completed("remote-done")).await; + let spawned = manager + .create_prepared_child( + prepared, + HostThreadSpawnReservation { + thread_id: child_id, + control: destination_control, + }, + ) + .await?; + assert_eq!( + ( + spawned.agent.thread_id, + &spawned.task_name, + spawned.config.parent_thread_id, + ), + (child_id, &child_path, Some(root_id)) + ); + assert_eq!( + ( + spawned.config.permission_profile, + spawned.config.disabled_plugin_ids, + ), + (captured.permission_profile, captured.disabled_plugin_ids), + ); + let worker = manager.get_thread(child_id).await?; + start_first_turn(&worker, TASK).await?; + analytics.flush().await; + let events = responses::received_analytics_events(&server, "codex_thread_initialized").await; + let child_starts: Vec<_> = events + .iter() + .map(|event| &event["event_params"]) + .filter(|params| params["thread_id"] == json!(child_id)) + .collect(); + let [start] = child_starts.as_slice() else { + anyhow::bail!("expected one initialization for {child_id}, got {child_starts:?}"); + }; + assert_eq!(start["parent_thread_id"], json!(root_id)); + assert_eq!(start["app_server_client"]["client_name"], "source-client"); + assert_eq!(start["app_server_client"]["client_version"], "1.0"); + let requests = responses::received_responses_requests(&server).await; + let history = context_snapshot::format_request_history_snapshot( + "A prepares a child. After A and its manager shut down, a second manager creates the idle child and then starts its first turn.", + &requests, + &ContextSnapshotOptions::default().rewrite_known_segments(), + ); + let snapshot_name = if fork_history { + "detached_spawn_full_history" + } else { + "detached_spawn_fresh_history" + }; + insta::assert_snapshot!(snapshot_name, history); + let child_request = child_request.single_request(); + assert_eq!( + child_request.body_json().get("prompt_cache_key").cloned(), + cache_key.map(|key| json!(key)), + ); + assert!( + child_request + .inputs_of_type("additional_tools") + .iter() + .any(|item| { + responses::namespace_child_tool(item, "functions", "exec") + .and_then(|tool| tool["description"].as_str()) + .is_some_and(|description| description.contains("host_spawn(args:")) + }) + ); + let user_messages = child_request.message_input_texts("user"); + assert!(user_messages.iter().any(|text| text.contains(PROJECT_RULE))); + worker.ensure_rollout_materialized().await; + worker.flush_rollout().await?; + let stored = destination + .thread_store + .load_latest_model_context(LoadThreadHistoryParams { + thread_id: child_id, + include_archived: false, + }) + .await?; + assert!(stored.items.iter().any(|item| matches!(item, + RolloutItem::SessionMeta(meta) if meta.meta.session_id == session_id + && meta.meta.parent_thread_id == Some(root_id) + ))); + for input in [EARLIER_INPUT, SPAWN_INPUT] { + let span = input_text(input); + let inherited = stored.items.iter().find_map(|item| match item { + RolloutItem::ResponseItem(envelope) + if matches!(&envelope.item, ResponseItem::Message { role, content, .. } + if role == "user" && content.contains(&span)) => + { + Some( + envelope + .metadata + .as_ref() + .is_some_and(|meta| meta.inherited_user_message), + ) + } + _ => None, + }); + assert_eq!(inherited, fork_history.then_some(true), "{input}"); + } + assert!(manager.get_thread(root_id).await.is_err()); + worker.shutdown_and_wait().await?; + Ok(()) +} + +struct ChildStartupGate { + child_id: ThreadId, + entered: Notify, +} + +impl ThreadLifecycleContributor for ChildStartupGate { + fn on_thread_ready<'a>( + &'a self, + input: ThreadReadyInput<'a, Config>, + ) -> ExtensionFuture<'a, ()> { + Box::pin(async move { + if input.thread_store.level_id() == self.child_id.to_string() { + self.entered.notify_one(); + std::future::pending::<()>().await; + } + }) + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn cancelled_creation_cleans_registered_child() -> Result<()> { + skip_if_no_network!(Ok(())); + let server = start_mock_server().await; + let child_id = ThreadId::new(); + let gate = Arc::new(ChildStartupGate { + child_id, + entered: Notify::new(), + }); + let mut extensions = ExtensionRegistryBuilder::new(); + extensions.thread_lifecycle_contributor(gate.clone()); + let (test, parent, controller) = host_spawn_fixture( + &server, + test_codex().with_extensions(Arc::new(extensions.build())), + NetworkPolicy::unmanaged(), + PromptCacheKey::SessionDerived, + /*thread_instructions_provider*/ None, + ) + .await?; + let invocation = + host_spawn_invocation(&server, &parent, "cancel-spawn", "Start a worker.").await?; + let context = parent.tool_invocation(&invocation.turn_id, &invocation.call_id)?; + let spawn = tokio::spawn({ + let parent = Arc::clone(&parent); + async move { + parent + .create_child( + &context, + create_request("cancelled", /*fork_history*/ true), + HostThreadSpawnReservation { + thread_id: child_id, + control: controller, + }, + ) + .await + } + }); + tokio::time::timeout(Duration::from_secs(10), gate.entered.notified()) + .await + .context("child never reached the startup gate")?; + let child = test.thread_manager.get_thread(child_id).await?; + parent.submit(Op::Interrupt).await?; + assert!(matches!( + tokio::time::timeout(Duration::from_secs(10), spawn) + .await + .context("host spawn did not observe invocation revocation")??, + Err(ThreadSpawnError::Preparation(_)) + )); + tokio::time::timeout(Duration::from_secs(10), child.wait_until_terminated()) + .await + .context("cancelled child's runtime did not terminate")?; + assert!(test.thread_manager.get_thread(child_id).await.is_err()); + parent.shutdown_and_wait().await?; + Ok(()) +} diff --git a/codex-rs/core/tests/suite/skills_extension.rs b/codex-rs/core/tests/suite/skills_extension.rs index 460b836c97..d941aae973 100644 --- a/codex-rs/core/tests/suite/skills_extension.rs +++ b/codex-rs/core/tests/suite/skills_extension.rs @@ -384,16 +384,7 @@ async fn wait_for_analytics_events( ) -> Vec { let deadline = Instant::now() + Duration::from_secs(10); loop { - let events = server - .received_requests() - .await - .unwrap_or_default() - .into_iter() - .filter(|request| request.url.path() == "/codex/analytics-events/events") - .filter_map(|request| serde_json::from_slice::(&request.body).ok()) - .flat_map(|payload| payload["events"].as_array().cloned().unwrap_or_default()) - .filter(|event| event["event_type"] == event_type) - .collect::>(); + let events = responses::received_analytics_events(server, event_type).await; if events.len() >= expected_count { return events; } diff --git a/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__reserved_spawn__detached_spawn_fresh_history.snap b/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__reserved_spawn__detached_spawn_fresh_history.snap new file mode 100644 index 0000000000..9715147925 --- /dev/null +++ b/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__reserved_spawn__detached_spawn_fresh_history.snap @@ -0,0 +1,58 @@ +--- +source: core/tests/suite/scenarios_reserved_spawn.rs +expression: history +--- +Scenario: A prepares a child. After A and its manager shut down, a second manager creates the idle child and then starts its first turn. + +## Window 1 +-- request 1 (turn) -- +00:additional_tools/developer (2; hash=e03063c2cb6ece3a): + - namespace/functions; hash=e58195dea5c51dea + - custom/exec + - function/wait + - function/request_user_input + - function/request_user_input_async + - namespace/clock: Tools for reading and waiting on time.; hash=f402e8c5e9b5e317 + - function/sleep +01:message/developer: + +02:message/developer[2]: + [01] + [02] +03:message/developer: + +04:message/developer: + +05:message/user: + +06:message/user: + Remember PRIOR_CONTEXT_5241. +07:configuration_update:reasoning={"effort":"low"} +-- request 2 (turn) -- +08:message/user: + Remember TRANSFERRED_CONTEXT_7391 and ask the host to start a worker. + +## Window 2 (after request 2: input diverged at item 00) +-- request 3 (turn) -- +00:additional_tools/developer (2; hash=adec8e1bf249cb8f): + - namespace/functions; hash=00aa6cf7871d56b9 + - custom/exec + - function/wait + - function/request_user_input + - namespace/clock: Tools for reading and waiting on time.; hash=f402e8c5e9b5e317 + - function/sleep +01:message/developer: + +02:message/developer[2]: + [01] + [02] +03:message/developer: + +04:message/developer: + +05:message/user[2]: + [01] + [02] +06:message/user: + Report TRANSFERRED_TASK_8812 and the remembered context. +07:configuration_update:reasoning={"effort":"low"} diff --git a/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__reserved_spawn__detached_spawn_full_history.snap b/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__reserved_spawn__detached_spawn_full_history.snap new file mode 100644 index 0000000000..573bbc5ebd --- /dev/null +++ b/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__reserved_spawn__detached_spawn_full_history.snap @@ -0,0 +1,50 @@ +--- +source: core/tests/suite/scenarios_reserved_spawn.rs +expression: history +--- +Scenario: A prepares a child. After A and its manager shut down, a second manager creates the idle child and then starts its first turn. + +## Window 1 +-- request 1 (turn) -- +00:additional_tools/developer (2; hash=e03063c2cb6ece3a): + - namespace/functions; hash=e58195dea5c51dea + - custom/exec + - function/wait + - function/request_user_input + - function/request_user_input_async + - namespace/clock: Tools for reading and waiting on time.; hash=f402e8c5e9b5e317 + - function/sleep +01:message/developer: + +02:message/developer[2]: + [01] + [02] +03:message/developer: + +04:message/developer: + +05:message/user: + +06:message/user: + Remember PRIOR_CONTEXT_5241. +07:configuration_update:reasoning={"effort":"low"} +-- request 2 (turn) -- +08:message/user: + Remember TRANSFERRED_CONTEXT_7391 and ask the host to start a worker. +-- request 3 (turn) -- +09:function_call/host_spawn:{} +10:function_call_output:aborted +11:message/developer: + You are now a newly spawned child agent. The preceding conversation is inherited parent context. Follow ...are not inherited. Report your result to your parent. [hash=0e586a5d405d39db] +12:message/developer: + +13:message/developer: + The following tools are no longer available. Do not call them: + - functions.request_user_input_async +14:message/user: + +15:message/developer: + +16:message/user: + Report TRANSFERRED_TASK_8812 and the remembered context. +17:configuration_update:reasoning={"effort":"low"} diff --git a/codex-rs/core/tests/suite/tool_parallelism.rs b/codex-rs/core/tests/suite/tool_parallelism.rs index c41ad11c4d..b497d955f5 100644 --- a/codex-rs/core/tests/suite/tool_parallelism.rs +++ b/codex-rs/core/tests/suite/tool_parallelism.rs @@ -2,6 +2,7 @@ #![allow(clippy::unwrap_used)] use codex_core::TurnInputRequest; +use codex_features::Feature; use core_test_support::test_codex::local_requests; use std::fs; use std::time::Duration; @@ -70,84 +71,45 @@ async fn run_turn(test: &TestCodex, prompt: &str) -> anyhow::Result<()> { Ok(()) } -async fn run_turn_and_measure(test: &TestCodex, prompt: &str) -> anyhow::Result { - let start = Instant::now(); - run_turn(test, prompt).await?; - Ok(start.elapsed()) -} - async fn build_codex_with_test_tool(server: &wiremock::MockServer) -> anyhow::Result { let mut builder = test_codex().with_model("test-gpt-5.1-codex"); builder.build(server).await } -fn assert_parallel_duration(actual: Duration) { - // Allow headroom for slow CI scheduling; barrier synchronization already enforces overlap. - assert!( - actual < Duration::from_millis(1_600), - "expected parallel execution to finish quickly, got {actual:?}" - ); -} - #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn read_file_tools_run_in_parallel() -> anyhow::Result<()> { +async fn native_v2_tool_calls_with_the_same_id_run_in_parallel() -> anyhow::Result<()> { skip_if_no_network!(Ok(())); let server = start_mock_server().await; - let test = build_codex_with_test_tool(&server).await?; - - let warmup_args = json!({ - "sleep_after_ms": 10, - "barrier": { - "id": "parallel-test-sync-warmup", - "participants": 2, - "timeout_ms": 1_000, - } + let test = test_codex() + .with_model("test-gpt-5.1-codex") + .with_config(|config| { + config.features.enable(Feature::MultiAgentV2).unwrap(); + }) + .build(&server) + .await?; + let args = json!({ + "barrier": { "id": "parallel-test-sync", "participants": 2, "timeout_ms": 5_000 } }) .to_string(); - - let parallel_args = json!({ - "sleep_after_ms": 300, - "barrier": { - "id": "parallel-test-sync", - "participants": 2, - "timeout_ms": 1_000, - } - }) - .to_string(); - - let warmup_first = sse(vec![ - json!({"type": "response.created", "response": {"id": "resp-warm-1"}}), - ev_function_call("warm-call-1", "test_sync_tool", &warmup_args), - ev_function_call("warm-call-2", "test_sync_tool", &warmup_args), - ev_completed("resp-warm-1"), - ]); - let warmup_second = sse(vec![ - ev_assistant_message("warm-msg-1", "warmup complete"), - ev_completed("resp-warm-2"), - ]); - - let first_response = sse(vec![ - json!({"type": "response.created", "response": {"id": "resp-1"}}), - ev_function_call("call-1", "test_sync_tool", ¶llel_args), - ev_function_call("call-2", "test_sync_tool", ¶llel_args), + let calls = sse(vec![ + ev_response_created("resp-1"), + ev_function_call("shared", "test_sync_tool", &args), + ev_function_call("shared", "test_sync_tool", &args), ev_completed("resp-1"), ]); - let second_response = sse(vec![ + let done = sse(vec![ ev_assistant_message("msg-1", "done"), ev_completed("resp-2"), ]); - mount_sse_sequence( - &server, - vec![warmup_first, warmup_second, first_response, second_response], - ) - .await; - - run_turn(&test, "warm up parallel tool").await?; - - let duration = run_turn_and_measure(&test, "exercise sync tool").await?; - assert_parallel_duration(duration); + let responses = mount_sse_sequence(&server, vec![calls, done]).await; + run_turn(&test, "exercise sync tool").await?; + // Either call can return "ok" only after both handlers reach the barrier. + assert_eq!( + responses.function_call_output_text("shared").as_deref(), + Some("ok") + ); Ok(()) } @@ -240,8 +202,13 @@ async fn mixed_parallel_tools_run_in_parallel() -> anyhow::Result<()> { ]); mount_sse_sequence(&server, vec![first_response, second_response]).await; - let duration = run_turn_and_measure(&test, "mix tools").await?; - assert_parallel_duration(duration); + let start = Instant::now(); + run_turn(&test, "mix tools").await?; + let elapsed = start.elapsed(); + assert!( + elapsed < Duration::from_millis(1_600), + "expected parallel execution to finish quickly, got {elapsed:?}" + ); Ok(()) } diff --git a/codex-rs/ext/extension-api/src/message_lifecycle.rs b/codex-rs/ext/extension-api/src/message_lifecycle.rs index c81391dca5..85d3e169e9 100644 --- a/codex-rs/ext/extension-api/src/message_lifecycle.rs +++ b/codex-rs/ext/extension-api/src/message_lifecycle.rs @@ -12,7 +12,7 @@ pub struct MessageLifecycleInput { pub data: ExtensionData, } -/// Observes messages published through local agent control, where the sender is known. +/// Observes messages from local agent control and the initial task of a host-spawned child. /// Callbacks are synchronous and must not block. Replaying history or submitting raw /// protocol input does not create a publication. pub trait MessageLifecycleContributor: Send + Sync {