From fdbe23da4556dcce488e5392a9f187734ba850da Mon Sep 17 00:00:00 2001 From: jif Date: Thu, 24 Sep 2026 10:28:14 +0000 Subject: [PATCH] Allow hosts to provide agent controllers through ThreadManager (#47808) ## What changed - Add `ThreadManager::with_agent_control_factory` to select a host-owned controller asynchronously using the final thread ID for new and cold-resumed threads. Keep local control as the default, reuse live controllers on warm resume, and preserve parent controllers for internal sessions. - Fail startup on factory errors or when a provided controller's identity does not match the session identity. - Defer root service-tier updates to host controllers until thread registration succeeds, preventing failed starts from changing settings on a shared controller. GitOrigin-RevId: 095afbecfee5d47329038d2a97d0cc0081f181df --- codex-rs/core/src/agent/control/runtime.rs | 10 +-- codex-rs/core/src/codex_delegate.rs | 2 +- .../core/src/guardian/review_session_setup.rs | 2 +- codex-rs/core/src/session/session.rs | 31 +++++--- codex-rs/core/src/thread_manager.rs | 73 ++++++++++++++++++- codex-rs/core/src/thread_manager_tests.rs | 2 +- 6 files changed, 97 insertions(+), 23 deletions(-) diff --git a/codex-rs/core/src/agent/control/runtime.rs b/codex-rs/core/src/agent/control/runtime.rs index 71f8b4f67d..287883785f 100644 --- a/codex-rs/core/src/agent/control/runtime.rs +++ b/codex-rs/core/src/agent/control/runtime.rs @@ -70,12 +70,12 @@ impl LocalAgentRuntime { } } -/// Local construction binds identity after reading history; internal children inherit -/// an already-bound controller without selecting a backend again. +/// Local construction binds identity after reading history. Hosts and internal children +/// provide an already-bound controller without selecting a backend again. #[derive(Clone)] pub(crate) enum AgentControlInit { Local(LocalAgentControl), - Inherited { + Provided { control: Arc, runtime: LocalAgentRuntime, }, @@ -91,14 +91,14 @@ impl AgentControlInit { pub(crate) fn runtime(&self) -> &LocalAgentRuntime { match self { Self::Local(control) => &control.runtime, - Self::Inherited { runtime, .. } => runtime, + Self::Provided { runtime, .. } => runtime, } } pub(crate) fn control(&self) -> &dyn AgentControl { match self { Self::Local(control) => control, - Self::Inherited { control, .. } => control.as_ref(), + Self::Provided { control, .. } => control.as_ref(), } } } diff --git a/codex-rs/core/src/codex_delegate.rs b/codex-rs/core/src/codex_delegate.rs index 4a792caebd..9d6b9656ba 100644 --- a/codex-rs/core/src/codex_delegate.rs +++ b/codex-rs/core/src/codex_delegate.rs @@ -114,7 +114,7 @@ pub(crate) async fn run_codex_thread_interactive( ThreadSource::Subagent }), originator: parent_ctx.originator.clone(), - agent_control: crate::agent::control::AgentControlInit::Inherited { + agent_control: crate::agent::control::AgentControlInit::Provided { control: Arc::clone(&parent_session.services.agent_control), runtime: parent_session.services.local_agent_runtime.clone(), }, diff --git a/codex-rs/core/src/guardian/review_session_setup.rs b/codex-rs/core/src/guardian/review_session_setup.rs index 06ab9eaa2e..92cc2487bb 100644 --- a/codex-rs/core/src/guardian/review_session_setup.rs +++ b/codex-rs/core/src/guardian/review_session_setup.rs @@ -102,7 +102,7 @@ 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: crate::agent::control::AgentControlInit::Inherited { + agent_control: crate::agent::control::AgentControlInit::Provided { control: Arc::clone(&self.parent.services.agent_control), runtime: self.parent.services.local_agent_runtime.clone(), }, diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index aa04cbb324..3f9cd15720 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -834,17 +834,6 @@ impl Session { .parent_thread_id .or_else(|| initial_history.get_resumed_parent_thread_id()); session_configuration.parent_thread_id = parent_thread_id; - if parent_thread_id.is_none() { - agent_control - .control() - .propagate_config_update(AgentConfigUpdate::ServiceTier( - session_configuration - .step_settings - .service_tier - .clone() - .or_else(|| config.service_tier.clone()), - )); - } let is_paginated_subagent = matches!( session_configuration.history_mode, ThreadHistoryMode::Paginated @@ -945,10 +934,28 @@ impl Session { .effective_agent_max_threads(MultiAgentVersion::V2) .unwrap_or(usize::MAX), ); + if parent_thread_id.is_none() { + control.propagate_config_update(AgentConfigUpdate::ServiceTier( + session_configuration + .step_settings + .service_tier + .clone() + .or_else(|| config.service_tier.clone()), + )); + } let runtime = control.runtime.clone(); (Arc::new(control), runtime) } - AgentControlInit::Inherited { control, runtime } => (control, runtime), + AgentControlInit::Provided { control, runtime } => { + let controller_id = control.identity(); + if controller_id != session_id { + return Err(CodexErr::InvalidRequest(format!( + "agent controller identity {controller_id} does not match session identity {session_id}" + )) + .into()); + } + (control, runtime) + } }; let time_provider = crate::current_time::resolve_time_provider( config.current_time_reminder.as_ref(), diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 09ad357867..82c0a11121 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -3,6 +3,8 @@ mod shared_instructions; use crate::CodexAppsToolsCache; use crate::agent::LocalAgentControl; +use crate::agent::api::AgentConfigUpdate; +use crate::agent::api::AgentControl; use crate::agent::control::AgentControlInit; use crate::agents_md_manager::SessionInstructions; use crate::attestation::AttestationProvider; @@ -96,9 +98,11 @@ use codex_thread_store::UpdateThreadMetadataParams; use codex_utils_absolute_path::AbsolutePathBuf; use codex_utils_git_discovery::GitRootDiscovery; use futures::StreamExt; +use futures::future::BoxFuture; use futures::stream::FuturesUnordered; use std::collections::HashMap; use std::collections::HashSet; +use std::future::Future; use std::path::PathBuf; use std::sync::Arc; use std::sync::atomic::AtomicBool; @@ -121,6 +125,8 @@ static FORCE_TEST_THREAD_MANAGER_BEHAVIOR: AtomicBool = AtomicBool::new(false); type CapturedOps = Vec<(ThreadId, Op)>; type SharedCapturedOps = Arc>; pub(crate) type ThreadIdGenerator = Arc ThreadId + Send + Sync>; +type AgentControlFactory = + Arc BoxFuture<'static, CodexResult>> + Send + Sync>; // `Op` is intentionally not `Clone`. Thread-manager tests only snapshot the // small subset of ops they inspect. @@ -399,6 +405,7 @@ pub(crate) struct ThreadManagerState { shared_thread_instructions: shared_instructions::SharedThreadInstructionsProviders, thread_created_tx: broadcast::Sender, thread_id_generator: ThreadIdGenerator, + agent_control_factory: Option, auth_manager: Arc, models_manager: SharedModelsManager, git_root_discovery: Arc, @@ -561,6 +568,7 @@ impl ThreadManager { shared_thread_instructions: Default::default(), thread_created_tx, thread_id_generator: default_thread_id_generator(), + agent_control_factory: None, models_manager, git_root_discovery: Arc::default(), environment_manager, @@ -599,6 +607,28 @@ impl ThreadManager { self } + /// Selects the host-owned agent controller for every new or cold-resumed thread. + /// + /// The factory receives the thread's final ID and must return the controller for + /// its agent tree. New roots use their thread ID as the controller identity; + /// resumed threads must match their persisted session identity. Factory errors + /// fail startup. Internal sessions inherit their parent's controller, and warm + /// resume retains the live controller without calling the factory. + /// + /// Install this before starting threads. Without a factory, the manager uses + /// the local backend. Host backends support MultiAgentV2; V1 requires local control. + pub fn with_agent_control_factory(mut self, factory: F) -> Self + where + F: Fn(ThreadId) -> Fut + Send + Sync + 'static, + Fut: Future>> + Send + 'static, + { + let Some(state) = Arc::get_mut(&mut self.state) else { + unreachable!("agent controller factory must be set before thread manager is shared"); + }; + state.agent_control_factory = Some(Arc::new(move |thread_id| Box::pin(factory(thread_id)))); + self + } + /// Replaces the process-wide provider before this manager is shared with threads. pub fn with_code_mode_session_provider( mut self, @@ -710,6 +740,7 @@ impl ThreadManager { shared_thread_instructions: Default::default(), thread_created_tx, thread_id_generator: default_thread_id_generator(), + agent_control_factory: None, models_manager: create_model_provider(provider, Some(auth_manager.clone())) .models_manager(codex_home, /*config_model_catalog*/ None), git_root_discovery: Arc::default(), @@ -1060,7 +1091,7 @@ impl ThreadManager { options.internal_parent = Some(InternalSessionParent { thread_id: parent_thread_id, auth_manager: Arc::clone(&parent.session.services.auth_manager), - agent_control: AgentControlInit::Inherited { + agent_control: AgentControlInit::Provided { control: Arc::clone(&parent.session.services.agent_control), runtime: parent.session.services.local_agent_runtime.clone(), }, @@ -2028,7 +2059,7 @@ impl ThreadManagerState { user_instructions: supplied_user_instructions, mut thread_extension_init, client_mcp_extensions, - reserved_thread_id, + mut reserved_thread_id, disabled_plugin_ids, } = options; let inherited_environments = captured_environments.or(inherited_environments); @@ -2094,6 +2125,35 @@ impl ThreadManagerState { threads.remove(&resumed.conversation_id); } } + let agent_control = match (&self.agent_control_factory, agent_control) { + (Some(factory), AgentControlInit::Local(local)) => { + let thread_id = match &initial_history { + InitialHistory::Resumed(resumed) => resumed.conversation_id, + InitialHistory::New | InitialHistory::Cleared | InitialHistory::Forked(_) => { + *reserved_thread_id + .get_or_insert_with(|| local.runtime.generate_thread_id()) + } + }; + AgentControlInit::Provided { + control: factory(thread_id).await?, + runtime: local.runtime, + } + } + (_, control) => control, + }; + let parent_thread_id = parent_thread_id + .or_else(|| initial_history.get_resumed_parent_thread_id()) + .or(match &session_source { + SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id, .. + }) => Some(*parent_thread_id), + _ => None, + }); + // Host controllers can already be shared with live threads. Publish root settings + // only after registration succeeds, so a failed start cannot update their tree. + let initial_host_config = (self.agent_control_factory.is_some() + && parent_thread_id.is_none()) + .then(|| AgentConfigUpdate::ServiceTier(config.service_tier.clone())); // Both resume entry points must restore identities before children can be loaded lazily. if let InitialHistory::Resumed(resumed) = &initial_history && initial_history.get_multi_agent_version() == Some(MultiAgentVersion::V2) @@ -2267,7 +2327,7 @@ impl ThreadManagerState { session.services.mcp_runtime.enable_full_access_form_input(); } let new_thread = self - .finalize_thread_spawn(session, io, tracked_session_source) + .finalize_thread_spawn(session, io, tracked_session_source, initial_host_config) .await?; new_thread.thread.emit_thread_ready_lifecycle().await; if source_changed_during_startup.load(Ordering::Acquire) { @@ -2284,6 +2344,7 @@ impl ThreadManagerState { session: Arc, io: SessionIo, session_source: SessionSource, + initial_host_config: Option, ) -> CodexResult { let thread_id = session.thread_id(); let event = io.next_event().await?; @@ -2300,6 +2361,12 @@ impl ThreadManagerState { { let mut threads = self.threads.write().await; if let std::collections::hash_map::Entry::Vacant(e) = threads.entry(thread_id) { + if let Some(update) = initial_host_config { + session + .services + .agent_control + .propagate_config_update(update); + } let thread = Arc::new(CodexThread::new( session, io, diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index abbaa86d0d..81d74e636e 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -1504,7 +1504,7 @@ async fn spawn_internal_session_preserves_parent_lineage_without_forking_history internal_parent: Some(InternalSessionParent { thread_id: parent.thread_id, auth_manager: Arc::clone(&parent.thread.session.services.auth_manager), - agent_control: AgentControlInit::Inherited { + agent_control: AgentControlInit::Provided { control: Arc::clone(&parent.thread.session.services.agent_control), runtime: parent.thread.session.services.local_agent_runtime.clone(), },