From 70cdb17703a4310b7173642e011f7534d2b2624f Mon Sep 17 00:00:00 2001 From: jif-oai Date: Thu, 19 Mar 2026 10:21:25 +0000 Subject: [PATCH] feat: add graph representation of agent network (#15056) Add a representation of the agent graph. This is now used for: * Cascade close agents (when I close a parent, it close the kids) * Cascade resume (oposite) Later, this will also be used for post-compaction stuffing of the context Direct fix for: https://github.com/openai/codex/issues/14458 --- MODULE.bazel.lock | 1 + codex-rs/Cargo.lock | 16 + codex-rs/core/src/agent/control.rs | 290 ++++++- codex-rs/core/src/agent/control_tests.rs | 794 +++++++++++++++++- codex-rs/core/src/memories/phase2.rs | 2 +- .../core/src/tools/handlers/agent_jobs.rs | 8 +- .../handlers/multi_agents/close_agent.rs | 2 +- .../src/tools/handlers/multi_agents_tests.rs | 201 ++++- codex-rs/core/src/tools/spec.rs | 2 +- codex-rs/state/Cargo.toml | 1 + .../migrations/0021_thread_spawn_edges.sql | 8 + codex-rs/state/src/lib.rs | 1 + codex-rs/state/src/model/graph.rs | 11 + codex-rs/state/src/model/mod.rs | 2 + codex-rs/state/src/runtime/threads.rs | 274 ++++++ 15 files changed, 1561 insertions(+), 52 deletions(-) create mode 100644 codex-rs/state/migrations/0021_thread_spawn_edges.sql create mode 100644 codex-rs/state/src/model/graph.rs diff --git a/MODULE.bazel.lock b/MODULE.bazel.lock index 47e3ca9cf..b37695679 100644 --- a/MODULE.bazel.lock +++ b/MODULE.bazel.lock @@ -1295,6 +1295,7 @@ "strum_0.26.3": "{\"dependencies\":[{\"features\":[\"macros\"],\"name\":\"phf\",\"optional\":true,\"req\":\"^0.10\"},{\"name\":\"strum_macros\",\"optional\":true,\"req\":\"^0.26.3\"},{\"kind\":\"dev\",\"name\":\"strum_macros\",\"req\":\"^0.26\"}],\"features\":{\"default\":[\"std\"],\"derive\":[\"strum_macros\"],\"std\":[]}}", "strum_0.27.2": "{\"dependencies\":[{\"features\":[\"macros\"],\"name\":\"phf\",\"optional\":true,\"req\":\"^0.12\"},{\"name\":\"strum_macros\",\"optional\":true,\"req\":\"^0.27\"}],\"features\":{\"default\":[\"std\"],\"derive\":[\"strum_macros\"],\"std\":[]}}", "strum_macros_0.26.4": "{\"dependencies\":[{\"name\":\"heck\",\"req\":\"^0.5.0\"},{\"name\":\"proc-macro2\",\"req\":\"^1.0\"},{\"name\":\"quote\",\"req\":\"^1.0\"},{\"name\":\"rustversion\",\"req\":\"^1.0\"},{\"kind\":\"dev\",\"name\":\"strum\",\"req\":\"^0.26\"},{\"features\":[\"parsing\",\"extra-traits\"],\"name\":\"syn\",\"req\":\"^2.0\"}],\"features\":{}}", + "strum_macros_0.27.2": "{\"dependencies\":[{\"name\":\"heck\",\"req\":\"^0.5.0\"},{\"name\":\"proc-macro2\",\"req\":\"^1.0\"},{\"name\":\"quote\",\"req\":\"^1.0\"},{\"features\":[\"parsing\"],\"name\":\"syn\",\"req\":\"^2.0\"}],\"features\":{}}", "strum_macros_0.28.0": "{\"dependencies\":[{\"name\":\"heck\",\"req\":\"^0.5.0\"},{\"name\":\"proc-macro2\",\"req\":\"^1.0\"},{\"name\":\"quote\",\"req\":\"^1.0\"},{\"features\":[\"parsing\"],\"name\":\"syn\",\"req\":\"^2.0\"}],\"features\":{}}", "subtle_2.6.1": "{\"dependencies\":[{\"kind\":\"dev\",\"name\":\"rand\",\"req\":\"^0.8\"}],\"features\":{\"const-generics\":[],\"core_hint_black_box\":[],\"default\":[\"std\",\"i128\"],\"i128\":[],\"nightly\":[],\"std\":[]}}", "supports-color_2.1.0": "{\"dependencies\":[{\"name\":\"is-terminal\",\"req\":\"^0.4.0\"},{\"name\":\"is_ci\",\"req\":\"^1.1.1\"}],\"features\":{}}", diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 0f03c455c..52c16b0fb 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2494,6 +2494,7 @@ dependencies = [ "serde", "serde_json", "sqlx", + "strum 0.27.2", "tokio", "tracing", "tracing-subscriber", @@ -9390,6 +9391,9 @@ name = "strum" version = "0.27.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "af23d6f6c1a224baef9d3f61e287d2761385a5b88fdab4eb4c6f11aeb54c4bcf" +dependencies = [ + "strum_macros 0.27.2", +] [[package]] name = "strum_macros" @@ -9404,6 +9408,18 @@ dependencies = [ "syn 2.0.114", ] +[[package]] +name = "strum_macros" +version = "0.27.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7695ce3845ea4b33927c055a39dc438a45b059f7c1b3d91d38d10355fb8cbca7" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn 2.0.114", +] + [[package]] name = "strum_macros" version = "0.28.0" diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index eaee6e985..8abed2616 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/codex-rs/core/src/agent/control.rs @@ -6,6 +6,8 @@ use crate::agent::status::is_final; use crate::codex_thread::ThreadConfigSnapshot; use crate::error::CodexErr; use crate::error::Result as CodexResult; +use crate::features::Feature; +use crate::find_archived_thread_path_by_id_str; use crate::find_thread_path_by_id_str; use crate::rollout::RolloutRecorder; use crate::session_prefix::format_subagent_context_line; @@ -23,9 +25,13 @@ use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::SubAgentSource; use codex_protocol::protocol::TokenUsage; use codex_protocol::user_input::UserInput; +use codex_state::DirectionalThreadSpawnEdgeStatus; +use std::collections::HashMap; +use std::collections::VecDeque; use std::sync::Arc; use std::sync::Weak; use tokio::sync::watch; +use tracing::warn; const AGENT_NAMES: &str = include_str!("agent_names.txt"); const FORKED_SPAWN_AGENT_OUTPUT_MESSAGE: &str = "You are the newly spawned agent. The prior conversation history was forked from your parent agent. Treat the next user message as your new task, and use the forked history only as background context."; @@ -169,7 +175,7 @@ impl AgentControl { "parent thread rollout unavailable for fork: {parent_thread_id}" )) })?; - let mut forked_rollout_items = + let mut forked_rollout_items: Vec = RolloutRecorder::get_rollout_history(&rollout_path) .await? .get_rollout_items(); @@ -218,6 +224,13 @@ impl AgentControl { // TODO(jif) add helper for drain state.notify_thread_created(new_thread.thread_id); + self.persist_thread_spawn_edge_for_source( + new_thread.thread.as_ref(), + new_thread.thread_id, + notification_source.as_ref(), + ) + .await; + self.send_input(new_thread.thread_id, items).await?; self.maybe_start_completion_watcher(new_thread.thread_id, notification_source); @@ -231,6 +244,84 @@ impl AgentControl { thread_id: ThreadId, session_source: SessionSource, ) -> CodexResult { + let root_depth = thread_spawn_depth(&session_source).unwrap_or(0); + let resumed_thread_id = self + .resume_single_agent_from_rollout(config.clone(), thread_id, session_source) + .await?; + let state = self.upgrade()?; + let Ok(resumed_thread) = state.get_thread(resumed_thread_id).await else { + return Ok(resumed_thread_id); + }; + let Some(state_db_ctx) = resumed_thread.state_db() else { + return Ok(resumed_thread_id); + }; + + let mut resume_queue = VecDeque::from([(thread_id, root_depth)]); + while let Some((parent_thread_id, parent_depth)) = resume_queue.pop_front() { + let child_ids = match state_db_ctx + .list_thread_spawn_children_with_status( + parent_thread_id, + DirectionalThreadSpawnEdgeStatus::Open, + ) + .await + { + Ok(child_ids) => child_ids, + Err(err) => { + warn!( + "failed to load persisted thread-spawn children for {parent_thread_id}: {err}" + ); + continue; + } + }; + + for child_thread_id in child_ids { + let child_depth = parent_depth + 1; + let child_resumed = if state.get_thread(child_thread_id).await.is_ok() { + true + } else { + let child_session_source = + SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id, + depth: child_depth, + agent_nickname: None, + agent_role: None, + }); + match self + .resume_single_agent_from_rollout( + config.clone(), + child_thread_id, + child_session_source, + ) + .await + { + Ok(_) => true, + Err(err) => { + warn!("failed to resume descendant thread {child_thread_id}: {err}"); + false + } + } + }; + if child_resumed { + resume_queue.push_back((child_thread_id, child_depth)); + } + } + } + + Ok(resumed_thread_id) + } + + async fn resume_single_agent_from_rollout( + &self, + mut config: crate::config::Config, + thread_id: ThreadId, + session_source: SessionSource, + ) -> CodexResult { + if let SessionSource::SubAgent(SubAgentSource::ThreadSpawn { depth, .. }) = &session_source + && *depth >= config.agent_max_depth + { + let _ = config.features.disable(Feature::SpawnCsv); + let _ = config.features.disable(Feature::Collab); + } let state = self.upgrade()?; let mut reservation = self.state.reserve_spawn_slot(config.agent_max_threads)?; let session_source = match session_source { @@ -280,9 +371,17 @@ impl AgentControl { .inherited_exec_policy_for_source(&state, Some(&session_source), &config) .await; let rollout_path = - find_thread_path_by_id_str(config.codex_home.as_path(), &thread_id.to_string()) + match find_thread_path_by_id_str(config.codex_home.as_path(), &thread_id.to_string()) .await? - .ok_or_else(|| CodexErr::ThreadNotFound(thread_id))?; + { + Some(rollout_path) => rollout_path, + None => find_archived_thread_path_by_id_str( + config.codex_home.as_path(), + &thread_id.to_string(), + ) + .await? + .ok_or_else(|| CodexErr::ThreadNotFound(thread_id))?, + }; let resumed_thread = state .resume_thread_from_rollout_with_source( @@ -298,7 +397,16 @@ impl AgentControl { // Resumed threads are re-registered in-memory and need the same listener // attachment path as freshly spawned threads. state.notify_thread_created(resumed_thread.thread_id); - self.maybe_start_completion_watcher(resumed_thread.thread_id, Some(notification_source)); + self.maybe_start_completion_watcher( + resumed_thread.thread_id, + Some(notification_source.clone()), + ); + self.persist_thread_spawn_edge_for_source( + resumed_thread.thread.as_ref(), + resumed_thread.thread_id, + Some(¬ification_source), + ) + .await; Ok(resumed_thread.thread_id) } @@ -332,15 +440,48 @@ impl AgentControl { state.send_op(agent_id, Op::Interrupt).await } - /// Submit a shutdown request to an existing agent thread. - pub(crate) async fn shutdown_agent(&self, agent_id: ThreadId) -> CodexResult { + /// Submit a shutdown request for a live agent without marking it explicitly closed in + /// persisted spawn-edge state. + pub(crate) async fn shutdown_live_agent(&self, agent_id: ThreadId) -> CodexResult { let state = self.upgrade()?; + if let Ok(thread) = state.get_thread(agent_id).await { + thread.codex.session.ensure_rollout_materialized().await; + thread.codex.session.flush_rollout().await; + } let result = state.send_op(agent_id, Op::Shutdown {}).await; let _ = state.remove_thread(&agent_id).await; self.state.release_spawned_thread(agent_id); result } + /// Mark `agent_id` as explicitly closed in persisted spawn-edge state, then shut down the + /// agent and any live descendants reached from the in-memory tree. + pub(crate) async fn close_agent(&self, agent_id: ThreadId) -> CodexResult { + let state = self.upgrade()?; + if let Ok(thread) = state.get_thread(agent_id).await + && let Some(state_db_ctx) = thread.state_db() + && let Err(err) = state_db_ctx + .set_thread_spawn_edge_status(agent_id, DirectionalThreadSpawnEdgeStatus::Closed) + .await + { + warn!("failed to persist thread-spawn edge status for {agent_id}: {err}"); + } + self.shutdown_agent_tree(agent_id).await + } + + /// Shut down `agent_id` and any live descendants reachable from the in-memory spawn tree. + async fn shutdown_agent_tree(&self, agent_id: ThreadId) -> CodexResult { + let descendant_ids = self.live_thread_spawn_descendants(agent_id).await?; + let result = self.shutdown_live_agent(agent_id).await; + for descendant_id in descendant_ids { + match self.shutdown_live_agent(descendant_id).await { + Ok(_) | Err(CodexErr::ThreadNotFound(_)) | Err(CodexErr::InternalAgentDied) => {} + Err(err) => return Err(err), + } + } + result + } + /// Fetch the last known status for `agent_id`, returning `NotFound` when unavailable. pub(crate) async fn get_status(&self, agent_id: ThreadId) -> AgentStatus { let Ok(state) = self.upgrade() else { @@ -407,34 +548,17 @@ impl AgentControl { &self, parent_thread_id: ThreadId, ) -> String { - let Ok(state) = self.upgrade() else { + let Ok(agents) = self.open_thread_spawn_children(parent_thread_id).await else { return String::new(); }; - let mut agents = Vec::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 SessionSource::SubAgent(SubAgentSource::ThreadSpawn { - parent_thread_id: agent_parent_thread_id, - agent_nickname, - .. - }) = snapshot.session_source - else { - continue; - }; - if agent_parent_thread_id != parent_thread_id { - continue; - } - agents.push(format_subagent_context_line( - &thread_id.to_string(), - agent_nickname.as_deref(), - )); - } - agents.sort(); - agents.join("\n") + agents + .into_iter() + .map(|(thread_id, nickname)| { + format_subagent_context_line(&thread_id.to_string(), nickname.as_deref()) + }) + .collect::>() + .join("\n") } /// Starts a detached watcher for sub-agents spawned from another thread. @@ -532,6 +656,110 @@ impl AgentControl { &parent_thread.codex.session.services.exec_policy, )) } + + async fn open_thread_spawn_children( + &self, + parent_thread_id: ThreadId, + ) -> CodexResult)>> { + let mut children_by_parent = self.live_thread_spawn_children().await?; + Ok(children_by_parent + .remove(&parent_thread_id) + .unwrap_or_default()) + } + + async fn live_thread_spawn_children( + &self, + ) -> CodexResult)>>> { + 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; + }; + children_by_parent + .entry(parent_thread_id) + .or_default() + .push((thread_id, snapshot.session_source.get_nickname())); + } + + for children in children_by_parent.values_mut() { + children.sort_by(|left, right| left.0.to_string().cmp(&right.0.to_string())); + } + + Ok(children_by_parent) + } + + async fn persist_thread_spawn_edge_for_source( + &self, + thread: &crate::CodexThread, + child_thread_id: ThreadId, + session_source: Option<&SessionSource>, + ) { + let Some(parent_thread_id) = session_source.and_then(thread_spawn_parent_thread_id) else { + return; + }; + let Some(state_db_ctx) = thread.state_db() else { + return; + }; + if let Err(err) = state_db_ctx + .upsert_thread_spawn_edge( + parent_thread_id, + child_thread_id, + DirectionalThreadSpawnEdgeStatus::Open, + ) + .await + { + warn!("failed to persist thread-spawn edge: {err}"); + } + } + + async fn live_thread_spawn_descendants( + &self, + root_thread_id: ThreadId, + ) -> CodexResult> { + let mut children_by_parent = self.live_thread_spawn_children().await?; + let mut descendants = Vec::new(); + let mut stack = children_by_parent + .remove(&root_thread_id) + .unwrap_or_default() + .into_iter() + .map(|(child_thread_id, _)| child_thread_id) + .rev() + .collect::>(); + + while let Some(thread_id) = stack.pop() { + descendants.push(thread_id); + if let Some(children) = children_by_parent.remove(&thread_id) { + for (child_thread_id, _) in children.into_iter().rev() { + stack.push(child_thread_id); + } + } + } + + Ok(descendants) + } +} + +fn thread_spawn_parent_thread_id(session_source: &SessionSource) -> Option { + match session_source { + SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id, .. + }) => Some(*parent_thread_id), + _ => None, + } +} + +fn thread_spawn_depth(session_source: &SessionSource) -> Option { + match session_source { + SessionSource::SubAgent(SubAgentSource::ThreadSpawn { depth, .. }) => Some(*depth), + _ => None, + } } #[cfg(test)] #[path = "control_tests.rs"] diff --git a/codex-rs/core/src/agent/control_tests.rs b/codex-rs/core/src/agent/control_tests.rs index 26819309a..7c2c46b2f 100644 --- a/codex-rs/core/src/agent/control_tests.rs +++ b/codex-rs/core/src/agent/control_tests.rs @@ -10,6 +10,7 @@ use crate::config_loader::LoaderOverrides; use crate::contextual_user_message::SUBAGENT_NOTIFICATION_OPEN_TAG; use crate::features::Feature; use assert_matches::assert_matches; +use chrono::Utc; use codex_protocol::config_types::ModeKind; use codex_protocol::models::ContentItem; use codex_protocol::models::ResponseItem; @@ -143,6 +144,42 @@ async fn wait_for_subagent_notification(parent_thread: &Arc) -> boo timeout(Duration::from_secs(2), wait).await.is_ok() } +async fn persist_thread_for_tree_resume(thread: &Arc, message: &str) { + thread + .inject_user_message_without_turn(message.to_string()) + .await; + thread.codex.session.ensure_rollout_materialized().await; + thread.codex.session.flush_rollout().await; +} + +async fn wait_for_live_thread_spawn_children( + control: &AgentControl, + parent_thread_id: ThreadId, + expected_children: &[ThreadId], +) { + let mut expected_children = expected_children.to_vec(); + expected_children.sort_by_key(std::string::ToString::to_string); + + timeout(Duration::from_secs(5), async { + loop { + let mut child_ids = control + .open_thread_spawn_children(parent_thread_id) + .await + .expect("live child list should load") + .into_iter() + .map(|(thread_id, _)| thread_id) + .collect::>(); + child_ids.sort_by_key(std::string::ToString::to_string); + if child_ids == expected_children { + break; + } + sleep(Duration::from_millis(25)).await; + } + }) + .await + .expect("expected persisted child tree"); +} + #[tokio::test] async fn send_input_errors_when_manager_dropped() { let control = AgentControl::default(); @@ -453,7 +490,7 @@ async fn spawn_agent_can_fork_parent_thread_history() { let _ = harness .control - .shutdown_agent(child_thread_id) + .shutdown_live_agent(child_thread_id) .await .expect("child shutdown should submit"); let _ = parent_thread @@ -529,7 +566,7 @@ async fn spawn_agent_fork_injects_output_for_parent_spawn_call() { let _ = harness .control - .shutdown_agent(child_thread_id) + .shutdown_live_agent(child_thread_id) .await .expect("child shutdown should submit"); let _ = parent_thread @@ -606,7 +643,7 @@ async fn spawn_agent_fork_flushes_parent_rollout_before_loading_history() { let _ = harness .control - .shutdown_agent(child_thread_id) + .shutdown_live_agent(child_thread_id) .await .expect("child shutdown should submit"); let _ = parent_thread @@ -653,7 +690,7 @@ async fn spawn_agent_respects_max_threads_limit() { assert_eq!(seen_max_threads, max_threads); let _ = control - .shutdown_agent(first_agent_id) + .shutdown_live_agent(first_agent_id) .await .expect("shutdown agent"); } @@ -678,7 +715,7 @@ async fn spawn_agent_releases_slot_after_shutdown() { .await .expect("spawn_agent should succeed"); let _ = control - .shutdown_agent(first_agent_id) + .shutdown_live_agent(first_agent_id) .await .expect("shutdown agent"); @@ -687,7 +724,7 @@ async fn spawn_agent_releases_slot_after_shutdown() { .await .expect("spawn_agent should succeed after shutdown"); let _ = control - .shutdown_agent(second_agent_id) + .shutdown_live_agent(second_agent_id) .await .expect("shutdown agent"); } @@ -723,7 +760,7 @@ async fn spawn_agent_limit_shared_across_clones() { assert_eq!(max_threads, 1); let _ = control - .shutdown_agent(first_agent_id) + .shutdown_live_agent(first_agent_id) .await .expect("shutdown agent"); } @@ -748,7 +785,7 @@ async fn resume_agent_respects_max_threads_limit() { .await .expect("spawn_agent should succeed"); let _ = control - .shutdown_agent(resumable_id) + .shutdown_live_agent(resumable_id) .await .expect("shutdown resumable thread"); @@ -770,7 +807,7 @@ async fn resume_agent_respects_max_threads_limit() { assert_eq!(seen_max_threads, max_threads); let _ = control - .shutdown_agent(active_id) + .shutdown_live_agent(active_id) .await .expect("shutdown active thread"); } @@ -800,7 +837,7 @@ async fn resume_agent_releases_slot_after_resume_failure() { .await .expect("spawn should succeed after failed resume"); let _ = control - .shutdown_agent(resumed_id) + .shutdown_live_agent(resumed_id) .await .expect("shutdown resumed thread"); } @@ -1046,7 +1083,7 @@ async fn resume_thread_subagent_restores_stored_nickname_and_role() { let _ = harness .control - .shutdown_agent(child_thread_id) + .shutdown_live_agent(child_thread_id) .await .expect("child shutdown should submit"); @@ -1089,7 +1126,740 @@ async fn resume_thread_subagent_restores_stored_nickname_and_role() { let _ = harness .control - .shutdown_agent(resumed_thread_id) + .shutdown_live_agent(resumed_thread_id) .await .expect("resumed child shutdown should submit"); } + +#[tokio::test] +async fn resume_agent_from_rollout_reads_archived_rollout_path() { + let harness = AgentControlHarness::new().await; + let child_thread_id = harness + .control + .spawn_agent(harness.config.clone(), text_input("hello"), None) + .await + .expect("child spawn should succeed"); + + let child_thread = harness + .manager + .get_thread(child_thread_id) + .await + .expect("child thread should exist"); + persist_thread_for_tree_resume(&child_thread, "persist before archiving").await; + let rollout_path = child_thread + .rollout_path() + .expect("thread should have rollout path"); + let state_db = child_thread + .state_db() + .expect("thread should have state db handle"); + + let _ = harness + .control + .shutdown_live_agent(child_thread_id) + .await + .expect("child shutdown should succeed"); + + let archived_root = harness + .config + .codex_home + .join(crate::ARCHIVED_SESSIONS_SUBDIR); + tokio::fs::create_dir_all(&archived_root) + .await + .expect("archived root should exist"); + let archived_rollout_path = archived_root.join( + rollout_path + .file_name() + .expect("rollout file name should be present"), + ); + tokio::fs::rename(&rollout_path, &archived_rollout_path) + .await + .expect("rollout should move to archived path"); + state_db + .mark_archived(child_thread_id, archived_rollout_path.as_path(), Utc::now()) + .await + .expect("state db archive update should succeed"); + + let resumed_thread_id = harness + .control + .resume_agent_from_rollout(harness.config.clone(), child_thread_id, SessionSource::Exec) + .await + .expect("resume should find archived rollout"); + assert_eq!(resumed_thread_id, child_thread_id); + + let _ = harness + .control + .shutdown_live_agent(child_thread_id) + .await + .expect("resumed child shutdown should succeed"); +} + +#[tokio::test] +async fn shutdown_agent_tree_closes_live_descendants() { + let harness = AgentControlHarness::new().await; + let (parent_thread_id, _parent_thread) = harness.start_thread().await; + + let child_thread_id = harness + .control + .spawn_agent( + harness.config.clone(), + text_input("hello child"), + Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id, + depth: 1, + agent_nickname: None, + agent_role: Some("explorer".to_string()), + })), + ) + .await + .expect("child spawn should succeed"); + let grandchild_thread_id = harness + .control + .spawn_agent( + harness.config.clone(), + text_input("hello grandchild"), + Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id: child_thread_id, + depth: 2, + agent_nickname: None, + agent_role: Some("worker".to_string()), + })), + ) + .await + .expect("grandchild spawn should succeed"); + + let child_thread = harness + .manager + .get_thread(child_thread_id) + .await + .expect("child thread should exist"); + let grandchild_thread = harness + .manager + .get_thread(grandchild_thread_id) + .await + .expect("grandchild thread should exist"); + persist_thread_for_tree_resume(&child_thread, "child persisted").await; + persist_thread_for_tree_resume(&grandchild_thread, "grandchild persisted").await; + wait_for_live_thread_spawn_children(&harness.control, parent_thread_id, &[child_thread_id]) + .await; + wait_for_live_thread_spawn_children(&harness.control, child_thread_id, &[grandchild_thread_id]) + .await; + + let _ = harness + .control + .shutdown_agent_tree(parent_thread_id) + .await + .expect("tree shutdown should succeed"); + + assert_eq!( + harness.control.get_status(parent_thread_id).await, + AgentStatus::NotFound + ); + assert_eq!( + harness.control.get_status(child_thread_id).await, + AgentStatus::NotFound + ); + assert_eq!( + harness.control.get_status(grandchild_thread_id).await, + AgentStatus::NotFound + ); + + let shutdown_ids = harness + .manager + .captured_ops() + .into_iter() + .filter_map(|(thread_id, op)| matches!(op, Op::Shutdown).then_some(thread_id)) + .collect::>(); + let mut expected_shutdown_ids = vec![parent_thread_id, child_thread_id, grandchild_thread_id]; + expected_shutdown_ids.sort_by_key(std::string::ToString::to_string); + let mut shutdown_ids = shutdown_ids; + shutdown_ids.sort_by_key(std::string::ToString::to_string); + assert_eq!(shutdown_ids, expected_shutdown_ids); +} + +#[tokio::test] +async fn shutdown_agent_tree_closes_descendants_when_started_at_child() { + let harness = AgentControlHarness::new().await; + let (parent_thread_id, _parent_thread) = harness.start_thread().await; + + let child_thread_id = harness + .control + .spawn_agent( + harness.config.clone(), + text_input("hello child"), + Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id, + depth: 1, + agent_nickname: None, + agent_role: Some("explorer".to_string()), + })), + ) + .await + .expect("child spawn should succeed"); + let grandchild_thread_id = harness + .control + .spawn_agent( + harness.config.clone(), + text_input("hello grandchild"), + Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id: child_thread_id, + depth: 2, + agent_nickname: None, + agent_role: Some("worker".to_string()), + })), + ) + .await + .expect("grandchild spawn should succeed"); + + let child_thread = harness + .manager + .get_thread(child_thread_id) + .await + .expect("child thread should exist"); + let grandchild_thread = harness + .manager + .get_thread(grandchild_thread_id) + .await + .expect("grandchild thread should exist"); + persist_thread_for_tree_resume(&child_thread, "child persisted").await; + persist_thread_for_tree_resume(&grandchild_thread, "grandchild persisted").await; + wait_for_live_thread_spawn_children(&harness.control, parent_thread_id, &[child_thread_id]) + .await; + wait_for_live_thread_spawn_children(&harness.control, child_thread_id, &[grandchild_thread_id]) + .await; + + let _ = harness + .control + .close_agent(child_thread_id) + .await + .expect("child close should succeed"); + + let _ = harness + .control + .shutdown_agent_tree(parent_thread_id) + .await + .expect("tree shutdown should succeed"); + + assert_eq!( + harness.control.get_status(child_thread_id).await, + AgentStatus::NotFound + ); + assert_eq!( + harness.control.get_status(grandchild_thread_id).await, + AgentStatus::NotFound + ); + assert_eq!( + harness.control.get_status(parent_thread_id).await, + AgentStatus::NotFound + ); + + let shutdown_ids = harness + .manager + .captured_ops() + .into_iter() + .filter_map(|(thread_id, op)| matches!(op, Op::Shutdown).then_some(thread_id)) + .collect::>(); + let mut expected_shutdown_ids = vec![parent_thread_id, child_thread_id, grandchild_thread_id]; + expected_shutdown_ids.sort_by_key(std::string::ToString::to_string); + let mut shutdown_ids = shutdown_ids; + shutdown_ids.sort_by_key(std::string::ToString::to_string); + assert_eq!(shutdown_ids, expected_shutdown_ids); +} + +#[tokio::test] +async fn resume_agent_from_rollout_does_not_reopen_closed_descendants() { + let harness = AgentControlHarness::new().await; + let (parent_thread_id, parent_thread) = harness.start_thread().await; + + let child_thread_id = harness + .control + .spawn_agent( + harness.config.clone(), + text_input("hello child"), + Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id, + depth: 1, + agent_nickname: None, + agent_role: Some("explorer".to_string()), + })), + ) + .await + .expect("child spawn should succeed"); + let grandchild_thread_id = harness + .control + .spawn_agent( + harness.config.clone(), + text_input("hello grandchild"), + Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id: child_thread_id, + depth: 2, + agent_nickname: None, + agent_role: Some("worker".to_string()), + })), + ) + .await + .expect("grandchild spawn should succeed"); + + let child_thread = harness + .manager + .get_thread(child_thread_id) + .await + .expect("child thread should exist"); + let grandchild_thread = harness + .manager + .get_thread(grandchild_thread_id) + .await + .expect("grandchild thread should exist"); + persist_thread_for_tree_resume(&parent_thread, "parent persisted").await; + persist_thread_for_tree_resume(&child_thread, "child persisted").await; + persist_thread_for_tree_resume(&grandchild_thread, "grandchild persisted").await; + wait_for_live_thread_spawn_children(&harness.control, parent_thread_id, &[child_thread_id]) + .await; + wait_for_live_thread_spawn_children(&harness.control, child_thread_id, &[grandchild_thread_id]) + .await; + + let _ = harness + .control + .close_agent(child_thread_id) + .await + .expect("child close should succeed"); + let _ = harness + .control + .shutdown_live_agent(parent_thread_id) + .await + .expect("parent shutdown should succeed"); + + let resumed_parent_thread_id = harness + .control + .resume_agent_from_rollout( + harness.config.clone(), + parent_thread_id, + SessionSource::Exec, + ) + .await + .expect("single-thread resume should succeed"); + assert_eq!(resumed_parent_thread_id, parent_thread_id); + assert_ne!( + harness.control.get_status(parent_thread_id).await, + AgentStatus::NotFound + ); + assert_eq!( + harness.control.get_status(child_thread_id).await, + AgentStatus::NotFound + ); + assert_eq!( + harness.control.get_status(grandchild_thread_id).await, + AgentStatus::NotFound + ); + + let _ = harness + .control + .shutdown_agent_tree(parent_thread_id) + .await + .expect("tree shutdown after resume should succeed"); +} + +#[tokio::test] +async fn resume_closed_child_reopens_open_descendants() { + let harness = AgentControlHarness::new().await; + let (parent_thread_id, parent_thread) = harness.start_thread().await; + + let child_thread_id = harness + .control + .spawn_agent( + harness.config.clone(), + text_input("hello child"), + Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id, + depth: 1, + agent_nickname: None, + agent_role: Some("explorer".to_string()), + })), + ) + .await + .expect("child spawn should succeed"); + let grandchild_thread_id = harness + .control + .spawn_agent( + harness.config.clone(), + text_input("hello grandchild"), + Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id: child_thread_id, + depth: 2, + agent_nickname: None, + agent_role: Some("worker".to_string()), + })), + ) + .await + .expect("grandchild spawn should succeed"); + + let child_thread = harness + .manager + .get_thread(child_thread_id) + .await + .expect("child thread should exist"); + let grandchild_thread = harness + .manager + .get_thread(grandchild_thread_id) + .await + .expect("grandchild thread should exist"); + persist_thread_for_tree_resume(&parent_thread, "parent persisted").await; + persist_thread_for_tree_resume(&child_thread, "child persisted").await; + persist_thread_for_tree_resume(&grandchild_thread, "grandchild persisted").await; + wait_for_live_thread_spawn_children(&harness.control, parent_thread_id, &[child_thread_id]) + .await; + wait_for_live_thread_spawn_children(&harness.control, child_thread_id, &[grandchild_thread_id]) + .await; + + let _ = harness + .control + .close_agent(child_thread_id) + .await + .expect("child close should succeed"); + + let resumed_child_thread_id = harness + .control + .resume_agent_from_rollout( + harness.config.clone(), + child_thread_id, + SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id, + depth: 1, + agent_nickname: None, + agent_role: None, + }), + ) + .await + .expect("child resume should succeed"); + assert_eq!(resumed_child_thread_id, child_thread_id); + assert_ne!( + harness.control.get_status(child_thread_id).await, + AgentStatus::NotFound + ); + assert_ne!( + harness.control.get_status(grandchild_thread_id).await, + AgentStatus::NotFound + ); + + let _ = harness + .control + .close_agent(child_thread_id) + .await + .expect("child close after resume should succeed"); + let _ = harness + .control + .shutdown_live_agent(parent_thread_id) + .await + .expect("parent shutdown should succeed"); +} + +#[tokio::test] +async fn resume_agent_from_rollout_reopens_open_descendants_after_manager_shutdown() { + let harness = AgentControlHarness::new().await; + let (parent_thread_id, parent_thread) = harness.start_thread().await; + + let child_thread_id = harness + .control + .spawn_agent( + harness.config.clone(), + text_input("hello child"), + Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id, + depth: 1, + agent_nickname: None, + agent_role: Some("explorer".to_string()), + })), + ) + .await + .expect("child spawn should succeed"); + let grandchild_thread_id = harness + .control + .spawn_agent( + harness.config.clone(), + text_input("hello grandchild"), + Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id: child_thread_id, + depth: 2, + agent_nickname: None, + agent_role: Some("worker".to_string()), + })), + ) + .await + .expect("grandchild spawn should succeed"); + + let child_thread = harness + .manager + .get_thread(child_thread_id) + .await + .expect("child thread should exist"); + let grandchild_thread = harness + .manager + .get_thread(grandchild_thread_id) + .await + .expect("grandchild thread should exist"); + persist_thread_for_tree_resume(&parent_thread, "parent persisted").await; + persist_thread_for_tree_resume(&child_thread, "child persisted").await; + persist_thread_for_tree_resume(&grandchild_thread, "grandchild persisted").await; + wait_for_live_thread_spawn_children(&harness.control, parent_thread_id, &[child_thread_id]) + .await; + wait_for_live_thread_spawn_children(&harness.control, child_thread_id, &[grandchild_thread_id]) + .await; + + let report = harness + .manager + .shutdown_all_threads_bounded(Duration::from_secs(5)) + .await; + assert_eq!(report.submit_failed, Vec::::new()); + assert_eq!(report.timed_out, Vec::::new()); + + let resumed_parent_thread_id = harness + .control + .resume_agent_from_rollout( + harness.config.clone(), + parent_thread_id, + SessionSource::Exec, + ) + .await + .expect("tree resume should succeed"); + assert_eq!(resumed_parent_thread_id, parent_thread_id); + assert_ne!( + harness.control.get_status(parent_thread_id).await, + AgentStatus::NotFound + ); + assert_ne!( + harness.control.get_status(child_thread_id).await, + AgentStatus::NotFound + ); + assert_ne!( + harness.control.get_status(grandchild_thread_id).await, + AgentStatus::NotFound + ); + + let _ = harness + .control + .shutdown_agent_tree(parent_thread_id) + .await + .expect("tree shutdown after subtree resume should succeed"); +} + +#[tokio::test] +async fn resume_agent_from_rollout_uses_edge_data_when_descendant_metadata_source_is_stale() { + let harness = AgentControlHarness::new().await; + let (parent_thread_id, parent_thread) = harness.start_thread().await; + + let child_thread_id = harness + .control + .spawn_agent( + harness.config.clone(), + text_input("hello child"), + Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id, + depth: 1, + agent_nickname: None, + agent_role: Some("explorer".to_string()), + })), + ) + .await + .expect("child spawn should succeed"); + let grandchild_thread_id = harness + .control + .spawn_agent( + harness.config.clone(), + text_input("hello grandchild"), + Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id: child_thread_id, + depth: 2, + agent_nickname: None, + agent_role: Some("worker".to_string()), + })), + ) + .await + .expect("grandchild spawn should succeed"); + + let child_thread = harness + .manager + .get_thread(child_thread_id) + .await + .expect("child thread should exist"); + let grandchild_thread = harness + .manager + .get_thread(grandchild_thread_id) + .await + .expect("grandchild thread should exist"); + persist_thread_for_tree_resume(&parent_thread, "parent persisted").await; + persist_thread_for_tree_resume(&child_thread, "child persisted").await; + persist_thread_for_tree_resume(&grandchild_thread, "grandchild persisted").await; + wait_for_live_thread_spawn_children(&harness.control, parent_thread_id, &[child_thread_id]) + .await; + wait_for_live_thread_spawn_children(&harness.control, child_thread_id, &[grandchild_thread_id]) + .await; + + let state_db = grandchild_thread + .state_db() + .expect("sqlite state db should be available"); + let mut stale_metadata = state_db + .get_thread(grandchild_thread_id) + .await + .expect("grandchild metadata query should succeed") + .expect("grandchild metadata should exist"); + stale_metadata.source = + serde_json::to_string(&SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id: ThreadId::new(), + depth: 99, + agent_nickname: None, + agent_role: Some("worker".to_string()), + })) + .expect("stale session source should serialize"); + state_db + .upsert_thread(&stale_metadata) + .await + .expect("stale grandchild metadata should persist"); + + let report = harness + .manager + .shutdown_all_threads_bounded(Duration::from_secs(5)) + .await; + assert_eq!(report.submit_failed, Vec::::new()); + assert_eq!(report.timed_out, Vec::::new()); + + let resumed_parent_thread_id = harness + .control + .resume_agent_from_rollout( + harness.config.clone(), + parent_thread_id, + SessionSource::Exec, + ) + .await + .expect("tree resume should succeed"); + assert_eq!(resumed_parent_thread_id, parent_thread_id); + assert_ne!( + harness.control.get_status(parent_thread_id).await, + AgentStatus::NotFound + ); + assert_ne!( + harness.control.get_status(child_thread_id).await, + AgentStatus::NotFound + ); + assert_ne!( + harness.control.get_status(grandchild_thread_id).await, + AgentStatus::NotFound + ); + + let resumed_grandchild_snapshot = harness + .manager + .get_thread(grandchild_thread_id) + .await + .expect("resumed grandchild thread should exist") + .config_snapshot() + .await; + let SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id: resumed_parent_thread_id, + depth: resumed_depth, + .. + }) = resumed_grandchild_snapshot.session_source + else { + panic!("expected thread-spawn sub-agent source"); + }; + assert_eq!(resumed_parent_thread_id, child_thread_id); + assert_eq!(resumed_depth, 2); + + let _ = harness + .control + .shutdown_agent_tree(parent_thread_id) + .await + .expect("tree shutdown after subtree resume should succeed"); +} + +#[tokio::test] +async fn resume_agent_from_rollout_skips_descendants_when_parent_resume_fails() { + let harness = AgentControlHarness::new().await; + let (parent_thread_id, parent_thread) = harness.start_thread().await; + + let child_thread_id = harness + .control + .spawn_agent( + harness.config.clone(), + text_input("hello child"), + Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id, + depth: 1, + agent_nickname: None, + agent_role: Some("explorer".to_string()), + })), + ) + .await + .expect("child spawn should succeed"); + let grandchild_thread_id = harness + .control + .spawn_agent( + harness.config.clone(), + text_input("hello grandchild"), + Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id: child_thread_id, + depth: 2, + agent_nickname: None, + agent_role: Some("worker".to_string()), + })), + ) + .await + .expect("grandchild spawn should succeed"); + + let child_thread = harness + .manager + .get_thread(child_thread_id) + .await + .expect("child thread should exist"); + let grandchild_thread = harness + .manager + .get_thread(grandchild_thread_id) + .await + .expect("grandchild thread should exist"); + persist_thread_for_tree_resume(&parent_thread, "parent persisted").await; + persist_thread_for_tree_resume(&child_thread, "child persisted").await; + persist_thread_for_tree_resume(&grandchild_thread, "grandchild persisted").await; + wait_for_live_thread_spawn_children(&harness.control, parent_thread_id, &[child_thread_id]) + .await; + wait_for_live_thread_spawn_children(&harness.control, child_thread_id, &[grandchild_thread_id]) + .await; + + let child_rollout_path = child_thread + .rollout_path() + .expect("child thread should have rollout path"); + let report = harness + .manager + .shutdown_all_threads_bounded(Duration::from_secs(5)) + .await; + assert_eq!(report.submit_failed, Vec::::new()); + assert_eq!(report.timed_out, Vec::::new()); + tokio::fs::remove_file(&child_rollout_path) + .await + .expect("child rollout path should be removable"); + + let resumed_parent_thread_id = harness + .control + .resume_agent_from_rollout( + harness.config.clone(), + parent_thread_id, + SessionSource::Exec, + ) + .await + .expect("root resume should succeed"); + assert_eq!(resumed_parent_thread_id, parent_thread_id); + assert_ne!( + harness.control.get_status(parent_thread_id).await, + AgentStatus::NotFound + ); + assert_eq!( + harness.control.get_status(child_thread_id).await, + AgentStatus::NotFound + ); + assert_eq!( + harness.control.get_status(grandchild_thread_id).await, + AgentStatus::NotFound + ); + + let _ = harness + .control + .shutdown_agent_tree(parent_thread_id) + .await + .expect("tree shutdown after partial subtree resume should succeed"); +} diff --git a/codex-rs/core/src/memories/phase2.rs b/codex-rs/core/src/memories/phase2.rs index b2a78fffb..23933ca7f 100644 --- a/codex-rs/core/src/memories/phase2.rs +++ b/codex-rs/core/src/memories/phase2.rs @@ -379,7 +379,7 @@ mod agent { // Fire and forget close of the agent. if !matches!(final_status, AgentStatus::Shutdown | AgentStatus::NotFound) { tokio::spawn(async move { - if let Err(err) = agent_control.shutdown_agent(thread_id).await { + if let Err(err) = agent_control.shutdown_live_agent(thread_id).await { warn!( "failed to auto-close global memory consolidation agent {thread_id}: {err}" ); diff --git a/codex-rs/core/src/tools/handlers/agent_jobs.rs b/codex-rs/core/src/tools/handlers/agent_jobs.rs index 42cb5242d..639b21d06 100644 --- a/codex-rs/core/src/tools/handlers/agent_jobs.rs +++ b/codex-rs/core/src/tools/handlers/agent_jobs.rs @@ -673,7 +673,7 @@ async fn run_agent_job_loop( let _ = session .services .agent_control - .shutdown_agent(thread_id) + .shutdown_live_agent(thread_id) .await; continue; } @@ -833,7 +833,7 @@ async fn recover_running_items( let _ = session .services .agent_control - .shutdown_agent(thread_id) + .shutdown_live_agent(thread_id) .await; } continue; @@ -955,7 +955,7 @@ async fn reap_stale_active_items( let _ = session .services .agent_control - .shutdown_agent(thread_id) + .shutdown_live_agent(thread_id) .await; active_items.remove(&thread_id); } @@ -991,7 +991,7 @@ async fn finalize_finished_item( let _ = session .services .agent_control - .shutdown_agent(thread_id) + .shutdown_live_agent(thread_id) .await; Ok(()) } diff --git a/codex-rs/core/src/tools/handlers/multi_agents/close_agent.rs b/codex-rs/core/src/tools/handlers/multi_agents/close_agent.rs index ead71c704..a6e37ed6d 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents/close_agent.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents/close_agent.rs @@ -72,7 +72,7 @@ impl ToolHandler for Handler { session .services .agent_control - .shutdown_agent(agent_id) + .close_agent(agent_id) .await .map_err(|err| collab_agent_error(agent_id, err)) .map(|_| ()) diff --git a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs index 99afe8ac2..be34a1570 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs @@ -6,6 +6,7 @@ use crate::built_in_model_providers; use crate::codex::make_session_and_context; use crate::config::DEFAULT_AGENT_MAX_DEPTH; use crate::config::types::ShellEnvironmentPolicy; +use crate::features::Feature; use crate::function_tool::FunctionCallError; use crate::protocol::AskForApproval; use crate::protocol::FileSystemSandboxPolicy; @@ -672,7 +673,7 @@ async fn resume_agent_restores_closed_agent_and_accepts_send_input() { let agent_id = thread.thread_id; let _ = manager .agent_control() - .shutdown_agent(agent_id) + .shutdown_live_agent(agent_id) .await .expect("shutdown agent"); assert_eq!( @@ -720,7 +721,7 @@ async fn resume_agent_restores_closed_agent_and_accepts_send_input() { let _ = manager .agent_control() - .shutdown_agent(agent_id) + .shutdown_live_agent(agent_id) .await .expect("shutdown resumed agent"); } @@ -1006,6 +1007,202 @@ async fn close_agent_submits_shutdown_and_returns_previous_status() { assert_eq!(status_after, AgentStatus::NotFound); } +#[tokio::test] +async fn tool_handlers_cascade_close_and_resume_and_keep_explicitly_closed_subtrees_closed() { + let (_session, turn) = make_session_and_context().await; + let manager = thread_manager(); + let mut config = turn.config.as_ref().clone(); + config.agent_max_depth = 3; + config + .features + .enable(Feature::Sqlite) + .expect("test config should allow sqlite"); + + let parent = manager + .start_thread(config.clone()) + .await + .expect("parent thread should start"); + let parent_thread_id = parent.thread_id; + let parent_session = parent.thread.codex.session.clone(); + + let child_spawn_output = SpawnAgentHandler + .handle(invocation( + parent_session.clone(), + parent_session.new_default_turn().await, + "spawn_agent", + function_payload(json!({"message": "hello child"})), + )) + .await + .expect("child spawn should succeed"); + let (child_content, child_success) = expect_text_output(child_spawn_output); + let child_result: serde_json::Value = + serde_json::from_str(&child_content).expect("child spawn result should be json"); + let child_thread_id = agent_id( + child_result + .get("agent_id") + .and_then(serde_json::Value::as_str) + .expect("child spawn result should include agent_id"), + ) + .expect("child agent_id should be valid"); + assert_eq!(child_success, Some(true)); + + let child_thread = manager + .get_thread(child_thread_id) + .await + .expect("child thread should exist"); + let child_session = child_thread.codex.session.clone(); + let grandchild_spawn_output = SpawnAgentHandler + .handle(invocation( + child_session.clone(), + child_session.new_default_turn().await, + "spawn_agent", + function_payload(json!({"message": "hello grandchild"})), + )) + .await + .expect("grandchild spawn should succeed"); + let (grandchild_content, grandchild_success) = expect_text_output(grandchild_spawn_output); + let grandchild_result: serde_json::Value = + serde_json::from_str(&grandchild_content).expect("grandchild spawn result should be json"); + let grandchild_thread_id = agent_id( + grandchild_result + .get("agent_id") + .and_then(serde_json::Value::as_str) + .expect("grandchild spawn result should include agent_id"), + ) + .expect("grandchild agent_id should be valid"); + assert_eq!(grandchild_success, Some(true)); + + let close_output = CloseAgentHandler + .handle(invocation( + parent_session.clone(), + parent_session.new_default_turn().await, + "close_agent", + function_payload(json!({"id": child_thread_id.to_string()})), + )) + .await + .expect("close_agent should close the child subtree"); + let (close_content, close_success) = expect_text_output(close_output); + let close_result: close_agent::CloseAgentResult = + serde_json::from_str(&close_content).expect("close_agent result should be json"); + assert_ne!(close_result.previous_status, AgentStatus::NotFound); + assert_eq!(close_success, Some(true)); + assert_eq!( + manager.agent_control().get_status(child_thread_id).await, + AgentStatus::NotFound + ); + assert_eq!( + manager + .agent_control() + .get_status(grandchild_thread_id) + .await, + AgentStatus::NotFound + ); + + let child_resume_output = ResumeAgentHandler + .handle(invocation( + parent_session.clone(), + parent_session.new_default_turn().await, + "resume_agent", + function_payload(json!({"id": child_thread_id.to_string()})), + )) + .await + .expect("resume_agent should reopen the child subtree"); + let (child_resume_content, child_resume_success) = expect_text_output(child_resume_output); + let child_resume_result: resume_agent::ResumeAgentResult = + serde_json::from_str(&child_resume_content).expect("resume result should be json"); + assert_ne!(child_resume_result.status, AgentStatus::NotFound); + assert_eq!(child_resume_success, Some(true)); + assert_ne!( + manager.agent_control().get_status(child_thread_id).await, + AgentStatus::NotFound + ); + assert_ne!( + manager + .agent_control() + .get_status(grandchild_thread_id) + .await, + AgentStatus::NotFound + ); + + let close_again_output = CloseAgentHandler + .handle(invocation( + parent_session.clone(), + parent_session.new_default_turn().await, + "close_agent", + function_payload(json!({"id": child_thread_id.to_string()})), + )) + .await + .expect("close_agent should be repeatable for the child subtree"); + let (close_again_content, close_again_success) = expect_text_output(close_again_output); + let close_again_result: close_agent::CloseAgentResult = + serde_json::from_str(&close_again_content) + .expect("second close_agent result should be json"); + assert_ne!(close_again_result.previous_status, AgentStatus::NotFound); + assert_eq!(close_again_success, Some(true)); + assert_eq!( + manager.agent_control().get_status(child_thread_id).await, + AgentStatus::NotFound + ); + assert_eq!( + manager + .agent_control() + .get_status(grandchild_thread_id) + .await, + AgentStatus::NotFound + ); + + let operator = manager + .start_thread(config) + .await + .expect("operator thread should start"); + let operator_session = operator.thread.codex.session.clone(); + let _ = manager + .agent_control() + .shutdown_live_agent(parent_thread_id) + .await + .expect("parent shutdown should succeed"); + assert_eq!( + manager.agent_control().get_status(parent_thread_id).await, + AgentStatus::NotFound + ); + + let parent_resume_output = ResumeAgentHandler + .handle(invocation( + operator_session, + operator.thread.codex.session.new_default_turn().await, + "resume_agent", + function_payload(json!({"id": parent_thread_id.to_string()})), + )) + .await + .expect("resume_agent should reopen the parent thread"); + let (parent_resume_content, parent_resume_success) = expect_text_output(parent_resume_output); + let parent_resume_result: resume_agent::ResumeAgentResult = + serde_json::from_str(&parent_resume_content).expect("parent resume result should be json"); + assert_ne!(parent_resume_result.status, AgentStatus::NotFound); + assert_eq!(parent_resume_success, Some(true)); + assert_ne!( + manager.agent_control().get_status(parent_thread_id).await, + AgentStatus::NotFound + ); + assert_eq!( + manager.agent_control().get_status(child_thread_id).await, + AgentStatus::NotFound + ); + assert_eq!( + manager + .agent_control() + .get_status(grandchild_thread_id) + .await, + AgentStatus::NotFound + ); + + let shutdown_report = manager + .shutdown_all_threads_bounded(Duration::from_secs(5)) + .await; + assert_eq!(shutdown_report.submit_failed, Vec::::new()); + assert_eq!(shutdown_report.timed_out, Vec::::new()); +} + #[tokio::test] async fn build_agent_spawn_config_uses_turn_context_values() { fn pick_allowed_sandbox_policy( diff --git a/codex-rs/core/src/tools/spec.rs b/codex-rs/core/src/tools/spec.rs index 032ec608f..8992eb445 100644 --- a/codex-rs/core/src/tools/spec.rs +++ b/codex-rs/core/src/tools/spec.rs @@ -1540,7 +1540,7 @@ fn create_close_agent_tool() -> ToolSpec { ToolSpec::Function(ResponsesApiTool { name: "close_agent".to_string(), - description: "Close an agent when it is no longer needed and return its previous status before shutdown was requested. Don't keep agents open for too long if they are not needed anymore.".to_string(), + description: "Close an agent and any open descendants when they are no longer needed, and return the target agent's previous status before shutdown was requested. Don't keep agents open for too long if they are not needed anymore.".to_string(), strict: false, defer_loading: None, parameters: JsonSchema::Object { diff --git a/codex-rs/state/Cargo.toml b/codex-rs/state/Cargo.toml index d4106da88..bb80f60e3 100644 --- a/codex-rs/state/Cargo.toml +++ b/codex-rs/state/Cargo.toml @@ -15,6 +15,7 @@ owo-colors = { workspace = true } serde = { workspace = true, features = ["derive"] } serde_json = { workspace = true } sqlx = { workspace = true } +strum = { workspace = true, features = ["derive"] } tokio = { workspace = true, features = ["fs", "io-util", "macros", "rt-multi-thread", "sync", "time"] } tracing = { workspace = true } tracing-subscriber = { workspace = true } diff --git a/codex-rs/state/migrations/0021_thread_spawn_edges.sql b/codex-rs/state/migrations/0021_thread_spawn_edges.sql new file mode 100644 index 000000000..d6514c46e --- /dev/null +++ b/codex-rs/state/migrations/0021_thread_spawn_edges.sql @@ -0,0 +1,8 @@ +CREATE TABLE thread_spawn_edges ( + parent_thread_id TEXT NOT NULL, + child_thread_id TEXT NOT NULL PRIMARY KEY, + status TEXT NOT NULL +); + +CREATE INDEX idx_thread_spawn_edges_parent_status + ON thread_spawn_edges(parent_thread_id, status); diff --git a/codex-rs/state/src/lib.rs b/codex-rs/state/src/lib.rs index e90672295..5929dad94 100644 --- a/codex-rs/state/src/lib.rs +++ b/codex-rs/state/src/lib.rs @@ -35,6 +35,7 @@ pub use model::Anchor; pub use model::BackfillState; pub use model::BackfillStats; pub use model::BackfillStatus; +pub use model::DirectionalThreadSpawnEdgeStatus; pub use model::ExtractionOutcome; pub use model::SortKey; pub use model::Stage1JobClaim; diff --git a/codex-rs/state/src/model/graph.rs b/codex-rs/state/src/model/graph.rs new file mode 100644 index 000000000..4ab9f8ff4 --- /dev/null +++ b/codex-rs/state/src/model/graph.rs @@ -0,0 +1,11 @@ +use strum::AsRefStr; +use strum::Display; +use strum::EnumString; + +/// Status attached to a directional thread-spawn edge. +#[derive(Debug, Clone, Copy, PartialEq, Eq, AsRefStr, Display, EnumString)] +#[strum(serialize_all = "snake_case")] +pub enum DirectionalThreadSpawnEdgeStatus { + Open, + Closed, +} diff --git a/codex-rs/state/src/model/mod.rs b/codex-rs/state/src/model/mod.rs index efaf3f787..39f0e800f 100644 --- a/codex-rs/state/src/model/mod.rs +++ b/codex-rs/state/src/model/mod.rs @@ -1,5 +1,6 @@ mod agent_job; mod backfill_state; +mod graph; mod log; mod memories; mod thread_metadata; @@ -13,6 +14,7 @@ pub use agent_job::AgentJobProgress; pub use agent_job::AgentJobStatus; pub use backfill_state::BackfillState; pub use backfill_state::BackfillStatus; +pub use graph::DirectionalThreadSpawnEdgeStatus; pub use log::LogEntry; pub use log::LogQuery; pub use log::LogRow; diff --git a/codex-rs/state/src/runtime/threads.rs b/codex-rs/state/src/runtime/threads.rs index 6373568e2..1f62deb62 100644 --- a/codex-rs/state/src/runtime/threads.rs +++ b/codex-rs/state/src/runtime/threads.rs @@ -1,4 +1,5 @@ use super::*; +use codex_protocol::protocol::SessionSource; impl StateRuntime { pub async fn get_thread(&self, id: ThreadId) -> anyhow::Result> { @@ -78,6 +79,172 @@ ORDER BY position ASC Ok(Some(tools)) } + /// Persist or replace the directional parent-child edge for a spawned thread. + pub async fn upsert_thread_spawn_edge( + &self, + parent_thread_id: ThreadId, + child_thread_id: ThreadId, + status: crate::DirectionalThreadSpawnEdgeStatus, + ) -> anyhow::Result<()> { + sqlx::query( + r#" +INSERT INTO thread_spawn_edges ( + parent_thread_id, + child_thread_id, + status +) VALUES (?, ?, ?) +ON CONFLICT(child_thread_id) DO UPDATE SET + parent_thread_id = excluded.parent_thread_id, + status = excluded.status + "#, + ) + .bind(parent_thread_id.to_string()) + .bind(child_thread_id.to_string()) + .bind(status.as_ref()) + .execute(self.pool.as_ref()) + .await?; + Ok(()) + } + + /// Update the persisted lifecycle status of a spawned thread's incoming edge. + pub async fn set_thread_spawn_edge_status( + &self, + child_thread_id: ThreadId, + status: crate::DirectionalThreadSpawnEdgeStatus, + ) -> anyhow::Result<()> { + sqlx::query("UPDATE thread_spawn_edges SET status = ? WHERE child_thread_id = ?") + .bind(status.as_ref()) + .bind(child_thread_id.to_string()) + .execute(self.pool.as_ref()) + .await?; + Ok(()) + } + + /// List direct spawned children of `parent_thread_id` whose edge matches `status`. + pub async fn list_thread_spawn_children_with_status( + &self, + parent_thread_id: ThreadId, + status: crate::DirectionalThreadSpawnEdgeStatus, + ) -> anyhow::Result> { + self.list_thread_spawn_children_matching(parent_thread_id, Some(status)) + .await + } + + /// List spawned descendants of `root_thread_id` whose edges match `status`. + /// + /// Descendants are returned breadth-first by depth, then by thread id for stable ordering. + pub async fn list_thread_spawn_descendants_with_status( + &self, + root_thread_id: ThreadId, + status: crate::DirectionalThreadSpawnEdgeStatus, + ) -> anyhow::Result> { + self.list_thread_spawn_descendants_matching(root_thread_id, Some(status)) + .await + } + + async fn list_thread_spawn_children_matching( + &self, + parent_thread_id: ThreadId, + status: Option, + ) -> anyhow::Result> { + let mut query = String::from( + "SELECT child_thread_id FROM thread_spawn_edges WHERE parent_thread_id = ?", + ); + if status.is_some() { + query.push_str(" AND status = ?"); + } + query.push_str(" ORDER BY child_thread_id"); + + let mut sql = sqlx::query(query.as_str()).bind(parent_thread_id.to_string()); + if let Some(status) = status { + sql = sql.bind(status.to_string()); + } + + let rows = sql.fetch_all(self.pool.as_ref()).await?; + rows.into_iter() + .map(|row| { + ThreadId::try_from(row.try_get::("child_thread_id")?).map_err(Into::into) + }) + .collect() + } + + async fn list_thread_spawn_descendants_matching( + &self, + root_thread_id: ThreadId, + status: Option, + ) -> anyhow::Result> { + let status_filter = if status.is_some() { + " AND status = ?" + } else { + "" + }; + let query = format!( + r#" +WITH RECURSIVE subtree(child_thread_id, depth) AS ( + SELECT child_thread_id, 1 + FROM thread_spawn_edges + WHERE parent_thread_id = ?{status_filter} + UNION ALL + SELECT edge.child_thread_id, subtree.depth + 1 + FROM thread_spawn_edges AS edge + JOIN subtree ON edge.parent_thread_id = subtree.child_thread_id + WHERE 1 = 1{status_filter} +) +SELECT child_thread_id +FROM subtree +ORDER BY depth ASC, child_thread_id ASC + "# + ); + + let mut sql = sqlx::query(query.as_str()).bind(root_thread_id.to_string()); + if let Some(status) = status { + let status = status.to_string(); + sql = sql.bind(status.clone()).bind(status); + } + + let rows = sql.fetch_all(self.pool.as_ref()).await?; + rows.into_iter() + .map(|row| { + ThreadId::try_from(row.try_get::("child_thread_id")?).map_err(Into::into) + }) + .collect() + } + + async fn insert_thread_spawn_edge_if_absent( + &self, + parent_thread_id: ThreadId, + child_thread_id: ThreadId, + ) -> anyhow::Result<()> { + sqlx::query( + r#" +INSERT INTO thread_spawn_edges ( + parent_thread_id, + child_thread_id, + status +) VALUES (?, ?, ?) +ON CONFLICT(child_thread_id) DO NOTHING + "#, + ) + .bind(parent_thread_id.to_string()) + .bind(child_thread_id.to_string()) + .bind(crate::DirectionalThreadSpawnEdgeStatus::Open.as_ref()) + .execute(self.pool.as_ref()) + .await?; + Ok(()) + } + + async fn insert_thread_spawn_edge_from_source_if_absent( + &self, + child_thread_id: ThreadId, + source: &str, + ) -> anyhow::Result<()> { + let Some(parent_thread_id) = thread_spawn_parent_thread_id_from_source_str(source) else { + return Ok(()); + }; + self.insert_thread_spawn_edge_if_absent(parent_thread_id, child_thread_id) + .await + } + /// Find a rollout path by thread id using the underlying database. pub async fn find_rollout_path_by_id( &self, @@ -276,6 +443,8 @@ ON CONFLICT(id) DO NOTHING .bind("enabled") .execute(self.pool.as_ref()) .await?; + self.insert_thread_spawn_edge_from_source_if_absent(metadata.id, metadata.source.as_str()) + .await?; Ok(result.rows_affected() > 0) } @@ -420,6 +589,8 @@ ON CONFLICT(id) DO UPDATE SET .bind(creation_memory_mode.unwrap_or("enabled")) .execute(self.pool.as_ref()) .await?; + self.insert_thread_spawn_edge_from_source_if_absent(metadata.id, metadata.source.as_str()) + .await?; Ok(()) } @@ -602,6 +773,18 @@ pub(super) fn extract_memory_mode(items: &[RolloutItem]) -> Option { }) } +fn thread_spawn_parent_thread_id_from_source_str(source: &str) -> Option { + let parsed_source = serde_json::from_str(source) + .or_else(|_| serde_json::from_value::(Value::String(source.to_string()))); + match parsed_source.ok() { + Some(SessionSource::SubAgent(codex_protocol::protocol::SubAgentSource::ThreadSpawn { + parent_thread_id, + .. + })) => Some(parent_thread_id), + _ => None, + } +} + pub(super) fn push_thread_filters<'a>( builder: &mut QueryBuilder<'a, Sqlite>, archived_only: bool, @@ -680,6 +863,7 @@ pub(super) fn push_thread_order_and_limit( #[cfg(test)] mod tests { use super::*; + use crate::DirectionalThreadSpawnEdgeStatus; use crate::runtime::test_support::test_thread_metadata; use crate::runtime::test_support::unique_temp_dir; use codex_protocol::protocol::EventMsg; @@ -1072,4 +1256,94 @@ mod tests { assert_eq!(persisted.tokens_used, 321); assert_eq!(persisted.updated_at, override_updated_at); } + + #[tokio::test] + async fn thread_spawn_edges_track_directional_status() { + let codex_home = unique_temp_dir(); + let runtime = StateRuntime::init(codex_home, "test-provider".to_string()) + .await + .expect("state db should initialize"); + let parent_thread_id = + ThreadId::from_string("00000000-0000-0000-0000-000000000900").expect("valid thread id"); + let child_thread_id = + ThreadId::from_string("00000000-0000-0000-0000-000000000901").expect("valid thread id"); + let grandchild_thread_id = + ThreadId::from_string("00000000-0000-0000-0000-000000000902").expect("valid thread id"); + + runtime + .upsert_thread_spawn_edge( + parent_thread_id, + child_thread_id, + DirectionalThreadSpawnEdgeStatus::Open, + ) + .await + .expect("child edge insert should succeed"); + runtime + .upsert_thread_spawn_edge( + child_thread_id, + grandchild_thread_id, + DirectionalThreadSpawnEdgeStatus::Open, + ) + .await + .expect("grandchild edge insert should succeed"); + + let children = runtime + .list_thread_spawn_children_with_status( + parent_thread_id, + DirectionalThreadSpawnEdgeStatus::Open, + ) + .await + .expect("open child list should load"); + assert_eq!(children, vec![child_thread_id]); + + let descendants = runtime + .list_thread_spawn_descendants_with_status( + parent_thread_id, + DirectionalThreadSpawnEdgeStatus::Open, + ) + .await + .expect("open descendants should load"); + assert_eq!(descendants, vec![child_thread_id, grandchild_thread_id]); + + runtime + .set_thread_spawn_edge_status(child_thread_id, DirectionalThreadSpawnEdgeStatus::Closed) + .await + .expect("edge close should succeed"); + + let open_children = runtime + .list_thread_spawn_children_with_status( + parent_thread_id, + DirectionalThreadSpawnEdgeStatus::Open, + ) + .await + .expect("open child list should load"); + assert_eq!(open_children, Vec::::new()); + + let closed_children = runtime + .list_thread_spawn_children_with_status( + parent_thread_id, + DirectionalThreadSpawnEdgeStatus::Closed, + ) + .await + .expect("closed child list should load"); + assert_eq!(closed_children, vec![child_thread_id]); + + let closed_descendants = runtime + .list_thread_spawn_descendants_with_status( + parent_thread_id, + DirectionalThreadSpawnEdgeStatus::Closed, + ) + .await + .expect("closed descendants should load"); + assert_eq!(closed_descendants, vec![child_thread_id]); + + let open_descendants_from_child = runtime + .list_thread_spawn_descendants_with_status( + child_thread_id, + DirectionalThreadSpawnEdgeStatus::Open, + ) + .await + .expect("open descendants from child should load"); + assert_eq!(open_descendants_from_child, vec![grandchild_thread_id]); + } }