mirror of
https://github.com/openai/codex.git
synced 2026-11-10 21:00:47 +00:00
Add host-reserved V2 child creation APIs (#57586)
## What changed - Add `CodexThread::create_child` to create an idle child from a live tool invocation using a host-reserved thread ID and controller. - Add `prepare_child` and `ThreadManager::create_prepared_child` to capture inherited settings and fresh or forked history, then start the child on another manager even after the parent shuts down. Reject detached preparation when live bindings cannot be transferred. - Add opt-in `HostToolInvocationLookup` to track live calls and reject duplicate active call IDs within a turn. Revoke invocation access when dispatch ends and clean up provisional children when startup fails or is cancelled. ## Testing Add tests for invocation lifetime and call ID reservation, inherited policy and live bindings, detached startup with fresh and full history, and cancellation cleanup. GitOrigin-RevId: 71f17ecb4cd7636517e675df59289a8e8c7105b8
This commit is contained in:
rhan-oai
@@ -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<ThreadManagerState>,
|
||||
session_source: Option<&SessionSource>,
|
||||
) -> Option<TurnEnvironmentSnapshot> {
|
||||
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<Arc<crate::exec_policy::ExecPolicyManager>> {
|
||||
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<Arc<crate::exec_policy::ExecPolicyManager>> {
|
||||
let parent_config = parent_thread.session.get_config().await;
|
||||
if !crate::exec_policy::child_uses_parent_exec_policy(&parent_config, child_config) {
|
||||
return None;
|
||||
|
||||
@@ -130,7 +130,7 @@ impl LocalAgentRuntime {
|
||||
Ok(descendants)
|
||||
}
|
||||
|
||||
pub(super) fn upgrade(&self) -> CodexResult<Arc<ThreadManagerState>> {
|
||||
pub(in crate::agent) fn upgrade(&self) -> CodexResult<Arc<ThreadManagerState>> {
|
||||
self.manager.upgrade().ok_or_else(|| {
|
||||
CodexErr::UnsupportedOperation("thread manager dropped".to_string())
|
||||
.with_agent_context(AgentErrorContext::ManagerUnavailable)
|
||||
|
||||
@@ -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<ThreadConfigSnapshot> {
|
||||
pub(in crate::agent) async fn effective_spawn_config(
|
||||
thread: &CodexThread,
|
||||
) -> CodexResult<ThreadConfigSnapshot> {
|
||||
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<ThreadHistoryMode>, Vec<DynamicToolSpec>) {
|
||||
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 {
|
||||
|
||||
@@ -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<ThreadManagerState>,
|
||||
pending: Option<(ThreadId, AgentTreeMembership)>,
|
||||
edge_write: Option<JoinHandle<()>>,
|
||||
owner: SpawnOwner,
|
||||
}
|
||||
|
||||
impl PendingSpawn {
|
||||
pub(super) fn new(
|
||||
pub(in crate::agent) fn new(
|
||||
state: Arc<ThreadManagerState>,
|
||||
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
|
||||
|
||||
@@ -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<StepContext>,
|
||||
finished: Mutex<CancellationToken>,
|
||||
}
|
||||
|
||||
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<HashMap<(String, String), Weak<InvocationContext>>>);
|
||||
|
||||
/// 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<Session>,
|
||||
call_id: String,
|
||||
context: Arc<InvocationContext>,
|
||||
}
|
||||
|
||||
impl ToolInvocationContext {
|
||||
/// Order publication with revocation of the issuing tool call.
|
||||
pub(super) fn lock_active(
|
||||
&self,
|
||||
) -> Result<MutexGuard<'_, CancellationToken>, 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<ToolInvocationContext> {
|
||||
let context = thread
|
||||
.session
|
||||
.services
|
||||
.thread_extension_data
|
||||
.get::<InvocationContexts>()
|
||||
.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<InvocationContexts>,
|
||||
key: (String, String),
|
||||
}
|
||||
|
||||
/// Keeps the dispatch's ID reserved until its terminal callback returns.
|
||||
#[derive(Default)]
|
||||
pub(crate) struct InvocationReservation(Mutex<Option<ReservedInvocation>>);
|
||||
|
||||
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::<HostToolInvocationLookup>()
|
||||
.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<InvocationGuard> {
|
||||
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<ReservedInvocation> {
|
||||
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<InvocationContext>);
|
||||
|
||||
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<ThreadSpawnResult, ThreadSpawnError> {
|
||||
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<PreparedThreadSpawn, ThreadSpawnError> {
|
||||
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<T, F, Fut>(
|
||||
thread: &CodexThread,
|
||||
invocation: &ToolInvocationContext,
|
||||
request: ThreadCreateRequest,
|
||||
operation: F,
|
||||
) -> Result<T, ThreadSpawnError>
|
||||
where
|
||||
F: FnOnce(ChildSpawnSettings, AgentPath) -> Fut,
|
||||
Fut: Future<Output = Result<T, ThreadSpawnError>>,
|
||||
{
|
||||
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;
|
||||
@@ -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(())
|
||||
}
|
||||
@@ -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;
|
||||
@@ -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<PreparedThreadSpawn> {
|
||||
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<ThreadSpawnResult, ThreadSpawnError> {
|
||||
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;
|
||||
@@ -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<SpawnAgentForkMode>) -> 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::<HostToolInvocationLookup>()
|
||||
.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(())
|
||||
}
|
||||
@@ -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<dyn AgentControl>,
|
||||
}
|
||||
|
||||
/// 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<Arc<ExecPolicyManager>>,
|
||||
pub(crate) originator: String,
|
||||
pub(crate) attachment_source: Option<ThreadId>,
|
||||
}
|
||||
@@ -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;
|
||||
|
||||
Reference in New Issue
Block a user