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/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");