jif-oai
2026-04-07 12:47:37 +01:00
committed by GitHub
parent 99f167e6bf
commit 89f1a44afa
4 changed files with 285 additions and 28 deletions
+9
View File
@@ -730,6 +730,15 @@ impl AgentControl {
self.state.agent_metadata_for_thread(agent_id)
}
pub(crate) async fn list_live_agent_subtree_thread_ids(
&self,
agent_id: ThreadId,
) -> CodexResult<Vec<ThreadId>> {
let mut thread_ids = vec![agent_id];
thread_ids.extend(self.live_thread_spawn_descendants(agent_id).await?);
Ok(thread_ids)
}
pub(crate) async fn get_agent_config_snapshot(
&self,
agent_id: ThreadId,
+126
View File
@@ -1687,6 +1687,132 @@ async fn resume_agent_from_rollout_reads_archived_rollout_path() {
.expect("resumed child shutdown should succeed");
}
#[tokio::test]
async fn list_agent_subtree_thread_ids_includes_anonymous_and_closed_descendants() {
let harness = AgentControlHarness::new().await;
let (parent_thread_id, _parent_thread) = harness.start_thread().await;
let worker_path = AgentPath::root().join("worker").expect("worker path");
let reviewer_path = AgentPath::root().join("reviewer").expect("reviewer path");
let worker_thread_id = harness
.control
.spawn_agent(
harness.config.clone(),
text_input("hello worker"),
Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn {
parent_thread_id,
depth: 1,
agent_path: Some(worker_path.clone()),
agent_nickname: None,
agent_role: Some("worker".to_string()),
})),
)
.await
.expect("worker spawn should succeed");
let worker_child_thread_id = harness
.control
.spawn_agent(
harness.config.clone(),
text_input("hello worker child"),
Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn {
parent_thread_id: worker_thread_id,
depth: 2,
agent_path: Some(
worker_path
.join("child")
.expect("worker child path should be valid"),
),
agent_nickname: None,
agent_role: Some("worker".to_string()),
})),
)
.await
.expect("worker child spawn should succeed");
let no_path_child_thread_id = harness
.control
.spawn_agent(
harness.config.clone(),
text_input("hello anonymous child"),
Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn {
parent_thread_id: worker_thread_id,
depth: 2,
agent_path: None,
agent_nickname: None,
agent_role: Some("worker".to_string()),
})),
)
.await
.expect("no-path child spawn should succeed");
let no_path_grandchild_thread_id = harness
.control
.spawn_agent(
harness.config.clone(),
text_input("hello anonymous grandchild"),
Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn {
parent_thread_id: no_path_child_thread_id,
depth: 3,
agent_path: None,
agent_nickname: None,
agent_role: Some("worker".to_string()),
})),
)
.await
.expect("no-path grandchild spawn should succeed");
let _reviewer_thread_id = harness
.control
.spawn_agent(
harness.config.clone(),
text_input("hello reviewer"),
Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn {
parent_thread_id,
depth: 1,
agent_path: Some(reviewer_path),
agent_nickname: None,
agent_role: Some("reviewer".to_string()),
})),
)
.await
.expect("reviewer spawn should succeed");
let _ = harness
.control
.shutdown_live_agent(no_path_grandchild_thread_id)
.await
.expect("no-path grandchild shutdown should succeed");
let mut worker_subtree_thread_ids = harness
.manager
.list_agent_subtree_thread_ids(worker_thread_id)
.await
.expect("worker subtree thread ids should load");
worker_subtree_thread_ids.sort_by_key(ToString::to_string);
let mut expected_worker_subtree_thread_ids = vec![
worker_thread_id,
worker_child_thread_id,
no_path_child_thread_id,
no_path_grandchild_thread_id,
];
expected_worker_subtree_thread_ids.sort_by_key(ToString::to_string);
assert_eq!(
worker_subtree_thread_ids,
expected_worker_subtree_thread_ids
);
let mut no_path_child_subtree_thread_ids = harness
.manager
.list_agent_subtree_thread_ids(no_path_child_thread_id)
.await
.expect("no-path subtree thread ids should load");
no_path_child_subtree_thread_ids.sort_by_key(ToString::to_string);
let mut expected_no_path_child_subtree_thread_ids =
vec![no_path_child_thread_id, no_path_grandchild_thread_id];
expected_no_path_child_subtree_thread_ids.sort_by_key(ToString::to_string);
assert_eq!(
no_path_child_subtree_thread_ids,
expected_no_path_child_subtree_thread_ids
);
}
#[tokio::test]
async fn shutdown_agent_tree_closes_live_descendants() {
let harness = AgentControlHarness::new().await;
+49
View File
@@ -43,9 +43,11 @@ use codex_protocol::protocol::SessionSource;
use codex_protocol::protocol::TurnAbortReason;
use codex_protocol::protocol::TurnAbortedEvent;
use codex_protocol::protocol::W3cTraceContext;
use codex_state::DirectionalThreadSpawnEdgeStatus;
use futures::StreamExt;
use futures::stream::FuturesUnordered;
use std::collections::HashMap;
use std::collections::HashSet;
use std::path::PathBuf;
use std::sync::Arc;
use std::sync::atomic::AtomicBool;
@@ -401,6 +403,53 @@ impl ThreadManager {
self.state.get_thread(thread_id).await
}
/// List `thread_id` plus all known descendants in its spawn subtree.
pub async fn list_agent_subtree_thread_ids(
&self,
thread_id: ThreadId,
) -> CodexResult<Vec<ThreadId>> {
let thread = self.state.get_thread(thread_id).await?;
let mut subtree_thread_ids = Vec::new();
let mut seen_thread_ids = HashSet::new();
subtree_thread_ids.push(thread_id);
seen_thread_ids.insert(thread_id);
if let Some(state_db_ctx) = thread.state_db() {
for status in [
DirectionalThreadSpawnEdgeStatus::Open,
DirectionalThreadSpawnEdgeStatus::Closed,
] {
for descendant_id in state_db_ctx
.list_thread_spawn_descendants_with_status(thread_id, status)
.await
.map_err(|err| {
CodexErr::Fatal(format!("failed to load thread-spawn descendants: {err}"))
})?
{
if seen_thread_ids.insert(descendant_id) {
subtree_thread_ids.push(descendant_id);
}
}
}
}
for descendant_id in thread
.codex
.session
.services
.agent_control
.list_live_agent_subtree_thread_ids(thread_id)
.await?
{
if seen_thread_ids.insert(descendant_id) {
subtree_thread_ids.push(descendant_id);
}
}
Ok(subtree_thread_ids)
}
pub async fn start_thread(&self, config: Config) -> CodexResult<NewThread> {
// Box delegated thread-spawn futures so these convenience wrappers do
// not inline the full spawn path into every caller's async state.