openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-14 16:38:43 +00:00
Code Issues Packages Projects Releases Wiki Activity

Separate local agent runtime state from controller handles (#47301)

Browse Source 
## 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
This commit is contained in:
jif
2026-09-22 14:48:28 +00:00
committed by copyberry
parent c6c09fe29f
commit 5c07856e25
19 changed files with 285 additions and 149 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/agent/control.rs
+30 -54
View File
@@ -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<ThreadManagerState>,
/// Captured at construction so delegates retain their manager's allocation policy.
thread_id_generator: ThreadIdGenerator,
state: Arc<AgentRegistry>,
v2_residency: Arc<V2Residency>,
agent_execution_limiter: Arc<AgentExecutionLimiter>,
/// Session-scoped state shared by the root thread and every cloned sub-agent control handle.
rollout_budget: Arc<RolloutBudget>,
/// The user-selected root routing tier, shared by the entire agent tree.
root_service_tier: Arc<ArcSwapOption<String>>,
/// Retains the root's opt-in instruction provider even when the root is unloaded.
shared_thread_instructions_provider: Arc<OnceLock<Arc<dyn ThreadInstructionsProvider>>>,
pub(crate) runtime: LocalAgentRuntime,
}
impl Default for LocalAgentControl {
@@ -132,26 +111,15 @@ impl LocalAgentControl {
thread_id_generator: ThreadIdGenerator,
rollout_budget: Option<RolloutBudgetConfig>,
) -> 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<Arc<dyn ThreadInstructionsProvider>>,
) -> Option<Arc<dyn ThreadInstructionsProvider>> {
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<ThreadId>,
) {
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<AgentMetadata> {
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<AgentMetadata> {
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<String>,
) -> 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<Arc<ThreadManagerState>> {
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),
codex-rs/core/src/agent/control/budget.rs
+6 -3
View File
@@ -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<RolloutBudgetReminder> {
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);
}
}
codex-rs/core/src/agent/control/execution.rs
+3 -3
View File
@@ -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<AgentExecutionGuard> {
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())
}
}
codex-rs/core/src/agent/control/execution_tests.rs
+5 -1
View File
@@ -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()));
codex-rs/core/src/agent/control/legacy.rs
+1 -1
View File
@@ -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
}
codex-rs/core/src/agent/control/residency.rs
+12 -9
View File
@@ -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);
codex-rs/core/src/agent/control/runtime.rs
+69
View File
@@ -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<ThreadManagerState>,
/// Captured at construction so delegates retain their manager's allocation policy.
pub(super) thread_id_generator: ThreadIdGenerator,
pub(super) agent_execution_limiter: Arc<AgentExecutionLimiter>,
/// Session-scoped state shared by the root thread and every cloned sub-agent control handle.
pub(super) rollout_budget: Arc<RolloutBudget>,
/// The user-selected root routing tier, shared by the entire agent tree.
pub(super) root_service_tier: Arc<ArcSwapOption<String>>,
/// Retains the root's opt-in instruction provider even when the root is unloaded.
pub(super) shared_thread_instructions_provider:
Arc<OnceLock<Arc<dyn ThreadInstructionsProvider>>>,
pub(super) registry: Arc<AgentRegistry>,
pub(super) residency: Arc<V2Residency>,
}
impl LocalAgentRuntime {
pub(super) fn new(
manager: Weak<ThreadManagerState>,
thread_id_generator: ThreadIdGenerator,
rollout_budget: Option<RolloutBudgetConfig>,
) -> 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(),
}
}
}
codex-rs/core/src/agent/control/service_tier.rs
+5 -2
View File
@@ -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<String> {
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<String>) {
self.root_service_tier.store(service_tier.map(Arc::new));
self.runtime
.root_service_tier
.store(service_tier.map(Arc::new));
}
}
codex-rs/core/src/agent/control/spawn.rs
+38 -12
View File
@@ -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,
codex-rs/core/src/agent/control/target.rs
+1 -1
View File
@@ -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!(
codex-rs/core/src/agent/control/user_authorization.rs
+4 -1
View File
@@ -38,7 +38,10 @@ impl LocalAgentControl {
&self,
thread_id: ThreadId,
) -> Option<GuardianRootSnapshot> {
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;
}
codex-rs/core/src/agent/control_tests.rs
+7 -1
View File
@@ -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 {
codex-rs/core/src/codex_delegate.rs
+4 -1
View File
@@ -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,
codex-rs/core/src/guardian/review_session_setup.rs
+5 -1
View File
@@ -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.
codex-rs/core/src/session/session.rs
+1
View File
@@ -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,
codex-rs/core/src/session/tests.rs
+2
View File
@@ -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(),
codex-rs/core/src/state/service.rs
+2
View File
@@ -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<SelectedCapabilityRoot>,
pub(crate) mcp_thread_init: ExtensionDataInit,
pub(crate) agent_control: LocalAgentControl,
pub(crate) local_agent_runtime: LocalAgentRuntime,
pub(crate) network_proxy: ArcSwapOption<StartedNetworkProxy>,
pub(crate) network_proxy_audit_metadata: NetworkProxyAuditMetadata,
pub(crate) managed_network_requirements_configured: bool,
codex-rs/core/src/thread_manager.rs
+11 -3
View File
@@ -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 {
codex-rs/core/src/tools/handlers/multi_agents_tests.rs
+79 -56
View File
@@ -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<crate::session::session::Session>,
turn: Arc<TurnContext>,
@@ -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()))