diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index 42981a5d0a..a9a90be54a 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/codex-rs/core/src/agent/control.rs @@ -1165,24 +1165,16 @@ impl AgentControl { let state = self.upgrade()?; let mut children_by_parent = HashMap::>::new(); - for thread_id in state.list_thread_ids().await { - let Ok(thread) = state.get_thread(thread_id).await else { - continue; - }; - let snapshot = thread.config_snapshot().await; - let Some(parent_thread_id) = thread_spawn_parent_thread_id(&snapshot.session_source) - else { - continue; - }; + for (parent_thread_id, child_thread_id) in state.list_live_thread_spawn_edges().await { children_by_parent .entry(parent_thread_id) .or_default() .push(( - thread_id, + child_thread_id, self.state - .agent_metadata_for_thread(thread_id) + .agent_metadata_for_thread(child_thread_id) .unwrap_or(AgentMetadata { - agent_id: Some(thread_id), + agent_id: Some(child_thread_id), ..Default::default() }), )); diff --git a/codex-rs/core/src/agent/control_tests.rs b/codex-rs/core/src/agent/control_tests.rs index 7ff50ffef8..54114f05d1 100644 --- a/codex-rs/core/src/agent/control_tests.rs +++ b/codex-rs/core/src/agent/control_tests.rs @@ -2123,6 +2123,64 @@ async fn list_agent_subtree_thread_ids_includes_anonymous_and_closed_descendants ); } +#[tokio::test] +async fn list_agent_subtree_thread_ids_includes_live_descendants_without_state_db() { + let (_home, config) = test_config().await; + let manager = ThreadManager::with_models_provider_home_and_state_for_tests( + CodexAuth::from_api_key("dummy"), + config.model_provider.clone(), + config.codex_home.to_path_buf(), + std::sync::Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), + /*state_db*/ None, + ); + let control = manager.agent_control(); + let parent_thread_id = manager + .start_thread(config.clone()) + .await + .expect("parent should start") + .thread_id; + + let child_thread_id = control + .spawn_agent( + config.clone(), + text_input("hello child"), + Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id, + depth: 1, + agent_path: None, + agent_nickname: None, + agent_role: Some("explorer".to_string()), + })), + ) + .await + .expect("child spawn should succeed"); + let grandchild_thread_id = control + .spawn_agent( + config, + text_input("hello grandchild"), + Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id: child_thread_id, + depth: 2, + agent_path: None, + agent_nickname: None, + agent_role: Some("worker".to_string()), + })), + ) + .await + .expect("grandchild spawn should succeed"); + + let mut subtree_thread_ids = manager + .list_agent_subtree_thread_ids(parent_thread_id) + .await + .expect("live subtree should load"); + subtree_thread_ids.sort_by_key(ToString::to_string); + let mut expected_subtree_thread_ids = + vec![parent_thread_id, child_thread_id, grandchild_thread_id]; + expected_subtree_thread_ids.sort_by_key(ToString::to_string); + + assert_eq!(subtree_thread_ids, expected_subtree_thread_ids); +} + #[tokio::test] async fn shutdown_agent_tree_closes_live_descendants() { let harness = AgentControlHarness::new().await; diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index e5b4fcf6a2..232f12ca71 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -928,6 +928,27 @@ impl ThreadManagerState { .collect() } + /// List parent-child edges for currently loaded thread-spawn agents. + pub(crate) async fn list_live_thread_spawn_edges(&self) -> Vec<(ThreadId, ThreadId)> { + self.threads + .read() + .await + .iter() + .filter_map(|(thread_id, thread)| { + if thread.session_source.is_internal() { + return None; + } + match &thread.session_source { + SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id, + .. + }) => Some((*parent_thread_id, *thread_id)), + _ => None, + } + }) + .collect() + } + /// Fetch a thread by ID or return ThreadNotFound. pub(crate) async fn get_thread(&self, thread_id: ThreadId) -> CodexResult> { let threads = self.threads.read().await;