From 5c07856e25b565c86c6266ff51edff58299fd4e6 Mon Sep 17 00:00:00 2001 From: jif Date: Tue, 22 Sep 2026 14:38:02 +0000 Subject: [PATCH] Separate local agent runtime state from controller handles (#47301) ## What changed Extract `LocalAgentRuntime` from `LocalAgentControl` to hold the shared registry, residency state, execution limiter, rollout budget, service tier, instruction provider, and thread-manager bindings. Keep session identity on the controller handle. Store the runtime in `SessionServices` and use it to bind controllers for child startup, resume, delegation, and guardian reviews. Route ownership checks and residency handling through the shared runtime, preserving registry identity across handles. Update existing test fixtures to keep the controller and runtime paired. GitOrigin-RevId: 846094657cbdf736da3149f0a1c96b4e01423dfe --- codex-rs/core/src/agent/control.rs | 84 ++++------- codex-rs/core/src/agent/control/budget.rs | 9 +- codex-rs/core/src/agent/control/execution.rs | 6 +- .../core/src/agent/control/execution_tests.rs | 6 +- codex-rs/core/src/agent/control/legacy.rs | 2 +- codex-rs/core/src/agent/control/residency.rs | 21 +-- codex-rs/core/src/agent/control/runtime.rs | 69 +++++++++ .../core/src/agent/control/service_tier.rs | 7 +- codex-rs/core/src/agent/control/spawn.rs | 50 +++++-- codex-rs/core/src/agent/control/target.rs | 2 +- .../src/agent/control/user_authorization.rs | 5 +- codex-rs/core/src/agent/control_tests.rs | 8 +- codex-rs/core/src/codex_delegate.rs | 5 +- .../core/src/guardian/review_session_setup.rs | 6 +- codex-rs/core/src/session/session.rs | 1 + codex-rs/core/src/session/tests.rs | 2 + codex-rs/core/src/state/service.rs | 2 + codex-rs/core/src/thread_manager.rs | 14 +- .../src/tools/handlers/multi_agents_tests.rs | 135 ++++++++++-------- 19 files changed, 285 insertions(+), 149 deletions(-) create mode 100644 codex-rs/core/src/agent/control/runtime.rs 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()))