github/codex
4
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-22 01:34:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

Allow hosts to provide agent controllers through ThreadManager (#47808)

Browse Source 
## 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
This commit is contained in:
jif
2026-09-24 10:39:10 +00:00
committed by copyberry
parent c098f97e53
commit fdbe23da45
6 changed files with 97 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/agent/control/runtime.rs
+5 -5
View File
@@ -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<dyn AgentControl>,
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(),
}
}
}
codex-rs/core/src/codex_delegate.rs
+1 -1
View File
@@ -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(),
},
codex-rs/core/src/guardian/review_session_setup.rs
+1 -1
View File
@@ -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(),
},
codex-rs/core/src/session/session.rs
+19 -12
View File
@@ -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(),
codex-rs/core/src/thread_manager.rs
+70 -3
View File
@@ -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<std::sync::Mutex<CapturedOps>>;
pub(crate) type ThreadIdGenerator = Arc<dyn Fn() -> ThreadId + Send + Sync>;
type AgentControlFactory =
Arc<dyn Fn(ThreadId) -> BoxFuture<'static, CodexResult<Arc<dyn AgentControl>>> + 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<ThreadId>,
thread_id_generator: ThreadIdGenerator,
agent_control_factory: Option<AgentControlFactory>,
auth_manager: Arc<AuthManager>,
models_manager: SharedModelsManager,
git_root_discovery: Arc<GitRootDiscovery>,
@@ -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<F, Fut>(mut self, factory: F) -> Self
where
F: Fn(ThreadId) -> Fut + Send + Sync + 'static,
Fut: Future<Output = CodexResult<Arc<dyn AgentControl>>> + 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<Session>,
io: SessionIo,
session_source: SessionSource,
initial_host_config: Option<AgentConfigUpdate>,
) -> CodexResult<NewThread> {
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,
codex-rs/core/src/thread_manager_tests.rs
+1 -1
View File
@@ -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(),
},