diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index ca27e0208d..c9a4516494 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/codex-rs/core/src/agent/control.rs @@ -2,7 +2,6 @@ use crate::TurnInputRequest; use crate::TurnInputSubmission; use crate::TurnStartOptions; use crate::agent::AgentStatus; -use crate::agent::registry::AgentRegistry; use crate::agent::role::DEFAULT_ROLE_NAME; use crate::agent::role::resolve_role_config; use crate::agent::status::is_final; @@ -14,7 +13,6 @@ use crate::config::Config; use crate::config::RolloutBudgetConfig; use crate::context::SubagentNotification; use crate::environment_selection::TurnEnvironmentSnapshot; -use crate::rollout_budget::RolloutBudget; use crate::session::emit_subagent_session_started; use crate::session_prefix::format_inter_agent_completion_message; use crate::session_prefix::format_subagent_context_line; @@ -24,7 +22,6 @@ use crate::thread_manager::ThreadManagerState; use crate::thread_manager::default_thread_id_generator; use crate::thread_rollout_truncation::truncate_rollout_to_last_n_fork_turns; use crate::turn_timing::now_unix_timestamp_ms; -use arc_swap::ArcSwapOption; use codex_extension_api::ThreadInstructionsProvider; use codex_history::InitialHistory; use codex_history::ResumedHistory; @@ -60,13 +57,11 @@ use std::collections::HashMap; use std::collections::HashSet; use std::collections::VecDeque; use std::sync::Arc; -use std::sync::OnceLock; use std::sync::Weak; use tracing::warn; use uuid::Uuid; -use self::execution::AgentExecutionLimiter; -use self::residency::V2Residency; +pub(crate) use self::runtime::LocalAgentRuntime; mod api; mod budget; @@ -78,6 +73,7 @@ mod interrupt; mod legacy; mod residency; mod resume; +mod runtime; mod sender_context; mod service_tier; mod spawn; @@ -88,31 +84,14 @@ mod watch; const MAX_ENVIRONMENT_SUBAGENTS: usize = 8; const MAX_ENVIRONMENT_SUBAGENT_BYTES: usize = 1_024; -/// Control-plane handle for multi-agent operations. -/// `LocalAgentControl` is held by each session (via `SessionServices`). It provides capability to -/// spawn new agents and the inter-agent communication layer. -/// An `LocalAgentControl` instance is intended to be created at most once per root thread/session -/// tree. That same `LocalAgentControl` is then shared with every sub-agent spawned from that root, -/// which keeps the registry scoped to that root thread rather than the entire `ThreadManager`. +/// Per-session controller handle for a local agent tree. +/// Handles retain a session identity and share their tree's `LocalAgentRuntime`. +/// Local startup preserves that state when creating or resuming children. #[derive(Clone)] pub(crate) struct LocalAgentControl { /// session_id is equal to the root thread's ID. session_id: SessionId, - /// Weak handle back to the global thread registry/state. - /// This is `Weak` to avoid reference cycles and shadow persistence of the form - /// `ThreadManagerState -> CodexThread -> Session -> SessionServices -> ThreadManagerState`. - manager: Weak, - /// Captured at construction so delegates retain their manager's allocation policy. - thread_id_generator: ThreadIdGenerator, - state: Arc, - v2_residency: Arc, - agent_execution_limiter: Arc, - /// Session-scoped state shared by the root thread and every cloned sub-agent control handle. - rollout_budget: Arc, - /// The user-selected root routing tier, shared by the entire agent tree. - root_service_tier: Arc>, - /// Retains the root's opt-in instruction provider even when the root is unloaded. - shared_thread_instructions_provider: Arc>>, + pub(crate) runtime: LocalAgentRuntime, } impl Default for LocalAgentControl { @@ -132,26 +111,15 @@ impl LocalAgentControl { thread_id_generator: ThreadIdGenerator, rollout_budget: Option, ) -> Self { - let control = Self { + Self { session_id: SessionId::default(), - manager, - thread_id_generator, - state: Arc::default(), - v2_residency: Arc::default(), - agent_execution_limiter: Arc::default(), - rollout_budget: Arc::default(), - root_service_tier: Arc::new(ArcSwapOption::from(None)), - shared_thread_instructions_provider: Arc::default(), - }; - if let Some(rollout_budget) = rollout_budget { - control.rollout_budget.configure(rollout_budget); + runtime: LocalAgentRuntime::new(manager, thread_id_generator, rollout_budget), } - control } pub(crate) fn with_session_id(mut self, session_id: SessionId, max_threads: usize) -> Self { self.session_id = session_id; - self.agent_execution_limiter.initialize(max_threads); + self.runtime.agent_execution_limiter.initialize(max_threads); self } @@ -160,7 +128,7 @@ impl LocalAgentControl { } pub(crate) fn generate_thread_id(&self) -> ThreadId { - (self.thread_id_generator)() + (self.runtime.thread_id_generator)() } pub(crate) fn root_thread_instructions_provider( @@ -168,7 +136,7 @@ impl LocalAgentControl { root_thread_id: ThreadId, provider: Option>, ) -> Option> { - let provider = match self.manager.upgrade() { + let provider = match self.runtime.manager.upgrade() { Some(manager) => manager.shared_thread_instructions_provider(root_thread_id, provider), None => provider, }; @@ -177,6 +145,7 @@ impl LocalAgentControl { .filter(|provider| provider.share_with_subagents()) { let _ = self + .runtime .shared_thread_instructions_provider .set(Arc::clone(provider)); } @@ -381,7 +350,7 @@ impl LocalAgentControl { { let _ = state.remove_thread(&agent_id).await; self.forget_v2_residency(agent_id); - self.state.release_spawned_thread(agent_id); + self.runtime.registry.release_spawned_thread(agent_id); } result } @@ -404,16 +373,19 @@ impl LocalAgentControl { current_parent_thread_id: Option, ) { if current_parent_thread_id.is_none() { - self.state.register_root_thread(current_thread_id); + self.runtime + .registry + .register_root_thread(current_thread_id); } } pub(crate) fn get_agent_metadata(&self, agent_id: ThreadId) -> Option { - self.state.agent_metadata_for_thread(agent_id) + self.runtime.registry.agent_metadata_for_thread(agent_id) } pub(crate) fn ensure_agent_known(&self, agent_id: ThreadId) -> CodexResult { - self.state + self.runtime + .registry .agent_metadata_for_thread(agent_id) .ok_or_else(|| CodexErr::ThreadNotFound(agent_id)) } @@ -451,7 +423,8 @@ impl LocalAgentControl { } let Some(parent_path) = self - .state + .runtime + .registry .agent_metadata_for_thread(parent_thread_id) .and_then(|metadata| metadata.agent_path) else { @@ -459,7 +432,8 @@ impl LocalAgentControl { }; let parent_prefix = format!("{parent_path}/"); let mut agent_paths = self - .state + .runtime + .registry .live_agents() .into_iter() .filter_map(|metadata| metadata.agent_path) @@ -512,7 +486,7 @@ impl LocalAgentControl { }) .transpose()?; - let mut live_agents = self.state.live_agents(); + let mut live_agents = self.runtime.registry.live_agents(); live_agents.sort_by(|left, right| { left.agent_path .as_deref() @@ -531,7 +505,7 @@ impl LocalAgentControl { if resolved_prefix .as_ref() .is_none_or(|prefix| agent_matches_prefix(Some(&root_path), prefix)) - && let Some(root_thread_id) = self.state.agent_id_for_path(&root_path) + && let Some(root_thread_id) = self.runtime.registry.agent_id_for_path(&root_path) && let Ok(root_thread) = state.get_thread(root_thread_id).await { agents.push(LiveAgent { @@ -706,7 +680,7 @@ impl LocalAgentControl { preferred_agent_nickname: Option, ) -> CodexResult<(SessionSource, AgentMetadata)> { if depth == 1 { - self.state.register_root_thread(parent_thread_id); + self.runtime.registry.register_root_thread(parent_thread_id); } let agent_metadata = self.prepare_agent_metadata( reservation, @@ -726,7 +700,8 @@ impl LocalAgentControl { } fn upgrade(&self) -> CodexResult> { - self.manager + self.runtime + .manager .upgrade() .ok_or_else(|| CodexErr::UnsupportedOperation("thread manager dropped".to_string())) } @@ -798,7 +773,8 @@ impl LocalAgentControl { .or_default() .push(( child_thread_id, - self.state + self.runtime + .registry .agent_metadata_for_thread(child_thread_id) .unwrap_or(AgentMetadata { agent_id: Some(child_thread_id), diff --git a/codex-rs/core/src/agent/control/budget.rs b/codex-rs/core/src/agent/control/budget.rs index 855072b9c4..3dc0b7f1d9 100644 --- a/codex-rs/core/src/agent/control/budget.rs +++ b/codex-rs/core/src/agent/control/budget.rs @@ -10,7 +10,7 @@ use codex_protocol::protocol::TokenUsage; impl LocalAgentControl { pub(crate) fn record_rollout_budget_usage(&self, usage: &TokenUsage) -> CodexResult<()> { - if self.rollout_budget.record_usage(usage)? { + if self.runtime.rollout_budget.record_usage(usage)? { return Err(CodexErr::SessionBudgetExceeded); } Ok(()) @@ -21,7 +21,9 @@ impl LocalAgentControl { thread_id: ThreadId, window_id: &str, ) -> Option { - self.rollout_budget.pending_reminder(thread_id, window_id) + self.runtime + .rollout_budget + .pending_reminder(thread_id, window_id) } pub(crate) fn mark_budget_reminder_delivered( @@ -30,7 +32,8 @@ impl LocalAgentControl { window_id: &str, reminder: RolloutBudgetReminder, ) { - self.rollout_budget + self.runtime + .rollout_budget .mark_reminder_delivered(thread_id, window_id, reminder); } } diff --git a/codex-rs/core/src/agent/control/execution.rs b/codex-rs/core/src/agent/control/execution.rs index 500905c3b8..9b605bd3cc 100644 --- a/codex-rs/core/src/agent/control/execution.rs +++ b/codex-rs/core/src/agent/control/execution.rs @@ -53,8 +53,8 @@ impl LocalAgentControl { if !is_execution_limited(multi_agent_version, session_source) { return Ok(()); } - let max_threads = self.agent_execution_limiter.max_threads(); - if self.agent_execution_limiter.has_capacity() { + let max_threads = self.runtime.agent_execution_limiter.max_threads(); + if self.runtime.agent_execution_limiter.has_capacity() { Ok(()) } else { Err(CodexErr::new(CodexErrorDetails::AgentLimitReached { @@ -69,7 +69,7 @@ impl LocalAgentControl { session_source: &SessionSource, ) -> Option { is_execution_limited(multi_agent_version, session_source) - .then(|| Arc::clone(&self.agent_execution_limiter).guard()) + .then(|| Arc::clone(&self.runtime.agent_execution_limiter).guard()) } } diff --git a/codex-rs/core/src/agent/control/execution_tests.rs b/codex-rs/core/src/agent/control/execution_tests.rs index 219844da10..f1b10840d1 100644 --- a/codex-rs/core/src/agent/control/execution_tests.rs +++ b/codex-rs/core/src/agent/control/execution_tests.rs @@ -7,7 +7,10 @@ use pretty_assertions::assert_eq; fn control_with_limit(max_threads: usize) -> LocalAgentControl { let control = LocalAgentControl::default(); - control.agent_execution_limiter.initialize(max_threads); + control + .runtime + .agent_execution_limiter + .initialize(max_threads); control } @@ -16,6 +19,7 @@ fn execution_guards_count_active_v2_subagent_turns() { let control = control_with_limit(/*max_threads*/ 1); // Child role configs cannot replace the root-derived session limit. control + .runtime .agent_execution_limiter .initialize(/*max_threads*/ 2); let source = SessionSource::SubAgent(SubAgentSource::Other("worker".to_string())); diff --git a/codex-rs/core/src/agent/control/legacy.rs b/codex-rs/core/src/agent/control/legacy.rs index e9952a6237..6586f308e5 100644 --- a/codex-rs/core/src/agent/control/legacy.rs +++ b/codex-rs/core/src/agent/control/legacy.rs @@ -40,7 +40,7 @@ impl LocalAgentControl { }; let _ = state.remove_thread(&agent_id).await; self.forget_v2_residency(agent_id); - self.state.release_spawned_thread(agent_id); + self.runtime.registry.release_spawned_thread(agent_id); result } diff --git a/codex-rs/core/src/agent/control/residency.rs b/codex-rs/core/src/agent/control/residency.rs index 9e8d5ad331..f1041c9c4e 100644 --- a/codex-rs/core/src/agent/control/residency.rs +++ b/codex-rs/core/src/agent/control/residency.rs @@ -1,4 +1,5 @@ use super::LocalAgentControl; +use super::LocalAgentRuntime; use crate::agent::AgentStatus; use crate::codex_thread::CodexThread; use crate::config::Config; @@ -55,7 +56,7 @@ impl LocalAgentControl { let capacity = config .effective_agent_max_threads(MultiAgentVersion::V2) .unwrap_or(usize::MAX); - Arc::clone(&self.v2_residency) + Arc::clone(&self.runtime.residency) .reserve_slot(state, capacity, protected_thread_id) .await } @@ -66,10 +67,16 @@ impl LocalAgentControl { thread_id: ThreadId, ) { if let Ok(thread) = state.get_thread(thread_id).await { - let _ = self.pin_v2_residency(state, &thread).await; + let _ = self.runtime.pin_v2_residency(state, &thread).await; } } + pub(super) fn forget_v2_residency(&self, thread_id: ThreadId) { + self.runtime.residency.remove(thread_id); + } +} + +impl LocalAgentRuntime { /// Pins and touches the registered runtime without waiting for unrelated eviction. pub(crate) async fn pin_v2_residency( &self, @@ -84,13 +91,9 @@ impl LocalAgentControl { if !Arc::ptr_eq(thread, &state.get_thread(thread_id).await?) { return Err(CodexErr::ThreadNotFound(thread_id)); } - self.v2_residency.touch(thread_id); + self.residency.touch(thread_id); Ok(Some(guard)) } - - pub(super) fn forget_v2_residency(&self, thread_id: ThreadId) { - self.v2_residency.remove(thread_id); - } } impl V2Residency { @@ -195,8 +198,8 @@ impl V2Residency { candidate_thread .session .services - .agent_control - .state + .local_agent_runtime + .registry .save_evicted_environments(candidate_thread_id, environments); // Keep publication excluded until both entries have been removed. threads.remove(&candidate_thread_id); diff --git a/codex-rs/core/src/agent/control/runtime.rs b/codex-rs/core/src/agent/control/runtime.rs new file mode 100644 index 0000000000..6af3642396 --- /dev/null +++ b/codex-rs/core/src/agent/control/runtime.rs @@ -0,0 +1,69 @@ +//! Shared state and startup bindings for one local agent tree. +//! Registry identity is allocation identity; cloning this handle preserves ownership checks. + +use super::LocalAgentControl; +use super::execution::AgentExecutionLimiter; +use super::residency::V2Residency; +use crate::agent::registry::AgentRegistry; +use crate::config::RolloutBudgetConfig; +use crate::rollout_budget::RolloutBudget; +use crate::thread_manager::ThreadIdGenerator; +use crate::thread_manager::ThreadManagerState; +use arc_swap::ArcSwapOption; +use codex_extension_api::ThreadInstructionsProvider; +use codex_protocol::SessionId; +use std::sync::Arc; +use std::sync::OnceLock; +use std::sync::Weak; + +/// Local tree state, kept separate from the shared agent operation interface. +#[derive(Clone)] +pub(crate) struct LocalAgentRuntime { + /// Weak handle back to the global thread registry/state. + /// This is `Weak` to avoid reference cycles and shadow persistence of the form + /// `ThreadManagerState -> CodexThread -> Session -> SessionServices -> ThreadManagerState`. + pub(super) manager: Weak, + /// Captured at construction so delegates retain their manager's allocation policy. + pub(super) thread_id_generator: ThreadIdGenerator, + pub(super) agent_execution_limiter: Arc, + /// Session-scoped state shared by the root thread and every cloned sub-agent control handle. + pub(super) rollout_budget: Arc, + /// The user-selected root routing tier, shared by the entire agent tree. + pub(super) root_service_tier: Arc>, + /// Retains the root's opt-in instruction provider even when the root is unloaded. + pub(super) shared_thread_instructions_provider: + Arc>>, + pub(super) registry: Arc, + pub(super) residency: Arc, +} + +impl LocalAgentRuntime { + pub(super) fn new( + manager: Weak, + thread_id_generator: ThreadIdGenerator, + rollout_budget: Option, + ) -> Self { + let runtime = Self { + manager, + thread_id_generator, + registry: Arc::default(), + residency: Arc::default(), + agent_execution_limiter: Arc::default(), + rollout_budget: Arc::default(), + root_service_tier: Arc::new(ArcSwapOption::from(None)), + shared_thread_instructions_provider: Arc::default(), + }; + if let Some(rollout_budget) = rollout_budget { + runtime.rollout_budget.configure(rollout_budget); + } + runtime + } + + /// Bind local startup to the same tree state with this session's identity. + pub(crate) fn control(&self, session_id: SessionId) -> LocalAgentControl { + LocalAgentControl { + session_id, + runtime: self.clone(), + } + } +} diff --git a/codex-rs/core/src/agent/control/service_tier.rs b/codex-rs/core/src/agent/control/service_tier.rs index 65e76bf7d0..3324cebfbb 100644 --- a/codex-rs/core/src/agent/control/service_tier.rs +++ b/codex-rs/core/src/agent/control/service_tier.rs @@ -6,13 +6,16 @@ use std::sync::Arc; impl LocalAgentControl { /// Returns the latest user-selected tier for this root and all its descendants. pub(crate) fn root_service_tier(&self) -> Option { - self.root_service_tier + self.runtime + .root_service_tier .load_full() .map(|service_tier| (*service_tier).clone()) } /// Publishes a root-owned tier without mutating individual child sessions. pub(crate) fn set_root_service_tier(&self, service_tier: Option) { - self.root_service_tier.store(service_tier.map(Arc::new)); + self.runtime + .root_service_tier + .store(service_tier.map(Arc::new)); } } diff --git a/codex-rs/core/src/agent/control/spawn.rs b/codex-rs/core/src/agent/control/spawn.rs index c495c3d6c6..536ca7f554 100644 --- a/codex-rs/core/src/agent/control/spawn.rs +++ b/codex-rs/core/src/agent/control/spawn.rs @@ -185,7 +185,7 @@ impl LocalAgentControl { config: &Config, root_thread_id: ThreadId, ) { - self.state.register_root_thread(root_thread_id); + self.runtime.registry.register_root_thread(root_thread_id); let Ok(state) = self.upgrade() else { return; @@ -208,7 +208,12 @@ impl LocalAgentControl { }; for thread_id in descendant_ids { - if self.state.agent_metadata_for_thread(thread_id).is_some() { + if self + .runtime + .registry + .agent_metadata_for_thread(thread_id) + .is_some() + { continue; } let restore_result = async { @@ -227,7 +232,10 @@ impl LocalAgentControl { .map_err(|err| { CodexErr::InvalidRequest(format!("invalid stored agent path: {err}")) })?; - let mut reservation = self.state.reserve_spawn_slot(/*max_threads*/ None)?; + let mut reservation = self + .runtime + .registry + .reserve_spawn_slot(/*max_threads*/ None)?; let mut metadata = self.prepare_agent_metadata( &mut reservation, config, @@ -276,7 +284,10 @@ impl LocalAgentControl { if thread.is_running() && thread.multi_agent_version() == Some(MultiAgentVersion::V2) && thread.session_source.parent_thread_id() == Some(parent_thread_id) - && Arc::ptr_eq(&self.state, &thread.session.services.agent_control.state) + && Arc::ptr_eq( + &self.runtime.registry, + &thread.session.services.local_agent_runtime.registry, + ) { return Ok(()); } @@ -303,7 +314,10 @@ impl LocalAgentControl { .is_some_and(|registered| Arc::ptr_eq(registered, parent)) || !parent.is_running() || parent.multi_agent_version() != Some(MultiAgentVersion::V2) - || !Arc::ptr_eq(&self.state, &parent.session.services.agent_control.state) + || !Arc::ptr_eq( + &self.runtime.registry, + &parent.session.services.local_agent_runtime.registry, + ) { return Err(CodexErr::InvalidRequest(format!( "cannot resume multi-agent v2 child {thread_id}: parent ownership is unavailable; resume the parent first" @@ -314,10 +328,15 @@ impl LocalAgentControl { self.touch_loaded_v2_residency(&state, thread_id).await; return Ok(()); } - if self.state.agent_metadata_for_thread(thread_id).is_none() { + if self + .runtime + .registry + .agent_metadata_for_thread(thread_id) + .is_none() + { return Err(CodexErr::ThreadNotFound(thread_id)); } - let mut environment_selections = self.state.evicted_environments(thread_id); + let mut environment_selections = self.runtime.registry.evicted_environments(thread_id); let stored_thread = state .read_stored_thread(ReadThreadParams { @@ -544,7 +563,8 @@ impl LocalAgentControl { { Some(parent.session.inherited_instructions().await) } else { - self.shared_thread_instructions_provider + self.runtime + .shared_thread_instructions_provider .get() .map(|provider| SessionInstructions { thread_provider: Some(Arc::clone(provider)), @@ -576,7 +596,7 @@ impl LocalAgentControl { if let Some(parent_thread_id) = owner_thread_id { self.validate_loaded_v2_child(&reloaded_thread.thread, parent_thread_id)?; } - self.state.clear_evicted_environments(thread_id); + self.runtime.registry.clear_evicted_environments(thread_id); residency_slot.commit(reloaded_thread.thread_id); state.notify_thread_created(reloaded_thread.thread_id); Ok(()) @@ -586,7 +606,7 @@ impl LocalAgentControl { if let Some(parent_thread_id) = owner_thread_id { self.validate_loaded_v2_child(&thread, parent_thread_id)?; } - self.state.clear_evicted_environments(thread_id); + self.runtime.registry.clear_evicted_environments(thread_id); drop(residency_slot); self.touch_loaded_v2_residency(&state, thread_id).await; return Ok(()); @@ -634,7 +654,10 @@ impl LocalAgentControl { } else { agent_max_threads }; - let mut reservation = self.state.reserve_spawn_slot(reservation_max_threads)?; + let mut reservation = self + .runtime + .registry + .reserve_spawn_slot(reservation_max_threads)?; let inheritance = SpawnAgentThreadInheritance { environments: self .inherited_environments_for_source(&state, session_source.as_ref()) @@ -1246,7 +1269,10 @@ impl LocalAgentControl { ) .await; let agent_max_threads = config.effective_agent_max_threads(multi_agent_version); - let mut reservation = self.state.reserve_spawn_slot(agent_max_threads)?; + let mut reservation = self + .runtime + .registry + .reserve_spawn_slot(agent_max_threads)?; let (session_source, agent_metadata) = match session_source { SessionSource::SubAgent(SubAgentSource::ThreadSpawn { parent_thread_id, diff --git a/codex-rs/core/src/agent/control/target.rs b/codex-rs/core/src/agent/control/target.rs index 33585b561b..44d07873a3 100644 --- a/codex-rs/core/src/agent/control/target.rs +++ b/codex-rs/core/src/agent/control/target.rs @@ -47,7 +47,7 @@ impl LocalAgentControl { let agent_path = current_agent_path .resolve(agent_reference) .map_err(CodexErr::UnsupportedOperation)?; - if let Some(thread_id) = self.state.agent_id_for_path(&agent_path) { + if let Some(thread_id) = self.runtime.registry.agent_id_for_path(&agent_path) { return Ok(thread_id); } Err(CodexErr::UnsupportedOperation(format!( diff --git a/codex-rs/core/src/agent/control/user_authorization.rs b/codex-rs/core/src/agent/control/user_authorization.rs index f8d2a374ed..9eb52d217d 100644 --- a/codex-rs/core/src/agent/control/user_authorization.rs +++ b/codex-rs/core/src/agent/control/user_authorization.rs @@ -38,7 +38,10 @@ impl LocalAgentControl { &self, thread_id: ThreadId, ) -> Option { - let root_thread_id = self.state.agent_id_for_path(&AgentPath::root())?; + let root_thread_id = self + .runtime + .registry + .agent_id_for_path(&AgentPath::root())?; if root_thread_id == thread_id { return None; } diff --git a/codex-rs/core/src/agent/control_tests.rs b/codex-rs/core/src/agent/control_tests.rs index df4a842ab9..d2bd966517 100644 --- a/codex-rs/core/src/agent/control_tests.rs +++ b/codex-rs/core/src/agent/control_tests.rs @@ -209,7 +209,13 @@ fn register_session_root_skips_threads_with_explicit_parent() { control.register_session_root(ThreadId::new(), Some(ThreadId::new())); - assert_eq!(control.state.agent_id_for_path(&AgentPath::root()), None); + assert_eq!( + control + .runtime + .registry + .agent_id_for_path(&AgentPath::root()), + None + ); } fn spawn_agent_call(call_id: &str) -> ResponseItem { diff --git a/codex-rs/core/src/codex_delegate.rs b/codex-rs/core/src/codex_delegate.rs index 49d8ee4c18..8f942cc182 100644 --- a/codex-rs/core/src/codex_delegate.rs +++ b/codex-rs/core/src/codex_delegate.rs @@ -114,7 +114,10 @@ pub(crate) async fn run_codex_thread_interactive( ThreadSource::Subagent }), originator: parent_ctx.originator.clone(), - agent_control: parent_session.services.agent_control.clone(), + agent_control: parent_session + .services + .local_agent_runtime + .control(parent_session.session_id()), dynamic_tools: Vec::new(), metrics_service_name: None, user_shell_override: None, diff --git a/codex-rs/core/src/guardian/review_session_setup.rs b/codex-rs/core/src/guardian/review_session_setup.rs index acdb5df36d..69814751eb 100644 --- a/codex-rs/core/src/guardian/review_session_setup.rs +++ b/codex-rs/core/src/guardian/review_session_setup.rs @@ -102,7 +102,11 @@ impl PreparedGuardianContext { internal_parent: Some(crate::thread_manager::InternalSessionParent { thread_id: self.parent.thread_id(), auth_manager: Arc::clone(&self.parent.services.auth_manager), - agent_control: self.parent.services.agent_control.clone(), + agent_control: self + .parent + .services + .local_agent_runtime + .control(self.parent.session_id()), originator: self.context.turn().originator.clone(), // Review the same applied instructions captured by the reuse key. // A live provider could advance independently while reviewing this action. diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index b65e77d767..917acf1577 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -1680,6 +1680,7 @@ impl Session { selected_capability_roots, mcp_thread_init, client_mcp_extensions, + local_agent_runtime: agent_control.runtime.clone(), agent_control, network_proxy: arc_swap::ArcSwapOption::from(network_proxy.map(Arc::new)), network_proxy_audit_metadata, diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 74b3c04d61..4454f1c9a7 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -6366,6 +6366,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { selected_capability_roots: Vec::new(), mcp_thread_init: codex_extension_api::ExtensionDataInit::default(), client_mcp_extensions: ClientMcpExtensions::default(), + local_agent_runtime: agent_control.runtime.clone(), agent_control, network_proxy: arc_swap::ArcSwapOption::from(None), network_proxy_audit_metadata: crate::config::NetworkProxyAuditMetadata::default(), @@ -8632,6 +8633,7 @@ where selected_capability_roots: Vec::new(), mcp_thread_init: codex_extension_api::ExtensionDataInit::default(), client_mcp_extensions: ClientMcpExtensions::default(), + local_agent_runtime: agent_control.runtime.clone(), agent_control, network_proxy: arc_swap::ArcSwapOption::from(None), network_proxy_audit_metadata: crate::config::NetworkProxyAuditMetadata::default(), diff --git a/codex-rs/core/src/state/service.rs b/codex-rs/core/src/state/service.rs index 5bed6d03d7..c92069c692 100644 --- a/codex-rs/core/src/state/service.rs +++ b/codex-rs/core/src/state/service.rs @@ -2,6 +2,7 @@ use std::path::PathBuf; use std::sync::Arc; use crate::agent::LocalAgentControl; +use crate::agent::control::LocalAgentRuntime; use crate::agents_md_manager::AgentsMdManager; use crate::attestation::AttestationProvider; use crate::client::ModelClient; @@ -83,6 +84,7 @@ pub(crate) struct SessionServices { pub(crate) selected_capability_roots: Vec, pub(crate) mcp_thread_init: ExtensionDataInit, pub(crate) agent_control: LocalAgentControl, + pub(crate) local_agent_runtime: LocalAgentRuntime, pub(crate) network_proxy: ArcSwapOption, pub(crate) network_proxy_audit_metadata: NetworkProxyAuditMetadata, pub(crate) managed_network_requirements_configured: bool, diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index fcf5c1b173..6725939d2f 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -1091,7 +1091,11 @@ impl ThreadManager { options.internal_parent = Some(InternalSessionParent { thread_id: parent_thread_id, auth_manager: Arc::clone(&parent.session.services.auth_manager), - agent_control: parent.session.services.agent_control.clone(), + agent_control: parent + .session + .services + .local_agent_runtime + .control(parent.session.session_id()), originator: parent.config_snapshot().await.originator, inherited_instructions, }); @@ -1234,7 +1238,11 @@ impl ThreadManager { )) })?; let config = parent.session.get_config().await.as_ref().clone(); - let agent_control = parent.session.services.agent_control.clone(); + let agent_control = parent + .session + .services + .local_agent_runtime + .control(parent.session.session_id()); agent_control .ensure_v2_agent_loaded(config, child_thread_id, Some(parent)) .await @@ -1695,7 +1703,7 @@ impl ThreadManagerState { thread .session .services - .agent_control + .local_agent_runtime .pin_v2_residency(self, &thread) .await? } else { diff --git a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs index 16a1e4d9b0..0bba76c7f2 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs @@ -84,6 +84,14 @@ use tokio::sync::Mutex; use tokio::time::timeout; use tokio_util::sync::CancellationToken; +fn set_agent_control( + session: &mut crate::session::session::Session, + control: crate::agent::LocalAgentControl, +) { + session.services.local_agent_runtime = control.runtime.clone(); + session.services.agent_control = control; +} + fn invocation( session: Arc, turn: Arc, @@ -296,7 +304,7 @@ async fn spawn_agent_uses_explorer_role_and_preserves_approval_policy() { let (mut session, mut turn) = make_session_and_context().await; let manager = thread_manager(); - session.services.agent_control = manager.agent_control(); + set_agent_control(&mut session, manager.agent_control()); let mut config = (*turn.config).clone(); let provider_info = built_in_model_providers(/* openai_base_url */ /*openai_base_url*/ None)["ollama"].clone(); @@ -352,7 +360,7 @@ async fn spawn_agent_fork_context_rejects_agent_type_override() { .start_thread(StartThreadOptions::new((*turn.config).clone())) .await .expect("root thread should start"); - session.services.agent_control = manager.agent_control(); + set_agent_control(&mut session, manager.agent_control()); session.thread_id = root.thread_id; let err = SpawnAgentHandler::default() .handle(invocation( @@ -386,7 +394,7 @@ async fn multi_agent_v2_spawn_fork_turns_all_applies_agent_type_override() { .start_thread(StartThreadOptions::new((*turn.config).clone())) .await .expect("root thread should start"); - session.services.agent_control = manager.agent_control(); + set_agent_control(&mut session, manager.agent_control()); session.thread_id = root.thread_id; let mut config = (*turn.config).clone(); config @@ -475,7 +483,10 @@ async fn spawn_agent_service_tier_inheritance_uses_root_preference_and_child_mod .start_thread(StartThreadOptions::new((*turn.config).clone())) .await .expect("root thread should start"); - session.services.agent_control = root.thread.session.services.agent_control.clone(); + set_agent_control( + &mut session, + root.thread.session.services.agent_control.clone(), + ); session.services.models_manager = Arc::new(StaticModelsManager::new( /*auth_manager*/ None, service_tier_test_catalog(), @@ -521,7 +532,10 @@ async fn spawn_agent_service_tier_inheritance_uses_root_preference_and_child_mod .start_thread(StartThreadOptions::new((*turn.config).clone())) .await .expect("root thread should start"); - session.services.agent_control = root.thread.session.services.agent_control.clone(); + set_agent_control( + &mut session, + root.thread.session.services.agent_control.clone(), + ); session.services.models_manager = Arc::new(StaticModelsManager::new( /*auth_manager*/ None, service_tier_test_catalog(), @@ -588,7 +602,7 @@ service_tier = "priority" .start_thread(StartThreadOptions::new((*turn.config).clone())) .await .expect("root thread should start"); - session.services.agent_control = manager.agent_control(); + set_agent_control(&mut session, manager.agent_control()); session.thread_id = root.thread_id; let output = SpawnAgentHandler::default() @@ -658,7 +672,10 @@ service_tier = "turbo" .start_thread(StartThreadOptions::new((*turn.config).clone())) .await .expect("root thread should start"); - session.services.agent_control = root.thread.session.services.agent_control.clone(); + set_agent_control( + &mut session, + root.thread.session.services.agent_control.clone(), + ); session.thread_id = root.thread_id; let output = SpawnAgentHandler::default() @@ -708,7 +725,10 @@ async fn spawn_agent_full_history_fork_inherits_root_service_tier() { .start_thread(StartThreadOptions::new((*turn.config).clone())) .await .expect("root thread should start"); - session.services.agent_control = root.thread.session.services.agent_control.clone(); + set_agent_control( + &mut session, + root.thread.session.services.agent_control.clone(), + ); session.thread_id = root.thread_id; let output = SpawnAgentHandler::default() @@ -762,7 +782,10 @@ async fn multi_agent_v2_full_history_fork_inherits_root_service_tier() { .start_thread(StartThreadOptions::new((*turn.config).clone())) .await .expect("root thread should start"); - session.services.agent_control = root.thread.session.services.agent_control.clone(); + set_agent_control( + &mut session, + root.thread.session.services.agent_control.clone(), + ); session.thread_id = root.thread_id; let session = Arc::new(session); let turn = Arc::new(turn); @@ -814,7 +837,7 @@ async fn multi_agent_v2_spawn_partial_fork_turns_allows_agent_type_override() { .start_thread(StartThreadOptions::new((*turn.config).clone())) .await .expect("root thread should start"); - session.services.agent_control = manager.agent_control(); + set_agent_control(&mut session, manager.agent_control()); session.thread_id = root.thread_id; let mut config = (*turn.config).clone(); config @@ -866,7 +889,7 @@ async fn multi_agent_v2_spawn_partial_fork_turns_allows_agent_type_override() { async fn spawn_agent_returns_agent_id_without_task_name() { let (mut session, turn) = make_session_and_context().await; let manager = thread_manager(); - session.services.agent_control = manager.agent_control(); + set_agent_control(&mut session, manager.agent_control()); let output = SpawnAgentHandler::default() .handle(invocation( @@ -897,7 +920,7 @@ async fn multi_agent_v2_spawn_requires_task_name() { .start_thread(StartThreadOptions::new((*turn.config).clone())) .await .expect("root thread should start"); - session.services.agent_control = manager.agent_control(); + set_agent_control(&mut session, manager.agent_control()); session.thread_id = root.thread_id; let mut config = (*turn.config).clone(); config @@ -931,7 +954,7 @@ async fn multi_agent_v2_spawn_rejects_legacy_items_field() { .start_thread(StartThreadOptions::new((*turn.config).clone())) .await .expect("root thread should start"); - session.services.agent_control = manager.agent_control(); + set_agent_control(&mut session, manager.agent_control()); session.thread_id = root.thread_id; let mut config = (*turn.config).clone(); config @@ -991,7 +1014,7 @@ async fn multi_agent_v2_spawn_returns_path_and_send_message_accepts_relative_pat .start_thread(StartThreadOptions::new((*turn.config).clone())) .await .expect("root thread should start"); - session.services.agent_control = manager.agent_control(); + set_agent_control(&mut session, manager.agent_control()); session.thread_id = root.thread_id; let mut config = (*turn.config).clone(); config @@ -1086,7 +1109,7 @@ async fn multi_agent_v2_spawn_rejects_legacy_fork_context() { .start_thread(StartThreadOptions::new((*turn.config).clone())) .await .expect("root thread should start"); - session.services.agent_control = manager.agent_control(); + set_agent_control(&mut session, manager.agent_control()); session.thread_id = root.thread_id; let mut config = (*turn.config).clone(); config @@ -1126,7 +1149,7 @@ async fn multi_agent_v2_spawn_rejects_invalid_fork_turns_string() { .start_thread(StartThreadOptions::new((*turn.config).clone())) .await .expect("root thread should start"); - session.services.agent_control = manager.agent_control(); + set_agent_control(&mut session, manager.agent_control()); session.thread_id = root.thread_id; let mut config = (*turn.config).clone(); config @@ -1166,7 +1189,7 @@ async fn multi_agent_v2_spawn_rejects_zero_fork_turns() { .start_thread(StartThreadOptions::new((*turn.config).clone())) .await .expect("root thread should start"); - session.services.agent_control = manager.agent_control(); + set_agent_control(&mut session, manager.agent_control()); session.thread_id = root.thread_id; let mut config = (*turn.config).clone(); config @@ -1212,7 +1235,7 @@ async fn multi_agent_v2_send_message_accepts_root_target_from_child() { .start_thread(StartThreadOptions::new((*turn.config).clone())) .await .expect("root thread should start"); - session.services.agent_control = manager.agent_control(); + set_agent_control(&mut session, manager.agent_control()); session.thread_id = root.thread_id; let child_path = AgentPath::try_from("/root/worker").expect("agent path"); @@ -1288,7 +1311,7 @@ async fn multi_agent_v2_followup_task_rejects_root_target_from_child() { .start_thread(StartThreadOptions::new((*turn.config).clone())) .await .expect("root thread should start"); - session.services.agent_control = manager.agent_control(); + set_agent_control(&mut session, manager.agent_control()); session.thread_id = root.thread_id; let child_path = AgentPath::try_from("/root/worker").expect("agent path"); @@ -1364,7 +1387,7 @@ async fn multi_agent_v2_list_agents_returns_completed_status() { .start_thread(StartThreadOptions::new((*turn.config).clone())) .await .expect("root thread should start"); - session.services.agent_control = manager.agent_control(); + set_agent_control(&mut session, manager.agent_control()); session.thread_id = root.thread_id; let mut config = (*turn.config).clone(); let _ = config.features.enable(Feature::MultiAgentV2); @@ -1452,7 +1475,7 @@ async fn multi_agent_v2_list_agents_filters_by_relative_path_prefix() { .start_thread(StartThreadOptions::new((*turn.config).clone())) .await .expect("root thread should start"); - session.services.agent_control = manager.agent_control(); + set_agent_control(&mut session, manager.agent_control()); session.thread_id = root.thread_id; let researcher_path = AgentPath::from_string("/root/researcher".to_string()).expect("path"); @@ -1533,7 +1556,7 @@ async fn multi_agent_v2_list_agents_omits_closed_agents() { .start_thread(StartThreadOptions::new((*turn.config).clone())) .await .expect("root thread should start"); - session.services.agent_control = manager.agent_control(); + set_agent_control(&mut session, manager.agent_control()); session.thread_id = root.thread_id; let mut config = (*turn.config).clone(); let _ = config.features.enable(Feature::MultiAgentV2); @@ -1593,7 +1616,7 @@ async fn multi_agent_v2_list_agents_keeps_interrupted_resident_agents() { .start_thread(StartThreadOptions::new((*turn.config).clone())) .await .expect("root thread should start"); - session.services.agent_control = manager.agent_control(); + set_agent_control(&mut session, manager.agent_control()); session.thread_id = root.thread_id; let mut config = (*turn.config).clone(); let _ = config.features.enable(Feature::MultiAgentV2); @@ -1665,7 +1688,7 @@ async fn multi_agent_v2_send_message_rejects_legacy_items_field() { .start_thread(StartThreadOptions::new((*turn.config).clone())) .await .expect("root thread should start"); - session.services.agent_control = manager.agent_control(); + set_agent_control(&mut session, manager.agent_control()); session.thread_id = root.thread_id; let mut config = turn.config.as_ref().clone(); let _ = config.features.enable(Feature::MultiAgentV2); @@ -1721,7 +1744,7 @@ async fn multi_agent_v2_send_message_rejects_interrupt_parameter() { .start_thread(StartThreadOptions::new((*turn.config).clone())) .await .expect("root thread should start"); - session.services.agent_control = manager.agent_control(); + set_agent_control(&mut session, manager.agent_control()); session.thread_id = root.thread_id; let mut config = turn.config.as_ref().clone(); let _ = config.features.enable(Feature::MultiAgentV2); @@ -1801,7 +1824,7 @@ async fn multi_agent_v2_followup_task_completion_notifies_parent_on_every_turn() // Production spawn_agent calls happen after the parent turn has resolved // and stored its runtime; mirror that before using the synthetic handler. root.thread.session.new_default_turn().await; - session.services.agent_control = manager.agent_control(); + set_agent_control(&mut session, manager.agent_control()); session.thread_id = root.thread_id; let session = Arc::new(session); let turn = Arc::new(turn); @@ -1951,7 +1974,7 @@ async fn multi_agent_v2_followup_task_rejects_legacy_items_field() { .start_thread(StartThreadOptions::new((*turn.config).clone())) .await .expect("root thread should start"); - session.services.agent_control = manager.agent_control(); + set_agent_control(&mut session, manager.agent_control()); session.thread_id = root.thread_id; let mut config = turn.config.as_ref().clone(); let _ = config.features.enable(Feature::MultiAgentV2); @@ -2004,7 +2027,7 @@ async fn multi_agent_v2_interrupted_turn_does_not_notify_parent() { .start_thread(StartThreadOptions::new((*turn.config).clone())) .await .expect("root thread should start"); - session.services.agent_control = manager.agent_control(); + set_agent_control(&mut session, manager.agent_control()); session.thread_id = root.thread_id; let mut config = turn.config.as_ref().clone(); let _ = config.features.enable(Feature::MultiAgentV2); @@ -2081,7 +2104,7 @@ async fn multi_agent_v2_spawn_omits_agent_id_when_named() { .start_thread(StartThreadOptions::new((*turn.config).clone())) .await .expect("root thread should start"); - session.services.agent_control = manager.agent_control(); + set_agent_control(&mut session, manager.agent_control()); session.thread_id = root.thread_id; let mut config = (*turn.config).clone(); config @@ -2120,7 +2143,7 @@ async fn multi_agent_v2_spawn_surfaces_task_name_validation_errors() { .start_thread(StartThreadOptions::new((*turn.config).clone())) .await .expect("root thread should start"); - session.services.agent_control = manager.agent_control(); + set_agent_control(&mut session, manager.agent_control()); session.thread_id = root.thread_id; let mut config = (*turn.config).clone(); config @@ -2172,7 +2195,7 @@ async fn spawn_agent_reapplies_runtime_sandbox_after_role_config() { turn.config.codex_home.to_path_buf(), Arc::clone(&environment_manager), ); - session.services.agent_control = manager.agent_control(); + set_agent_control(&mut session, manager.agent_control()); let expected_sandbox = turn.config.legacy_sandbox_policy(); #[allow(deprecated)] let mut expected_file_system_sandbox_policy = @@ -2298,7 +2321,7 @@ async fn spawn_agent_reapplies_runtime_sandbox_after_role_config() { async fn spawn_agent_rejects_when_depth_limit_exceeded() { let (mut session, mut turn) = make_session_and_context().await; let manager = thread_manager(); - session.services.agent_control = manager.agent_control(); + set_agent_control(&mut session, manager.agent_control()); let max_depth = turn.config.agent_max_depth; turn.session_source = SessionSource::SubAgent(SubAgentSource::ThreadSpawn { @@ -2336,7 +2359,7 @@ async fn spawn_agent_allows_depth_up_to_configured_max_depth() { let (mut session, mut turn) = make_session_and_context().await; let manager = thread_manager(); - session.services.agent_control = manager.agent_control(); + set_agent_control(&mut session, manager.agent_control()); let mut config = (*turn.config).clone(); config.agent_max_depth = DEFAULT_AGENT_MAX_DEPTH + 1; @@ -2392,7 +2415,7 @@ async fn multi_agent_v2_spawn_agent_ignores_configured_max_depth() { .start_thread(StartThreadOptions::new(config.clone())) .await .expect("root thread should start"); - session.services.agent_control = manager.agent_control(); + set_agent_control(&mut session, manager.agent_control()); session.thread_id = root.thread_id; set_turn_config(&mut turn, config); let parent_path = AgentPath::try_from("/root/parent").expect("agent path"); @@ -2490,7 +2513,7 @@ async fn send_input_rejects_invalid_id() { async fn send_input_reports_missing_agent() { let (mut session, turn) = make_session_and_context().await; let manager = thread_manager(); - session.services.agent_control = manager.agent_control(); + set_agent_control(&mut session, manager.agent_control()); let agent_id = ThreadId::new(); let invocation = invocation( Arc::new(session), @@ -2511,7 +2534,7 @@ async fn send_input_reports_missing_agent() { async fn send_input_interrupts_before_prompt() { let (mut session, turn) = make_session_and_context().await; let manager = thread_manager(); - session.services.agent_control = manager.agent_control(); + set_agent_control(&mut session, manager.agent_control()); let config = turn.config.as_ref().clone(); let thread = manager .start_thread(StartThreadOptions::new(config.clone())) @@ -2560,7 +2583,7 @@ async fn send_input_interrupts_before_prompt() { async fn send_input_accepts_structured_items() { let (mut session, turn) = make_session_and_context().await; let manager = thread_manager(); - session.services.agent_control = manager.agent_control(); + set_agent_control(&mut session, manager.agent_control()); let config = turn.config.as_ref().clone(); let thread = manager .start_thread(StartThreadOptions::new(config.clone())) @@ -2628,7 +2651,7 @@ async fn resume_agent_rejects_invalid_id() { async fn resume_agent_reports_missing_agent() { let (mut session, turn) = make_session_and_context().await; let manager = thread_manager(); - session.services.agent_control = manager.agent_control(); + set_agent_control(&mut session, manager.agent_control()); let agent_id = ThreadId::new(); let invocation = invocation( Arc::new(session), @@ -2649,7 +2672,7 @@ async fn resume_agent_reports_missing_agent() { async fn resume_agent_noops_for_active_agent() { let (mut session, turn) = make_session_and_context().await; let manager = thread_manager(); - session.services.agent_control = manager.agent_control(); + set_agent_control(&mut session, manager.agent_control()); let config = turn.config.as_ref().clone(); let thread = manager .start_thread(StartThreadOptions::new(config.clone())) @@ -2688,7 +2711,7 @@ async fn resume_agent_noops_for_active_agent() { async fn resume_agent_restores_closed_agent_and_accepts_send_input() { let (mut session, turn) = make_session_and_context().await; let manager = thread_manager(); - session.services.agent_control = manager.agent_control(); + set_agent_control(&mut session, manager.agent_control()); let config = turn.config.as_ref().clone(); let thread = manager .resume_thread_with_history( @@ -2771,7 +2794,7 @@ async fn resume_agent_restores_closed_agent_and_accepts_send_input() { async fn resume_agent_rejects_when_depth_limit_exceeded() { let (mut session, mut turn) = make_session_and_context().await; let manager = thread_manager(); - session.services.agent_control = manager.agent_control(); + set_agent_control(&mut session, manager.agent_control()); let max_depth = turn.config.agent_max_depth; turn.session_source = SessionSource::SubAgent(SubAgentSource::ThreadSpawn { @@ -2864,7 +2887,7 @@ async fn multi_agent_v2_wait_agent_accepts_timeout_only_argument() { .start_thread(StartThreadOptions::new((*turn.config).clone())) .await .expect("root thread should start"); - session.services.agent_control = manager.agent_control(); + set_agent_control(&mut session, manager.agent_control()); session.thread_id = root.thread_id; let mut config = (*turn.config).clone(); config @@ -3194,7 +3217,7 @@ async fn multi_agent_v2_wait_agent_accepts_explicit_timeout_at_configured_max() async fn wait_agent_returns_not_found_for_missing_agents() { let (mut session, turn) = make_session_and_context().await; let manager = thread_manager(); - session.services.agent_control = manager.agent_control(); + set_agent_control(&mut session, manager.agent_control()); let id_a = ThreadId::new(); let id_b = ThreadId::new(); let invocation = invocation( @@ -3230,7 +3253,7 @@ async fn wait_agent_returns_not_found_for_missing_agents() { async fn wait_agent_times_out_when_status_is_not_final() { let (mut session, turn) = make_session_and_context().await; let manager = thread_manager(); - session.services.agent_control = manager.agent_control(); + set_agent_control(&mut session, manager.agent_control()); let config = turn.config.as_ref().clone(); let thread = manager .start_thread(StartThreadOptions::new(config.clone())) @@ -3273,7 +3296,7 @@ async fn wait_agent_times_out_when_status_is_not_final() { async fn wait_agent_clamps_short_timeouts_to_minimum() { let (mut session, turn) = make_session_and_context().await; let manager = thread_manager(); - session.services.agent_control = manager.agent_control(); + set_agent_control(&mut session, manager.agent_control()); let config = turn.config.as_ref().clone(); let thread = manager .start_thread(StartThreadOptions::new(config.clone())) @@ -3311,7 +3334,7 @@ async fn wait_agent_clamps_short_timeouts_to_minimum() { async fn wait_agent_returns_final_status_without_timeout() { let (mut session, turn) = make_session_and_context().await; let manager = thread_manager(); - session.services.agent_control = manager.agent_control(); + set_agent_control(&mut session, manager.agent_control()); let config = turn.config.as_ref().clone(); let thread = manager .start_thread(StartThreadOptions::new(config.clone())) @@ -3363,7 +3386,7 @@ async fn multi_agent_v2_wait_agent_returns_summary_for_mailbox_activity() { .start_thread(StartThreadOptions::new((*turn.config).clone())) .await .expect("root thread should start"); - session.services.agent_control = manager.agent_control(); + set_agent_control(&mut session, manager.agent_control()); session.thread_id = root.thread_id; let mut config = (*turn.config).clone(); config @@ -3456,7 +3479,7 @@ async fn multi_agent_v2_wait_agent_returns_for_already_queued_mail() { .start_thread(StartThreadOptions::new((*turn.config).clone())) .await .expect("root thread should start"); - session.services.agent_control = manager.agent_control(); + set_agent_control(&mut session, manager.agent_control()); session.thread_id = root.thread_id; let mut config = (*turn.config).clone(); config @@ -3540,7 +3563,7 @@ async fn multi_agent_v2_wait_agent_wakes_on_any_mailbox_notification() { .start_thread(StartThreadOptions::new((*turn.config).clone())) .await .expect("root thread should start"); - session.services.agent_control = manager.agent_control(); + set_agent_control(&mut session, manager.agent_control()); session.thread_id = root.thread_id; let mut config = (*turn.config).clone(); config @@ -3634,7 +3657,7 @@ async fn multi_agent_v2_wait_agent_does_not_return_completed_content() { .start_thread(StartThreadOptions::new((*turn.config).clone())) .await .expect("root thread should start"); - session.services.agent_control = manager.agent_control(); + set_agent_control(&mut session, manager.agent_control()); session.thread_id = root.thread_id; let mut config = (*turn.config).clone(); config @@ -3726,7 +3749,7 @@ async fn multi_agent_v2_interrupt_agent_accepts_task_name_target() { .start_thread(StartThreadOptions::new((*turn.config).clone())) .await .expect("root thread should start"); - session.services.agent_control = manager.agent_control(); + set_agent_control(&mut session, manager.agent_control()); session.thread_id = root.thread_id; let mut config = (*turn.config).clone(); config @@ -3852,7 +3875,7 @@ async fn multi_agent_v2_interrupt_agent_accepts_unloaded_task_name_target() { .start_thread(StartThreadOptions::new(config.clone())) .await .expect("root thread should start"); - session.services.agent_control = manager.agent_control(); + set_agent_control(&mut session, manager.agent_control()); session.thread_id = root.thread_id; set_turn_config(&mut turn, config.clone()); @@ -3943,7 +3966,7 @@ async fn multi_agent_v2_interrupt_agent_rejects_root_target_and_id() { .start_thread(StartThreadOptions::new((*turn.config).clone())) .await .expect("root thread should start"); - session.services.agent_control = manager.agent_control(); + set_agent_control(&mut session, manager.agent_control()); session.thread_id = root.thread_id; let mut config = (*turn.config).clone(); config @@ -3999,7 +4022,7 @@ async fn multi_agent_v2_interrupt_agent_rejects_self_target_by_id() { .start_thread(StartThreadOptions::new((*turn.config).clone())) .await .expect("root thread should start"); - session.services.agent_control = manager.agent_control(); + set_agent_control(&mut session, manager.agent_control()); session.thread_id = root.thread_id; let child_path = AgentPath::try_from("/root/worker").expect("agent path"); @@ -4066,7 +4089,7 @@ async fn multi_agent_v2_interrupt_agent_rejects_self_target_by_task_name() { .start_thread(StartThreadOptions::new((*turn.config).clone())) .await .expect("root thread should start"); - session.services.agent_control = manager.agent_control(); + set_agent_control(&mut session, manager.agent_control()); session.thread_id = root.thread_id; let child_path = AgentPath::try_from("/root/worker").expect("agent path"); @@ -4123,7 +4146,7 @@ async fn multi_agent_v2_interrupt_agent_rejects_self_target_by_task_name() { async fn close_agent_submits_shutdown_and_returns_previous_status() { let (mut session, turn) = make_session_and_context().await; let manager = thread_manager(); - session.services.agent_control = manager.agent_control(); + set_agent_control(&mut session, manager.agent_control()); let config = turn.config.as_ref().clone(); let thread = manager .start_thread(StartThreadOptions::new(config.clone()))