From 9457f8f73412b22b26ea9cf6b8e8bcf625b381c8 Mon Sep 17 00:00:00 2001 From: keith thornhill Date: Fri, 7 Aug 2026 01:22:58 +0000 Subject: [PATCH] Allow `ThreadManager` to customize thread ID generation (#37350) ## What changed - Add `ThreadManager::with_thread_id_generator` to configure ID allocation for new root, child, and forked threads. - Keep the standard UUIDv7 generator as the default and preserve stored IDs when resuming threads. - Add `ThreadId::from_u128` for constructing identifiers from UUID-compatible values. ## Testing - Cover default allocation, custom IDs across roots, child agents, and forks, and ID preservation on resume. GitOrigin-RevId: 6c38c862e30a50bad944a994bcbccf7d1e7ada24 --- codex-rs/core/src/agent/control.rs | 28 ++++- codex-rs/core/src/session/session.rs | 2 +- codex-rs/core/src/thread_manager.rs | 32 ++++- codex-rs/core/src/thread_manager_tests.rs | 146 ++++++++++++++++++++++ codex-rs/protocol/src/thread_id.rs | 8 ++ 5 files changed, 211 insertions(+), 5 deletions(-) diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index 76a7d19b8b..1ba0317294 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/codex-rs/core/src/agent/control.rs @@ -16,7 +16,9 @@ use crate::session_prefix::format_inter_agent_completion_message; use crate::session_prefix::format_subagent_context_line; use crate::session_prefix::format_subagent_notification_message; use crate::thread_manager::ResumeThreadWithHistoryOptions; +use crate::thread_manager::ThreadIdGenerator; 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 codex_protocol::AgentPath; use codex_protocol::SessionId; @@ -93,7 +95,7 @@ pub(crate) struct ListedAgent { /// An `AgentControl` instance is intended to be created at most once per root thread/session /// tree. That same `AgentControl` 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`. -#[derive(Clone, Default)] +#[derive(Clone)] pub(crate) struct AgentControl { /// ID shared by the whole agent control session. This means every sub-agents from a common /// root share the same session ID. @@ -102,6 +104,8 @@ pub(crate) struct AgentControl { /// 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, @@ -109,15 +113,31 @@ pub(crate) struct AgentControl { rollout_budget: Arc, } +impl Default for AgentControl { + fn default() -> Self { + Self::new( + Weak::default(), + default_thread_id_generator(), + /*rollout_budget*/ None, + ) + } +} + impl AgentControl { /// Construct a new `AgentControl` that can spawn/message agents via the given manager state. pub(crate) fn new( manager: Weak, + thread_id_generator: ThreadIdGenerator, rollout_budget: Option, ) -> Self { let control = Self { + session_id: SessionId::default(), manager, - ..Default::default() + thread_id_generator, + state: Arc::default(), + v2_residency: Arc::default(), + agent_execution_limiter: Arc::default(), + rollout_budget: Arc::default(), }; if let Some(rollout_budget) = rollout_budget { control.rollout_budget.configure(rollout_budget); @@ -135,6 +155,10 @@ impl AgentControl { self.session_id } + pub(crate) fn generate_thread_id(&self) -> ThreadId { + (self.thread_id_generator)() + } + pub(crate) fn rollout_budget(&self) -> &RolloutBudget { self.rollout_budget.as_ref() } diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 265765eaa5..f390596bb6 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -569,7 +569,7 @@ impl Session { let thread_id = match &initial_history { InitialHistory::New | InitialHistory::Cleared | InitialHistory::Forked(_) => { - ThreadId::default() + agent_control.generate_thread_id() } InitialHistory::Resumed(resumed_history) => resumed_history.conversation_id, }; diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index 8053df8a34..f30f1f86b9 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -107,6 +107,11 @@ 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>; + +pub(crate) fn default_thread_id_generator() -> ThreadIdGenerator { + Arc::new(ThreadId::new) +} pub(crate) fn set_thread_manager_test_mode_for_tests(enabled: bool) { FORCE_TEST_THREAD_MANAGER_BEHAVIOR.store(enabled, Ordering::Relaxed); @@ -311,6 +316,7 @@ pub(crate) struct ResumeThreadWithHistoryOptions { pub(crate) struct ThreadManagerState { threads: Arc>>>, thread_created_tx: broadcast::Sender, + thread_id_generator: ThreadIdGenerator, auth_manager: Arc, models_manager: SharedModelsManager, environment_manager: Arc, @@ -437,6 +443,7 @@ impl ThreadManager { state: Arc::new(ThreadManagerState { threads: Arc::new(RwLock::new(HashMap::new())), thread_created_tx, + thread_id_generator: default_thread_id_generator(), models_manager, environment_manager, starting_mcp_runtimes: std::sync::Mutex::new(Vec::new()), @@ -461,6 +468,18 @@ impl ThreadManager { } } + /// Generate every new thread identifier with the caller-provided factory. + pub fn with_thread_id_generator( + mut self, + generator: impl Fn() -> ThreadId + Send + Sync + 'static, + ) -> Self { + let Some(state) = Arc::get_mut(&mut self.state) else { + unreachable!("thread ID generator must be set before thread manager is shared"); + }; + state.thread_id_generator = Arc::new(generator); + self + } + /// Replaces the process-wide provider before this manager is shared with threads. pub fn with_code_mode_session_provider( mut self, @@ -569,6 +588,7 @@ impl ThreadManager { state: Arc::new(ThreadManagerState { threads: Arc::new(RwLock::new(HashMap::new())), thread_created_tx, + thread_id_generator: default_thread_id_generator(), models_manager: create_model_provider(provider, Some(auth_manager.clone())) .models_manager(codex_home, /*config_model_catalog*/ None), environment_manager, @@ -1197,11 +1217,19 @@ impl ThreadManager { } pub(crate) fn agent_control(&self) -> AgentControl { - AgentControl::new(Arc::downgrade(&self.state), /*rollout_budget*/ None) + AgentControl::new( + Arc::downgrade(&self.state), + self.state.thread_id_generator.clone(), + /*rollout_budget*/ None, + ) } fn agent_control_for_config(&self, config: &Config) -> AgentControl { - AgentControl::new(Arc::downgrade(&self.state), config.rollout_budget.clone()) + AgentControl::new( + Arc::downgrade(&self.state), + self.state.thread_id_generator.clone(), + config.rollout_budget.clone(), + ) } #[cfg(test)] diff --git a/codex-rs/core/src/thread_manager_tests.rs b/codex-rs/core/src/thread_manager_tests.rs index 3e1c0422aa..fcca8b6036 100644 --- a/codex-rs/core/src/thread_manager_tests.rs +++ b/codex-rs/core/src/thread_manager_tests.rs @@ -1,4 +1,5 @@ use super::*; +use crate::agent::control::SpawnAgentOptions; use crate::config::test_config; use crate::init_state_db; use crate::installation_id::INSTALLATION_ID_FILENAME; @@ -32,6 +33,7 @@ use codex_protocol::protocol::SubAgentSource; use codex_protocol::protocol::ThreadSource; use codex_protocol::protocol::TurnStartedEvent; use codex_protocol::protocol::UserMessageEvent; +use codex_protocol::user_input::UserInput; use codex_utils_path_uri::PathUri; use core_test_support::PathBufExt; use core_test_support::PathExt; @@ -44,6 +46,150 @@ use wiremock::MockServer; const TEST_INSTALLATION_ID: &str = "11111111-1111-4111-8111-111111111111"; +/// Controls without a custom allocation policy still produce distinct thread identifiers. +#[test] +fn thread_id_generator_defaults_to_standard_ids() { + let agent_control = AgentControl::default(); + + assert_ne!( + agent_control.generate_thread_id(), + agent_control.generate_thread_id() + ); +} + +/// One custom ID factory supplies identifiers for roots, actual child agents, and forks. +#[tokio::test] +async fn thread_id_generator_applies_to_roots_children_and_forks() { + let temp_dir = tempdir().expect("tempdir"); + let mut config = test_config().await; + config.codex_home = temp_dir.path().join("codex-home").abs(); + config.cwd = config.codex_home.abs(); + std::fs::create_dir_all(&config.codex_home).expect("create codex home"); + + let generated_ids = [ + ThreadId::from_u128(/*value*/ 0x018f_0000_0000_7000_8000_0000_0000_0001), + ThreadId::from_u128(/*value*/ 0x018f_0000_0000_7000_8000_0000_0000_0002), + ThreadId::from_u128(/*value*/ 0x018f_0000_0000_7000_8000_0000_0000_0003), + ]; + let next_id = std::sync::atomic::AtomicUsize::new(0); + let manager = ThreadManager::with_models_provider_and_home_for_tests( + CodexAuth::from_api_key("dummy"), + config.model_provider.clone(), + config.codex_home.to_path_buf(), + Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), + ) + .with_thread_id_generator(move || generated_ids[next_id.fetch_add(1, Ordering::Relaxed)]); + let root = manager + .start_thread(StartThreadOptions::new(config.clone())) + .await + .expect("start root thread"); + let child = root + .thread + .session + .services + .agent_control + .spawn_agent_with_metadata( + config.clone(), + vec![UserInput::Text { + text: "child task".to_string(), + text_elements: Vec::new(), + }], + Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id: root.thread_id, + depth: 1, + agent_path: None, + agent_nickname: None, + agent_role: None, + })), + SpawnAgentOptions { + parent_thread_id: Some(root.thread_id), + ..Default::default() + }, + ) + .await + .expect("spawn actual child agent"); + let fork = manager + .spawn_subagent(root.thread_id, StartThreadOptions::new(config)) + .await + .expect("fork root thread"); + + assert_eq!( + [root.thread_id, child.thread_id, fork.thread_id], + generated_ids + ); + + let report = manager + .shutdown_all_threads_bounded(Duration::from_secs(10)) + .await; + assert_eq!(report.completed.len(), 3); +} + +/// Resuming a thread preserves its stored ID instead of invoking the new manager's factory. +#[tokio::test] +async fn thread_id_generator_does_not_replace_resumed_thread_id() { + let temp_dir = tempdir().expect("tempdir"); + let mut config = test_config().await; + config.codex_home = temp_dir.path().join("codex-home").abs(); + config.cwd = config.codex_home.abs(); + std::fs::create_dir_all(&config.codex_home).expect("create codex home"); + + let original_thread_id = + ThreadId::from_u128(/*value*/ 0x018f_0000_0000_7000_8000_0000_0000_0001); + let original_manager = ThreadManager::with_models_provider_and_home_for_tests( + CodexAuth::from_api_key("dummy"), + config.model_provider.clone(), + config.codex_home.to_path_buf(), + Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), + ) + .with_thread_id_generator(move || original_thread_id); + let original = original_manager + .start_thread(StartThreadOptions::new(config.clone())) + .await + .expect("start source thread"); + original.thread.ensure_rollout_materialized().await; + original + .thread + .flush_rollout() + .await + .expect("flush source rollout"); + let rollout_path = original + .thread + .rollout_path() + .expect("source rollout path should exist"); + assert_eq!(original.thread_id, original_thread_id); + original + .thread + .shutdown_and_wait() + .await + .expect("shut down source thread"); + let _ = original_manager.remove_thread(&original_thread_id).await; + + let resumed_manager = ThreadManager::with_models_provider_and_home_for_tests( + CodexAuth::from_api_key("dummy"), + config.model_provider.clone(), + config.codex_home.to_path_buf(), + Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), + ) + .with_thread_id_generator(|| panic!("resuming must not allocate a new thread ID")); + let resumed = resumed_manager + .resume_thread_from_rollout( + config, + rollout_path, + Arc::clone(&resumed_manager.state.auth_manager), + /*parent_trace*/ None, + ClientMcpExtensions::default(), + ) + .await + .expect("resume existing source thread"); + + assert_eq!(resumed.thread_id, original_thread_id); + resumed + .thread + .shutdown_and_wait() + .await + .expect("shut down resumed thread"); +} + #[tokio::test] async fn child_session_inherits_client_mcp_extensions() { let temp_dir = tempdir().expect("tempdir"); diff --git a/codex-rs/protocol/src/thread_id.rs b/codex-rs/protocol/src/thread_id.rs index dcfedc7f52..4395703a4f 100644 --- a/codex-rs/protocol/src/thread_id.rs +++ b/codex-rs/protocol/src/thread_id.rs @@ -24,6 +24,13 @@ impl ThreadId { } } + /// Construct an identifier from a UUID's 128-bit representation. + pub fn from_u128(value: u128) -> Self { + Self { + uuid: Uuid::from_u128(value), + } + } + pub fn from_string(s: &str) -> Result { Ok(Self { uuid: Uuid::parse_str(s)?, @@ -98,6 +105,7 @@ impl JsonSchema for ThreadId { #[cfg(test)] mod tests { use super::*; + #[test] fn test_thread_id_default_is_not_zeroes() { let id = ThreadId::default();