diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index 6317dff3b59e..a6e945b5f6ac 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/codex-rs/core/src/agent/control.rs @@ -1,12 +1,15 @@ use crate::agent::AgentStatus; use crate::agent::registry::AgentMetadata; use crate::agent::registry::AgentRegistry; +use crate::agent::residency::MIN_RESIDENT_SUBAGENTS; +use crate::agent::residency::ResidentAgents; use crate::agent::role::DEFAULT_ROLE_NAME; use crate::agent::role::resolve_role_config; use crate::agent::status::is_final; use crate::codex_thread::ThreadConfigSnapshot; use crate::config::Config; use crate::session::emit_subagent_session_started; +use crate::session::session::Session; use crate::session_prefix::format_subagent_context_line; use crate::session_prefix::format_subagent_notification_message; use crate::shell_snapshot::ShellSnapshot; @@ -34,16 +37,18 @@ use codex_protocol::protocol::TurnEnvironmentSelection; use codex_protocol::user_input::UserInput; use codex_state::DirectionalThreadSpawnEdgeStatus; use codex_thread_store::ReadThreadParams; +use futures::future::BoxFuture; use serde::Serialize; use std::collections::HashMap; -use std::collections::VecDeque; use std::sync::Arc; use std::sync::Weak; +use tokio::sync::Mutex; use tokio::sync::watch; use tracing::warn; -const AGENT_NAMES: &str = include_str!("agent_names.txt"); -const ROOT_LAST_TASK_MESSAGE: &str = "Main thread"; +mod legacy; +mod resident; +mod spawn; #[derive(Clone, Debug, PartialEq, Eq)] pub(crate) enum SpawnAgentForkMode { @@ -59,11 +64,6 @@ pub(crate) struct SpawnAgentOptions { pub(crate) environments: Option>, } -struct SpawnAgentThreadInheritance { - shell_snapshot: Option>, - exec_policy: Option>, -} - #[derive(Clone, Debug)] pub(crate) struct LiveAgent { pub(crate) thread_id: ThreadId, @@ -78,77 +78,30 @@ pub(crate) struct ListedAgent { pub(crate) last_task_message: Option, } -fn default_agent_nickname_list() -> Vec<&'static str> { - AGENT_NAMES - .lines() - .map(str::trim) - .filter(|name| !name.is_empty()) - .collect() -} - -fn agent_nickname_candidates(config: &Config, role_name: Option<&str>) -> Vec { - let role_name = role_name.unwrap_or(DEFAULT_ROLE_NAME); - if let Some(candidates) = - resolve_role_config(config, role_name).and_then(|role| role.nickname_candidates.clone()) - { - return candidates; - } - - default_agent_nickname_list() - .into_iter() - .map(ToOwned::to_owned) - .collect() -} - -fn keep_forked_rollout_item(item: &RolloutItem, preserve_reference_context_item: bool) -> bool { - match item { - RolloutItem::ResponseItem(ResponseItem::Message { role, phase, .. }) => match role.as_str() - { - "system" | "developer" | "user" => true, - "assistant" => *phase == Some(MessagePhase::FinalAnswer), - _ => false, +fn agent_metadata_from_session_source( + thread_id: ThreadId, + session_source: &SessionSource, +) -> AgentMetadata { + match session_source { + SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + agent_path, + agent_nickname, + agent_role, + .. + }) => AgentMetadata { + agent_id: Some(thread_id), + agent_path: agent_path.clone(), + agent_nickname: agent_nickname.clone(), + agent_role: agent_role.clone(), + last_task_message: None, + }, + _ => AgentMetadata { + agent_id: Some(thread_id), + ..Default::default() }, - RolloutItem::ResponseItem( - ResponseItem::AgentMessage { .. } - | ResponseItem::Reasoning { .. } - | ResponseItem::LocalShellCall { .. } - | ResponseItem::FunctionCall { .. } - | ResponseItem::ToolSearchCall { .. } - | ResponseItem::FunctionCallOutput { .. } - | ResponseItem::CustomToolCall { .. } - | ResponseItem::CustomToolCallOutput { .. } - | ResponseItem::ToolSearchOutput { .. } - | ResponseItem::WebSearchCall { .. } - | ResponseItem::ImageGenerationCall { .. } - | ResponseItem::Compaction { .. } - | ResponseItem::CompactionTrigger - | ResponseItem::ContextCompaction { .. } - | ResponseItem::Other, - ) => false, - // Full-history forks preserve the cached prompt prefix and can keep diffing - // from the parent's durable baseline. Truncated forks drop part of that prompt, - // so they must rebuild context on their first child turn. - RolloutItem::TurnContext(_) => preserve_reference_context_item, - RolloutItem::Compacted(_) | RolloutItem::EventMsg(_) | RolloutItem::SessionMeta(_) => true, } } -fn is_multi_agent_v2_usage_hint_message(item: &ResponseItem, usage_hint_texts: &[String]) -> bool { - let ResponseItem::Message { role, content, .. } = item else { - return false; - }; - if role != "developer" { - return false; - } - let [ContentItem::InputText { text }] = content.as_slice() else { - return false; - }; - - usage_hint_texts - .iter() - .any(|usage_hint_text| usage_hint_text == text) -} - /// Control-plane handle for multi-agent operations. /// `AgentControl` is held by each session (via `SessionServices`). It provides capability to /// spawn new agents and the inter-agent communication layer. @@ -165,6 +118,7 @@ pub(crate) struct AgentControl { /// `ThreadManagerState -> CodexThread -> Session -> SessionServices -> ThreadManagerState`. manager: Weak, state: Arc, + residency: Arc>, } impl AgentControl { @@ -185,531 +139,6 @@ impl AgentControl { self.session_id } - /// Spawn a new agent thread and submit the initial prompt. - #[cfg(test)] - pub(crate) async fn spawn_agent( - &self, - config: Config, - initial_operation: Op, - session_source: Option, - ) -> CodexResult { - let spawned_agent = Box::pin(self.spawn_agent_internal( - config, - initial_operation, - session_source, - SpawnAgentOptions::default(), - )) - .await?; - Ok(spawned_agent.thread_id) - } - - /// Spawn an agent thread with some metadata. - pub(crate) async fn spawn_agent_with_metadata( - &self, - config: Config, - initial_operation: Op, - session_source: Option, - options: SpawnAgentOptions, // TODO(jif) drop with new fork. - ) -> CodexResult { - Box::pin(self.spawn_agent_internal(config, initial_operation, session_source, options)) - .await - } - - async fn spawn_agent_internal( - &self, - config: Config, - initial_operation: Op, - session_source: Option, - options: SpawnAgentOptions, - ) -> CodexResult { - let state = self.upgrade()?; - let multi_agent_version = state - .effective_multi_agent_version_for_spawn( - &InitialHistory::New, - session_source.as_ref(), - options.parent_thread_id, - /*forked_from_thread_id*/ None, - &config, - ) - .await; - let agent_max_threads = config.effective_agent_max_threads(multi_agent_version); - let mut reservation = self.state.reserve_spawn_slot(agent_max_threads)?; - let inheritance = SpawnAgentThreadInheritance { - shell_snapshot: self - .inherited_shell_snapshot_for_source(&state, session_source.as_ref()) - .await, - exec_policy: self - .inherited_exec_policy_for_source(&state, session_source.as_ref(), &config) - .await, - }; - let (session_source, mut agent_metadata) = match session_source { - Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { - parent_thread_id, - depth, - agent_path, - agent_role, - .. - })) => { - let (session_source, agent_metadata) = self.prepare_thread_spawn( - &mut reservation, - &config, - parent_thread_id, - depth, - agent_path, - agent_role, - /*preferred_agent_nickname*/ None, - )?; - (Some(session_source), agent_metadata) - } - other => (other, AgentMetadata::default()), - }; - let notification_source = session_source.clone(); - - // The same `AgentControl` is sent to spawn the thread. - let new_thread = match (session_source, options.fork_mode.as_ref(), inheritance) { - (Some(session_source), Some(_), inheritance) => { - Box::pin(self.spawn_forked_thread( - &state, - config, - session_source, - &options, - inheritance, - multi_agent_version, - )) - .await? - } - (Some(session_source), None, inheritance) => { - Box::pin(state.spawn_new_thread_with_source( - config.clone(), - self.clone(), - session_source, - options.parent_thread_id, - /*forked_from_thread_id*/ None, - /*thread_source*/ Some(ThreadSource::Subagent), - /*metrics_service_name*/ None, - inheritance.shell_snapshot, - inheritance.exec_policy, - options.environments.clone(), - )) - .await? - } - (None, _, _) => Box::pin(state.spawn_new_thread(config.clone(), self.clone())).await?, - }; - agent_metadata.agent_id = Some(new_thread.thread_id); - reservation.commit(agent_metadata.clone()); - - if let Some(SessionSource::SubAgent( - subagent_source @ SubAgentSource::ThreadSpawn { - parent_thread_id, .. - }, - )) = notification_source.as_ref() - { - let client_metadata = match state.get_thread(*parent_thread_id).await { - Ok(parent_thread) => { - parent_thread - .codex - .session - .app_server_client_metadata() - .await - } - Err(error) => { - tracing::warn!( - error = %error, - parent_thread_id = %parent_thread_id, - "skipping subagent thread analytics: failed to load parent thread metadata" - ); - crate::session::session::AppServerClientMetadata { - client_name: None, - client_version: None, - } - } - }; - let thread_config = new_thread.thread.codex.thread_config_snapshot().await; - let parent_thread_id = thread_config.parent_thread_id; - emit_subagent_session_started( - &new_thread - .thread - .codex - .session - .services - .analytics_events_client, - client_metadata, - new_thread.thread.codex.session.session_id(), - new_thread.thread_id, - parent_thread_id, - thread_config, - subagent_source.clone(), - ); - } - - // Notify a new thread has been created. This notification will be processed by clients - // to subscribe or drain this newly created thread. - // 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, initial_operation) - .await?; - if multi_agent_version != MultiAgentVersion::V2 { - let child_reference = agent_metadata - .agent_path - .as_ref() - .map(ToString::to_string) - .unwrap_or_else(|| new_thread.thread_id.to_string()); - self.maybe_start_completion_watcher( - new_thread.thread_id, - notification_source, - child_reference, - agent_metadata.agent_path.clone(), - ); - } - - Ok(LiveAgent { - thread_id: new_thread.thread_id, - metadata: agent_metadata, - status: self.get_status(new_thread.thread_id).await, - }) - } - - async fn spawn_forked_thread( - &self, - state: &Arc, - config: Config, - session_source: SessionSource, - options: &SpawnAgentOptions, - inheritance: SpawnAgentThreadInheritance, - multi_agent_version: MultiAgentVersion, - ) -> CodexResult { - let SpawnAgentThreadInheritance { - shell_snapshot: inherited_shell_snapshot, - exec_policy: inherited_exec_policy, - } = inheritance; - if options.fork_parent_spawn_call_id.is_none() { - return Err(CodexErr::Fatal( - "spawn_agent fork requires a parent spawn call id".to_string(), - )); - } - let Some(fork_mode) = options.fork_mode.as_ref() else { - return Err(CodexErr::Fatal( - "spawn_agent fork requires a fork mode".to_string(), - )); - }; - let SessionSource::SubAgent(SubAgentSource::ThreadSpawn { - parent_thread_id, .. - }) = &session_source - else { - return Err(CodexErr::Fatal( - "spawn_agent fork requires a thread-spawn session source".to_string(), - )); - }; - - let parent_thread_id = *parent_thread_id; - let parent_thread = state.get_thread(parent_thread_id).await.ok(); - if let Some(parent_thread) = parent_thread.as_ref() { - // `record_conversation_items` only queues persistence writes asynchronously. - // Flush before snapshotting store history for a fork. - parent_thread.ensure_rollout_materialized().await; - parent_thread.flush_rollout().await?; - } - - let parent_history = state - .read_stored_thread(ReadThreadParams { - thread_id: parent_thread_id, - include_archived: true, - include_history: true, - }) - .await? - .history - .ok_or_else(|| { - CodexErr::Fatal(format!( - "parent thread history unavailable for fork: {parent_thread_id}" - )) - })?; - - let mut forked_rollout_items = parent_history.items; - if let SpawnAgentForkMode::LastNTurns(last_n_turns) = fork_mode { - forked_rollout_items = - truncate_rollout_to_last_n_fork_turns(&forked_rollout_items, *last_n_turns); - } - let multi_agent_v2_usage_hint_texts_to_filter: Vec = - if let Some(parent_thread) = parent_thread.as_ref() { - if multi_agent_version == MultiAgentVersion::V2 { - let parent_config = parent_thread.codex.session.get_config().await; - [ - parent_config - .multi_agent_v2 - .root_agent_usage_hint_text - .clone(), - parent_config - .multi_agent_v2 - .subagent_usage_hint_text - .clone(), - ] - .into_iter() - .flatten() - .collect() - } else { - Vec::new() - } - } else if multi_agent_version == MultiAgentVersion::V2 { - [ - config.multi_agent_v2.root_agent_usage_hint_text.clone(), - config.multi_agent_v2.subagent_usage_hint_text.clone(), - ] - .into_iter() - .flatten() - .collect() - } else { - Vec::new() - }; - let preserve_reference_context_item = matches!(fork_mode, SpawnAgentForkMode::FullHistory); - forked_rollout_items.retain(|item| { - keep_forked_rollout_item(item, preserve_reference_context_item) - && !matches!( - item, - RolloutItem::ResponseItem(response_item) - if is_multi_agent_v2_usage_hint_message( - response_item, - &multi_agent_v2_usage_hint_texts_to_filter, - ) - ) - }); - for item in &mut forked_rollout_items { - if let RolloutItem::Compacted(compacted) = item - && let Some(replacement_history) = compacted.replacement_history.as_mut() - { - replacement_history.retain(|response_item| { - !is_multi_agent_v2_usage_hint_message( - response_item, - &multi_agent_v2_usage_hint_texts_to_filter, - ) - }); - } - } - if preserve_reference_context_item - && multi_agent_version == MultiAgentVersion::V2 - && config.multi_agent_v2.usage_hint_enabled - && let Some(subagent_usage_hint_text) = - config.multi_agent_v2.subagent_usage_hint_text.clone() - && let Some(subagent_usage_hint_message) = - crate::context_manager::updates::build_developer_update_item(vec![ - subagent_usage_hint_text, - ]) - { - forked_rollout_items.push(RolloutItem::ResponseItem(subagent_usage_hint_message)); - } - - state - .fork_thread_with_source( - config.clone(), - InitialHistory::Forked(forked_rollout_items), - self.clone(), - session_source, - /*thread_source*/ Some(ThreadSource::Subagent), - /*parent_thread_id*/ Some(parent_thread_id), - /*forked_from_thread_id*/ Some(parent_thread_id), - inherited_shell_snapshot, - inherited_exec_policy, - options.environments.clone(), - ) - .await - } - - /// Resume an existing agent thread from a recorded rollout file. - pub(crate) async fn resume_agent_from_rollout( - &self, - config: Config, - thread_id: ThreadId, - session_source: SessionSource, - ) -> CodexResult { - let root_depth = thread_spawn_depth(&session_source).unwrap_or(0); - let resumed_thread_id = Box::pin(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_path: None, - agent_nickname: None, - agent_role: None, - }); - match Box::pin(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, - config: Config, - thread_id: ThreadId, - session_source: SessionSource, - ) -> CodexResult { - let state = self.upgrade()?; - let state_db_ctx = state.state_db(); - let stored_thread = state - .read_stored_thread(ReadThreadParams { - thread_id, - include_archived: true, - include_history: true, - }) - .await?; - let history = stored_thread - .history - .ok_or_else(|| CodexErr::ThreadNotFound(thread_id))? - .items; - let initial_history = InitialHistory::Resumed(ResumedHistory { - conversation_id: thread_id, - history, - rollout_path: stored_thread.rollout_path, - }); - let parent_thread_id = stored_thread.parent_thread_id; - let multi_agent_version = state - .effective_multi_agent_version_for_spawn( - &initial_history, - Some(&session_source), - parent_thread_id, - /*forked_from_thread_id*/ None, - &config, - ) - .await; - let agent_max_threads = config.effective_agent_max_threads(multi_agent_version); - let mut reservation = self.state.reserve_spawn_slot(agent_max_threads)?; - let (session_source, agent_metadata) = match session_source { - SessionSource::SubAgent(SubAgentSource::ThreadSpawn { - parent_thread_id, - depth, - agent_path, - agent_role: _, - agent_nickname: _, - }) => { - let (resumed_agent_nickname, resumed_agent_role) = - if let Some(state_db_ctx) = state_db_ctx.as_ref() { - match state_db_ctx.get_thread(thread_id).await { - Ok(Some(metadata)) => (metadata.agent_nickname, metadata.agent_role), - Ok(None) | Err(_) => (None, None), - } - } else { - (None, None) - }; - self.prepare_thread_spawn( - &mut reservation, - &config, - parent_thread_id, - depth, - agent_path, - resumed_agent_role, - resumed_agent_nickname, - )? - } - other => (other, AgentMetadata::default()), - }; - let notification_source = session_source.clone(); - let inherited_shell_snapshot = self - .inherited_shell_snapshot_for_source(&state, Some(&session_source)) - .await; - let inherited_exec_policy = self - .inherited_exec_policy_for_source(&state, Some(&session_source), &config) - .await; - - let resumed_thread = state - .resume_thread_with_history_with_source(ResumeThreadWithHistoryOptions { - config: config.clone(), - initial_history, - agent_control: self.clone(), - session_source, - parent_thread_id, - inherited_shell_snapshot, - inherited_exec_policy, - }) - .await?; - let mut agent_metadata = agent_metadata; - agent_metadata.agent_id = Some(resumed_thread.thread_id); - reservation.commit(agent_metadata.clone()); - // 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); - if multi_agent_version != MultiAgentVersion::V2 { - let child_reference = agent_metadata - .agent_path - .as_ref() - .map(ToString::to_string) - .unwrap_or_else(|| resumed_thread.thread_id.to_string()); - self.maybe_start_completion_watcher( - resumed_thread.thread_id, - Some(notification_source.clone()), - child_reference, - agent_metadata.agent_path.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) - } - /// Send rich user input items to an existing agent thread. pub(crate) async fn send_input( &self, @@ -753,7 +182,7 @@ impl AgentControl { agent_id, &state, state - .send_op(agent_id, Op::InterAgentCommunication { communication }) + .enqueue_inter_agent_communication(agent_id, communication) .await, ) .await; @@ -771,7 +200,15 @@ impl AgentControl { /// Interrupt the current task for an existing agent thread. pub(crate) async fn interrupt_agent(&self, agent_id: ThreadId) -> CodexResult { let state = self.upgrade()?; - state.send_op(agent_id, Op::Interrupt).await + let result = state.send_op(agent_id, Op::Interrupt).await; + let agent_died = matches!(result, Err(CodexErr::InternalAgentDied)); + let result = self + .handle_thread_request_result(agent_id, &state, result) + .await; + if agent_died { + self.residency.lock().await.remove(agent_id); + } + result } async fn handle_thread_request_result( @@ -782,86 +219,9 @@ impl AgentControl { ) -> CodexResult { if matches!(result, Err(CodexErr::InternalAgentDied)) { let _ = state.remove_thread(&agent_id).await; - self.state.release_spawned_thread(agent_id); - } - result - } - - /// 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()?; - let result = 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 = if matches!(thread.agent_status().await, AgentStatus::Shutdown) { - Ok(String::new()) - } else { - state.send_op(agent_id, Op::Shutdown {}).await - }; - thread.wait_until_terminated().await; - result - } else { - 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()?; - let known_agent = self.state.agent_metadata_for_thread(agent_id).is_some(); - match state.get_thread(agent_id).await { - Ok(thread) => { - if 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}"); - } - } - Err(CodexErr::ThreadNotFound(_)) if known_agent => { - if let Some(state_db_ctx) = state.state_db() - && let Err(err) = state_db_ctx - .set_thread_spawn_edge_status( - agent_id, - DirectionalThreadSpawnEdgeStatus::Closed, - ) - .await - { - return Err(CodexErr::Fatal(format!( - "failed to persist stale thread-spawn edge status for {agent_id}: {err}" - ))); - } - } - Err(CodexErr::ThreadNotFound(_)) => {} - Err(err) => { - warn!("failed to inspect agent before close {agent_id}: {err}"); - } - } - match Box::pin(self.shutdown_agent_tree(agent_id)).await { - Err(CodexErr::ThreadNotFound(_)) | Err(CodexErr::InternalAgentDied) if known_agent => { - Ok(String::new()) - } - result => result, - } - } - - /// 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), + self.state.release_execution_slot(agent_id); + if let Ok(mut residency) = self.residency.try_lock() { + residency.remove(agent_id); } } result @@ -902,6 +262,28 @@ impl AgentControl { Ok(thread_ids) } + pub(crate) async fn format_environment_context_subagents( + &self, + parent_thread_id: ThreadId, + ) -> String { + let Ok(agents) = self.open_thread_spawn_children(parent_thread_id).await else { + return String::new(); + }; + + agents + .into_iter() + .map(|(thread_id, metadata)| { + let reference = metadata + .agent_path + .as_ref() + .map(|agent_path| agent_path.name().to_string()) + .unwrap_or_else(|| thread_id.to_string()); + format_subagent_context_line(reference.as_str(), metadata.agent_nickname.as_deref()) + }) + .collect::>() + .join("\n") + } + pub(crate) async fn get_agent_config_snapshot( &self, agent_id: ThreadId, @@ -930,8 +312,28 @@ impl AgentControl { if let Some(thread_id) = self.state.agent_id_for_path(&agent_path) { return Ok(thread_id); } + let state = self.upgrade()?; + let root_thread_id = self + .state + .agent_id_for_path(&AgentPath::root()) + .ok_or_else(|| { + CodexErr::UnsupportedOperation("root agent is unavailable".to_string()) + })?; + if let Some(state_db_ctx) = state.state_db() + && let Some(thread_id) = state_db_ctx + .find_open_thread_spawn_descendant_by_path(root_thread_id, agent_path.as_str()) + .await + .map_err(|err| { + CodexErr::Fatal(format!( + "failed to resolve agent path `{agent_path}`: {err}" + )) + })? + { + self.ensure_agent_known(thread_id).await?; + return Ok(thread_id); + } Err(CodexErr::UnsupportedOperation(format!( - "live agent path `{}` not found", + "agent path `{}` not found", agent_path.as_str() ))) } @@ -946,269 +348,6 @@ impl AgentControl { Ok(thread.subscribe_status()) } - pub(crate) async fn format_environment_context_subagents( - &self, - parent_thread_id: ThreadId, - ) -> String { - let Ok(agents) = self.open_thread_spawn_children(parent_thread_id).await else { - return String::new(); - }; - - agents - .into_iter() - .map(|(thread_id, metadata)| { - let reference = metadata - .agent_path - .as_ref() - .map(|agent_path| agent_path.name().to_string()) - .unwrap_or_else(|| thread_id.to_string()); - format_subagent_context_line(reference.as_str(), metadata.agent_nickname.as_deref()) - }) - .collect::>() - .join("\n") - } - - pub(crate) async fn list_agents( - &self, - current_session_source: &SessionSource, - path_prefix: Option<&str>, - ) -> CodexResult> { - let state = self.upgrade()?; - let resolved_prefix = path_prefix - .map(|prefix| { - current_session_source - .get_agent_path() - .unwrap_or_else(AgentPath::root) - .resolve(prefix) - .map_err(CodexErr::UnsupportedOperation) - }) - .transpose()?; - - let mut live_agents = self.state.live_agents(); - live_agents.sort_by(|left, right| { - left.agent_path - .as_deref() - .unwrap_or_default() - .cmp(right.agent_path.as_deref().unwrap_or_default()) - .then_with(|| { - left.agent_id - .map(|id| id.to_string()) - .unwrap_or_default() - .cmp(&right.agent_id.map(|id| id.to_string()).unwrap_or_default()) - }) - }); - - let root_path = AgentPath::root(); - let mut agents = Vec::with_capacity(live_agents.len().saturating_add(1)); - if resolved_prefix - .as_ref() - .is_none_or(|prefix| agent_matches_prefix(Some(&root_path), prefix)) - && let Some(root_thread_id) = self.state.agent_id_for_path(&root_path) - && let Ok(root_thread) = state.get_thread(root_thread_id).await - { - agents.push(ListedAgent { - agent_name: root_path.to_string(), - agent_status: root_thread.agent_status().await, - last_task_message: Some(ROOT_LAST_TASK_MESSAGE.to_string()), - }); - } - - for metadata in live_agents { - let Some(thread_id) = metadata.agent_id else { - continue; - }; - if resolved_prefix - .as_ref() - .is_some_and(|prefix| !agent_matches_prefix(metadata.agent_path.as_ref(), prefix)) - { - continue; - } - - let Ok(thread) = state.get_thread(thread_id).await else { - continue; - }; - let agent_name = metadata - .agent_path - .as_ref() - .map(ToString::to_string) - .unwrap_or_else(|| thread_id.to_string()); - let last_task_message = metadata.last_task_message.clone(); - agents.push(ListedAgent { - agent_name, - agent_status: thread.agent_status().await, - last_task_message, - }); - } - - Ok(agents) - } - - /// Starts a detached watcher for sub-agents spawned from another thread. - /// - /// This is only enabled for `SubAgentSource::ThreadSpawn`, where a parent thread exists and - /// can receive completion notifications. - fn maybe_start_completion_watcher( - &self, - child_thread_id: ThreadId, - session_source: Option, - child_reference: String, - child_agent_path: Option, - ) { - let Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { - parent_thread_id, .. - })) = session_source - else { - return; - }; - let control = self.clone(); - tokio::spawn(async move { - let status = match control.subscribe_status(child_thread_id).await { - Ok(mut status_rx) => { - let mut status = status_rx.borrow().clone(); - while !is_final(&status) { - if status_rx.changed().await.is_err() { - status = control.get_status(child_thread_id).await; - break; - } - status = status_rx.borrow().clone(); - } - status - } - Err(_) => control.get_status(child_thread_id).await, - }; - if !is_final(&status) { - return; - } - - let Ok(state) = control.upgrade() else { - return; - }; - let child_thread = state.get_thread(child_thread_id).await.ok(); - let message = format_subagent_notification_message(child_reference.as_str(), &status); - let child_uses_multi_agent_v2 = match child_thread.as_ref() { - Some(child_thread) => { - child_thread.multi_agent_version() == Some(MultiAgentVersion::V2) - } - None => true, - }; - if child_agent_path.is_some() && child_uses_multi_agent_v2 { - let Some(child_agent_path) = child_agent_path.clone() else { - return; - }; - let Some(parent_agent_path) = child_agent_path - .as_str() - .rsplit_once('/') - .and_then(|(parent, _)| AgentPath::try_from(parent).ok()) - else { - return; - }; - let communication = InterAgentCommunication::new( - child_agent_path, - parent_agent_path, - Vec::new(), - message, - /*trigger_turn*/ false, - ); - let _ = control - .send_inter_agent_communication(parent_thread_id, communication) - .await; - return; - } - let Ok(parent_thread) = state.get_thread(parent_thread_id).await else { - return; - }; - parent_thread - .inject_user_message_without_turn(message) - .await; - }); - } - - #[allow(clippy::too_many_arguments)] - fn prepare_thread_spawn( - &self, - reservation: &mut crate::agent::registry::SpawnReservation, - config: &Config, - parent_thread_id: ThreadId, - depth: i32, - agent_path: Option, - agent_role: Option, - preferred_agent_nickname: Option, - ) -> CodexResult<(SessionSource, AgentMetadata)> { - if depth == 1 { - self.state.register_root_thread(parent_thread_id); - } - if let Some(agent_path) = agent_path.as_ref() { - reservation.reserve_agent_path(agent_path)?; - } - let candidate_names = agent_nickname_candidates(config, agent_role.as_deref()); - let candidate_name_refs: Vec<&str> = candidate_names.iter().map(String::as_str).collect(); - let agent_nickname = Some(reservation.reserve_agent_nickname_with_preference( - &candidate_name_refs, - preferred_agent_nickname.as_deref(), - )?); - let session_source = SessionSource::SubAgent(SubAgentSource::ThreadSpawn { - parent_thread_id, - depth, - agent_path: agent_path.clone(), - agent_nickname: agent_nickname.clone(), - agent_role: agent_role.clone(), - }); - let agent_metadata = AgentMetadata { - agent_id: None, - agent_path, - agent_nickname, - agent_role, - last_task_message: None, - }; - Ok((session_source, agent_metadata)) - } - - fn upgrade(&self) -> CodexResult> { - self.manager - .upgrade() - .ok_or_else(|| CodexErr::UnsupportedOperation("thread manager dropped".to_string())) - } - - async fn inherited_shell_snapshot_for_source( - &self, - state: &Arc, - session_source: Option<&SessionSource>, - ) -> Option> { - let Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { - parent_thread_id, .. - })) = session_source - else { - return None; - }; - - let parent_thread = state.get_thread(*parent_thread_id).await.ok()?; - parent_thread.codex.session.user_shell().shell_snapshot() - } - - async fn inherited_exec_policy_for_source( - &self, - state: &Arc, - session_source: Option<&SessionSource>, - child_config: &Config, - ) -> Option> { - let Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { - parent_thread_id, .. - })) = session_source - else { - return None; - }; - - let parent_thread = state.get_thread(*parent_thread_id).await.ok()?; - let parent_config = parent_thread.codex.session.get_config().await; - if !crate::exec_policy::child_uses_parent_exec_policy(&parent_config, child_config) { - return None; - } - - Some(Arc::clone( - &parent_thread.codex.session.services.exec_policy, - )) - } - async fn open_thread_spawn_children( &self, parent_thread_id: ThreadId, @@ -1254,31 +393,6 @@ impl AgentControl { 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(SessionSource::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, @@ -1304,20 +418,12 @@ impl AgentControl { Ok(descendants) } -} -fn agent_matches_prefix(agent_path: Option<&AgentPath>, prefix: &AgentPath) -> bool { - if prefix.is_root() { - return true; + pub(super) fn upgrade(&self) -> CodexResult> { + self.manager + .upgrade() + .ok_or_else(|| CodexErr::UnsupportedOperation("thread manager dropped".to_string())) } - - agent_path.is_some_and(|agent_path| { - agent_path == prefix - || agent_path - .as_str() - .strip_prefix(prefix.as_str()) - .is_some_and(|suffix| suffix.starts_with('/')) - }) } pub(crate) fn render_input_preview(initial_operation: &Op) -> String { @@ -1354,12 +460,6 @@ fn non_empty_task_message(message: String) -> Option { (!message.is_empty()).then_some(message) } -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"] mod tests; diff --git a/codex-rs/core/src/agent/control/legacy.rs b/codex-rs/core/src/agent/control/legacy.rs new file mode 100644 index 000000000000..dcaab20fd5ca --- /dev/null +++ b/codex-rs/core/src/agent/control/legacy.rs @@ -0,0 +1,153 @@ +use super::*; + +impl AgentControl { + /// Shut down a live agent and remove its logical identity from the registry. + pub(crate) async fn shutdown_and_forget_agent( + &self, + agent_id: ThreadId, + ) -> CodexResult { + let state = self.upgrade()?; + let result = 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 = if matches!(thread.agent_status().await, AgentStatus::Shutdown) { + Ok(String::new()) + } else { + state.send_op(agent_id, Op::Shutdown {}).await + }; + thread.wait_until_terminated().await; + result + } else { + state.send_op(agent_id, Op::Shutdown {}).await + }; + let _ = state.remove_thread(&agent_id).await; + self.residency.lock().await.remove(agent_id); + self.state.release_spawned_thread(agent_id); + result + } + + /// 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 { + self.shutdown_and_forget_agent(agent_id).await + } + + /// Permanently close an agent and all open descendants. + pub(crate) async fn close_agent(&self, agent_id: ThreadId) -> CodexResult { + self.close_agent_v1(agent_id).await + } + + /// Permanently close a V1 agent and all open descendants. + pub(crate) async fn close_agent_v1(&self, agent_id: ThreadId) -> CodexResult { + let state = self.upgrade()?; + let known_agent = self.state.agent_metadata_for_thread(agent_id).is_some(); + let mut agent_ids = vec![agent_id]; + if let Some(state_db_ctx) = state.state_db() { + let descendants = state_db_ctx + .list_thread_spawn_descendants_with_status( + agent_id, + DirectionalThreadSpawnEdgeStatus::Open, + ) + .await + .map_err(|err| { + CodexErr::Fatal(format!( + "failed to list open descendants for {agent_id}: {err}" + )) + })?; + agent_ids.extend(descendants); + for thread_id in &agent_ids { + state_db_ctx + .set_thread_spawn_edge_status( + *thread_id, + DirectionalThreadSpawnEdgeStatus::Closed, + ) + .await + .map_err(|err| { + CodexErr::Fatal(format!( + "failed to persist closed agent state for {thread_id}: {err}" + )) + })?; + } + } else { + agent_ids.extend(self.live_thread_spawn_descendants(agent_id).await?); + } + let mut target_result = Ok(String::new()); + for thread_id in agent_ids { + let result = self.shutdown_and_forget_agent(thread_id).await; + if thread_id == agent_id { + target_result = result; + } else if let Err(err) = result + && !matches!( + err, + CodexErr::ThreadNotFound(_) | CodexErr::InternalAgentDied + ) + { + return Err(err); + } + } + match target_result { + Err(CodexErr::ThreadNotFound(_)) | Err(CodexErr::InternalAgentDied) if known_agent => { + Ok(String::new()) + } + result => result, + } + } + + /// Shut down `agent_id` and any live descendants reachable from the in-memory spawn tree. + pub(crate) 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_and_forget_agent(agent_id).await; + for descendant_id in descendant_ids { + match self.shutdown_and_forget_agent(descendant_id).await { + Ok(_) | Err(CodexErr::ThreadNotFound(_)) | Err(CodexErr::InternalAgentDied) => {} + Err(err) => return Err(err), + } + } + result + } + + pub(super) fn maybe_start_completion_watcher( + &self, + child_thread_id: ThreadId, + session_source: Option, + child_reference: String, + ) { + let Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id, .. + })) = session_source + else { + return; + }; + let control = self.clone(); + tokio::spawn(async move { + let status = match control.subscribe_status(child_thread_id).await { + Ok(mut status_rx) => { + let mut status = status_rx.borrow().clone(); + while !is_final(&status) { + if status_rx.changed().await.is_err() { + status = control.get_status(child_thread_id).await; + break; + } + status = status_rx.borrow().clone(); + } + status + } + Err(_) => control.get_status(child_thread_id).await, + }; + if !is_final(&status) { + return; + } + + let Ok(state) = control.upgrade() else { + return; + }; + let message = format_subagent_notification_message(child_reference.as_str(), &status); + let Ok(parent_thread) = state.get_thread(parent_thread_id).await else { + return; + }; + parent_thread + .inject_user_message_without_turn(message) + .await; + }); + } +} diff --git a/codex-rs/core/src/agent/control/resident.rs b/codex-rs/core/src/agent/control/resident.rs new file mode 100644 index 000000000000..323dab014ab4 --- /dev/null +++ b/codex-rs/core/src/agent/control/resident.rs @@ -0,0 +1,368 @@ +use super::*; + +const ROOT_LAST_TASK_MESSAGE: &str = "Main thread"; + +fn max_resident_subagents(config: &Config) -> usize { + MIN_RESIDENT_SUBAGENTS.max( + config + .multi_agent_v2 + .max_concurrent_threads_per_session + .saturating_add(4), + ) +} + +impl AgentControl { + pub(crate) async fn reserve_execution_slot_for_turn_start( + &self, + session: &Session, + ) -> CodexResult> { + if session.multi_agent_version() != Some(MultiAgentVersion::V2) { + return Ok(None); + } + let metadata = self + .state + .agent_metadata_for_thread(session.thread_id) + .ok_or_else(|| { + CodexErr::Fatal(format!( + "missing agent metadata for V2 thread {}", + session.thread_id + )) + })?; + if metadata.agent_path.as_ref().is_some_and(AgentPath::is_root) { + return Ok(None); + } + let config = session.get_config().await; + let max_threads = config + .effective_agent_max_threads(MultiAgentVersion::V2) + .unwrap_or_default(); + self.state + .reserve_execution_slot(session.thread_id, max_threads) + .map(Some) + } + + pub(crate) async fn send_message_to_agent( + &self, + config: &Config, + agent_id: ThreadId, + communication: InterAgentCommunication, + ) -> CodexResult { + let metadata = self.ensure_agent_known(agent_id).await?; + let state = self.upgrade()?; + if metadata.agent_path.as_ref().is_some_and(AgentPath::is_root) { + state.get_thread(agent_id).await?; + return self + .send_inter_agent_communication(agent_id, communication) + .await; + } + let mut residency = self.residency.lock().await; + self.ensure_agent_loaded(&state, &mut residency, config, agent_id) + .await?; + self.send_inter_agent_communication(agent_id, communication) + .await + } + + pub(crate) async fn send_followup_to_agent( + &self, + config: &Config, + agent_id: ThreadId, + communication: InterAgentCommunication, + ) -> CodexResult { + self.ensure_agent_known(agent_id).await?; + let state = self.upgrade()?; + let max_threads = config + .effective_agent_max_threads(MultiAgentVersion::V2) + .unwrap_or_default(); + let mut residency = self.residency.lock().await; + let reservation = self.state.reserve_execution_slot(agent_id, max_threads)?; + self.ensure_agent_loaded(&state, &mut residency, config, agent_id) + .await?; + let sub_id = self + .send_inter_agent_communication(agent_id, communication) + .await?; + reservation.commit(); + drop(residency); + let thread = state.get_thread(agent_id).await?; + thread + .codex + .session + .maybe_start_turn_for_pending_work_with_sub_id(sub_id.clone()) + .await; + Ok(sub_id) + } + + pub(crate) async fn release_execution_slot_if_idle(&self, session: &Session) { + let Some(metadata) = self.state.agent_metadata_for_thread(session.thread_id) else { + return; + }; + if metadata.agent_path.as_ref().is_none_or(AgentPath::is_root) { + return; + } + let residency = self.residency.lock().await; + if session.active_turn.lock().await.is_none() + && !session.input_queue.has_trigger_turn_mailbox_items().await + && self.state.release_execution_slot(session.thread_id) + { + drop(residency); + let Ok(state) = self.upgrade() else { + return; + }; + for agent in self.state.live_agents() { + let Some(thread_id) = agent.agent_id else { + continue; + }; + if thread_id == session.thread_id { + continue; + } + let Ok(thread) = state.get_thread(thread_id).await else { + continue; + }; + if thread + .codex + .session + .input_queue + .has_trigger_turn_mailbox_items() + .await + { + thread + .codex + .session + .maybe_start_turn_for_pending_work() + .await; + } + } + } + } + + pub(crate) async fn ensure_agent_known( + &self, + agent_id: ThreadId, + ) -> CodexResult { + if let Some(metadata) = self.state.agent_metadata_for_thread(agent_id) { + return Ok(metadata); + } + + let state = self.upgrade()?; + let stored_thread = state + .read_stored_thread(ReadThreadParams { + thread_id: agent_id, + include_archived: true, + include_history: false, + }) + .await?; + if let Some(state_db_ctx) = state.state_db() + && !state_db_ctx + .is_thread_spawn_edge_open(agent_id) + .await + .map_err(|err| { + CodexErr::Fatal(format!( + "failed to read spawned agent state for {agent_id}: {err}" + )) + })? + { + return Err(CodexErr::InternalAgentDied); + } + if !matches!( + &stored_thread.source, + SessionSource::SubAgent(SubAgentSource::ThreadSpawn { .. }) + ) { + return Err(CodexErr::ThreadNotFound(agent_id)); + } + let metadata = agent_metadata_from_session_source(agent_id, &stored_thread.source); + self.state.register_agent(metadata.clone()); + Ok(metadata) + } + + async fn ensure_agent_loaded( + &self, + state: &Arc, + residency: &mut ResidentAgents, + config: &Config, + agent_id: ThreadId, + ) -> CodexResult<()> { + if state.get_thread(agent_id).await.is_ok() { + if !residency.contains(agent_id) { + self.make_resident_room(state, residency, config).await?; + } + residency.touch(agent_id); + return Ok(()); + } + + self.make_resident_room(state, residency, config).await?; + let stored_thread = state + .read_stored_thread(ReadThreadParams { + thread_id: agent_id, + include_archived: true, + include_history: false, + }) + .await?; + self.resume_single_agent_from_rollout(config.clone(), agent_id, stored_thread.source) + .await?; + residency.touch(agent_id); + Ok(()) + } + + pub(super) async fn make_resident_room( + &self, + state: &Arc, + residency: &mut ResidentAgents, + config: &Config, + ) -> CodexResult<()> { + if residency.len() < max_resident_subagents(config) { + return Ok(()); + } + let mut eviction_candidate = None; + for agent_id in residency.snapshot() { + if self.state.is_execution_active(agent_id) { + continue; + } + if let Ok(thread) = state.get_thread(agent_id).await + && thread + .codex + .session + .input_queue + .has_pending_mailbox_items() + .await + { + continue; + } + eviction_candidate = Some(agent_id); + break; + } + let Some(agent_id) = eviction_candidate else { + return Err(CodexErr::Fatal( + "resident agent cache is full with no idle agent available to unload".to_string(), + )); + }; + self.unload_resident_agent(state, residency, agent_id).await + } + + async fn unload_resident_agent( + &self, + state: &Arc, + residency: &mut ResidentAgents, + agent_id: ThreadId, + ) -> CodexResult<()> { + let Ok(thread) = state.get_thread(agent_id).await else { + residency.remove(agent_id); + return Ok(()); + }; + thread.ensure_rollout_materialized().await; + thread.flush_rollout().await?; + let mut wait_for_termination = false; + if !matches!(thread.agent_status().await, AgentStatus::Shutdown) { + match state.send_op(agent_id, Op::Shutdown {}).await { + Ok(_) => wait_for_termination = true, + Err(CodexErr::ThreadNotFound(_)) | Err(CodexErr::InternalAgentDied) => {} + Err(err) => return Err(err), + } + } + if wait_for_termination { + thread.wait_until_terminated().await; + } + let _ = state.remove_thread(&agent_id).await; + residency.remove(agent_id); + self.state.release_execution_slot(agent_id); + self.state.clear_last_task_message(agent_id); + Ok(()) + } + + pub(crate) async fn list_agents( + &self, + current_session_source: &SessionSource, + path_prefix: Option<&str>, + ) -> CodexResult> { + let state = self.upgrade()?; + let resolved_prefix = path_prefix + .map(|prefix| { + current_session_source + .get_agent_path() + .unwrap_or_else(AgentPath::root) + .resolve(prefix) + .map_err(CodexErr::UnsupportedOperation) + }) + .transpose()?; + + let resident_agent_ids = self.residency.lock().await.snapshot(); + let mut resident_agents = resident_agent_ids + .into_iter() + .map(|agent_id| { + self.state + .agent_metadata_for_thread(agent_id) + .unwrap_or(AgentMetadata { + agent_id: Some(agent_id), + ..Default::default() + }) + }) + .collect::>(); + resident_agents.sort_by(|left, right| { + left.agent_path + .as_deref() + .unwrap_or_default() + .cmp(right.agent_path.as_deref().unwrap_or_default()) + .then_with(|| { + left.agent_id + .map(|id| id.to_string()) + .unwrap_or_default() + .cmp(&right.agent_id.map(|id| id.to_string()).unwrap_or_default()) + }) + }); + + let root_path = AgentPath::root(); + let mut agents = Vec::with_capacity(resident_agents.len().saturating_add(1)); + if resolved_prefix + .as_ref() + .is_none_or(|prefix| agent_matches_prefix(Some(&root_path), prefix)) + && let Some(root_thread_id) = self.state.agent_id_for_path(&root_path) + && let Ok(root_thread) = state.get_thread(root_thread_id).await + { + agents.push(ListedAgent { + agent_name: root_path.to_string(), + agent_status: root_thread.agent_status().await, + last_task_message: Some(ROOT_LAST_TASK_MESSAGE.to_string()), + }); + } + + for metadata in resident_agents { + let Some(thread_id) = metadata.agent_id else { + continue; + }; + if resolved_prefix + .as_ref() + .is_some_and(|prefix| !agent_matches_prefix(metadata.agent_path.as_ref(), prefix)) + { + continue; + } + + let agent_name = metadata + .agent_path + .as_ref() + .map(ToString::to_string) + .unwrap_or_else(|| thread_id.to_string()); + let last_task_message = metadata.last_task_message.clone(); + let Ok(thread) = state.get_thread(thread_id).await else { + continue; + }; + agents.push(ListedAgent { + agent_name, + agent_status: thread.agent_status().await, + last_task_message, + }); + } + + Ok(agents) + } +} + +fn agent_matches_prefix(agent_path: Option<&AgentPath>, prefix: &AgentPath) -> bool { + if prefix.is_root() { + return true; + } + + agent_path.is_some_and(|agent_path| { + agent_path == prefix + || agent_path + .as_str() + .strip_prefix(prefix.as_str()) + .is_some_and(|suffix| suffix.starts_with('/')) + }) +} diff --git a/codex-rs/core/src/agent/control/spawn.rs b/codex-rs/core/src/agent/control/spawn.rs new file mode 100644 index 000000000000..170a9efbfd2c --- /dev/null +++ b/codex-rs/core/src/agent/control/spawn.rs @@ -0,0 +1,818 @@ +use super::*; + +const AGENT_NAMES: &str = include_str!("../agent_names.txt"); + +struct SpawnAgentThreadInheritance { + shell_snapshot: Option>, + exec_policy: Option>, +} + +fn default_agent_nickname_list() -> Vec<&'static str> { + AGENT_NAMES + .lines() + .map(str::trim) + .filter(|name| !name.is_empty()) + .collect() +} + +pub(super) fn agent_nickname_candidates(config: &Config, role_name: Option<&str>) -> Vec { + let role_name = role_name.unwrap_or(DEFAULT_ROLE_NAME); + if let Some(candidates) = + resolve_role_config(config, role_name).and_then(|role| role.nickname_candidates.clone()) + { + return candidates; + } + + default_agent_nickname_list() + .into_iter() + .map(ToOwned::to_owned) + .collect() +} + +fn resume_session_source( + requested_source: SessionSource, + stored_source: SessionSource, + multi_agent_version: MultiAgentVersion, +) -> SessionSource { + match multi_agent_version { + MultiAgentVersion::V2 => resume_session_source_v2(requested_source, stored_source), + MultiAgentVersion::Disabled | MultiAgentVersion::V1 => match stored_source { + source @ SessionSource::SubAgent(SubAgentSource::ThreadSpawn { .. }) => source, + _ => requested_source, + }, + } +} + +fn resume_session_source_v2( + requested_source: SessionSource, + stored_source: SessionSource, +) -> SessionSource { + let stored_thread_spawn = match stored_source { + SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id, + depth, + agent_path, + agent_role, + .. + }) => Some((parent_thread_id, depth, agent_path, agent_role)), + _ => None, + }; + + match requested_source { + SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id, + depth, + agent_path, + agent_role, + .. + }) => { + let stored_agent_path = stored_thread_spawn + .as_ref() + .and_then(|(_, _, agent_path, _)| agent_path.clone()); + let stored_agent_role = stored_thread_spawn + .as_ref() + .and_then(|(_, _, _, agent_role)| agent_role.clone()); + SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id, + depth, + agent_path: agent_path.or(stored_agent_path), + agent_nickname: None, + agent_role: agent_role.or(stored_agent_role), + }) + } + other => match stored_thread_spawn { + Some((parent_thread_id, depth, agent_path, agent_role)) => { + SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id, + depth, + agent_path, + agent_nickname: None, + agent_role, + }) + } + None => other, + }, + } +} + +fn keep_forked_rollout_item(item: &RolloutItem, preserve_reference_context_item: bool) -> bool { + match item { + RolloutItem::ResponseItem(ResponseItem::Message { role, phase, .. }) => match role.as_str() + { + "system" | "developer" | "user" => true, + "assistant" => *phase == Some(MessagePhase::FinalAnswer), + _ => false, + }, + RolloutItem::ResponseItem( + ResponseItem::AgentMessage { .. } + | ResponseItem::Reasoning { .. } + | ResponseItem::LocalShellCall { .. } + | ResponseItem::FunctionCall { .. } + | ResponseItem::ToolSearchCall { .. } + | ResponseItem::FunctionCallOutput { .. } + | ResponseItem::CustomToolCall { .. } + | ResponseItem::CustomToolCallOutput { .. } + | ResponseItem::ToolSearchOutput { .. } + | ResponseItem::WebSearchCall { .. } + | ResponseItem::ImageGenerationCall { .. } + | ResponseItem::Compaction { .. } + | ResponseItem::CompactionTrigger + | ResponseItem::ContextCompaction { .. } + | ResponseItem::Other, + ) => false, + // Full-history forks preserve the cached prompt prefix and can keep diffing + // from the parent's durable baseline. Truncated forks drop part of that prompt, + // so they must rebuild context on their first child turn. + RolloutItem::TurnContext(_) => preserve_reference_context_item, + RolloutItem::Compacted(_) | RolloutItem::EventMsg(_) | RolloutItem::SessionMeta(_) => true, + } +} + +fn is_multi_agent_v2_usage_hint_message(item: &ResponseItem, usage_hint_texts: &[String]) -> bool { + let ResponseItem::Message { role, content, .. } = item else { + return false; + }; + if role != "developer" { + return false; + } + let [ContentItem::InputText { text }] = content.as_slice() else { + return false; + }; + + usage_hint_texts + .iter() + .any(|usage_hint_text| usage_hint_text == text) +} + +impl AgentControl { + /// Spawn a new agent thread and submit the initial prompt. + #[cfg(test)] + pub(crate) async fn spawn_agent( + &self, + config: Config, + initial_operation: Op, + session_source: Option, + ) -> CodexResult { + let spawned_agent = Box::pin(self.spawn_agent_internal( + config, + initial_operation, + session_source, + SpawnAgentOptions::default(), + )) + .await?; + Ok(spawned_agent.thread_id) + } + + /// Spawn an agent thread with some metadata. + pub(crate) async fn spawn_agent_with_metadata( + &self, + config: Config, + initial_operation: Op, + session_source: Option, + options: SpawnAgentOptions, // TODO(jif) drop with new fork. + ) -> CodexResult { + Box::pin(self.spawn_agent_internal(config, initial_operation, session_source, options)) + .await + } + + async fn spawn_agent_internal( + &self, + config: Config, + initial_operation: Op, + session_source: Option, + options: SpawnAgentOptions, + ) -> CodexResult { + let state = self.upgrade()?; + let multi_agent_version = state + .effective_multi_agent_version_for_spawn( + &InitialHistory::New, + session_source.as_ref(), + options.parent_thread_id, + /*forked_from_thread_id*/ None, + &config, + ) + .await; + let agent_max_threads = config.effective_agent_max_threads(multi_agent_version); + let mut reservation = self.state.reserve_spawn_slot(agent_max_threads)?; + let inheritance = SpawnAgentThreadInheritance { + shell_snapshot: self + .inherited_shell_snapshot_for_source(&state, session_source.as_ref()) + .await, + exec_policy: self + .inherited_exec_policy_for_source(&state, session_source.as_ref(), &config) + .await, + }; + let (session_source, mut agent_metadata) = match session_source { + Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id, + depth, + agent_path, + agent_role, + .. + })) => { + if let Some(agent_path) = agent_path.as_ref() { + self.ensure_agent_path_available(&state, parent_thread_id, depth, agent_path) + .await?; + } + let (session_source, agent_metadata) = self.prepare_thread_spawn( + &mut reservation, + &config, + parent_thread_id, + depth, + agent_path, + agent_role, + /*preferred_agent_nickname*/ None, + multi_agent_version, + )?; + (Some(session_source), agent_metadata) + } + other => (other, AgentMetadata::default()), + }; + let notification_source = session_source.clone(); + + let new_thread = match (session_source, options.fork_mode.as_ref(), inheritance) { + (Some(session_source), Some(_), inheritance) => { + Box::pin(self.spawn_forked_thread( + &state, + config.clone(), + session_source, + &options, + inheritance, + multi_agent_version, + )) + .await? + } + (Some(session_source), None, inheritance) => { + Box::pin(state.spawn_new_thread_with_source( + config.clone(), + self.clone(), + session_source, + options.parent_thread_id, + /*forked_from_thread_id*/ None, + /*thread_source*/ Some(ThreadSource::Subagent), + /*metrics_service_name*/ None, + inheritance.shell_snapshot, + inheritance.exec_policy, + options.environments.clone(), + )) + .await? + } + (None, _, _) => Box::pin(state.spawn_new_thread(config.clone(), self.clone())).await?, + }; + agent_metadata.agent_id = Some(new_thread.thread_id); + if multi_agent_version == MultiAgentVersion::V2 { + let mut residency = self.residency.lock().await; + if let Err(err) = async { + if !residency.contains(new_thread.thread_id) { + self.make_resident_room(&state, &mut residency, &config) + .await?; + } + residency.touch(new_thread.thread_id); + CodexResult::<()>::Ok(()) + } + .await + { + if state + .send_op(new_thread.thread_id, Op::Shutdown {}) + .await + .is_ok() + { + new_thread.thread.wait_until_terminated().await; + } + let _ = state.remove_thread(&new_thread.thread_id).await; + return Err(err); + } + } + reservation.commit(agent_metadata.clone()); + + if let Some(SessionSource::SubAgent( + subagent_source @ SubAgentSource::ThreadSpawn { + parent_thread_id, .. + }, + )) = notification_source.as_ref() + { + let client_metadata = match state.get_thread(*parent_thread_id).await { + Ok(parent_thread) => { + parent_thread + .codex + .session + .app_server_client_metadata() + .await + } + Err(error) => { + tracing::warn!( + error = %error, + parent_thread_id = %parent_thread_id, + "skipping subagent thread analytics: failed to load parent thread metadata" + ); + crate::session::session::AppServerClientMetadata { + client_name: None, + client_version: None, + } + } + }; + let thread_config = new_thread.thread.codex.thread_config_snapshot().await; + let parent_thread_id = thread_config.parent_thread_id; + emit_subagent_session_started( + &new_thread + .thread + .codex + .session + .services + .analytics_events_client, + client_metadata, + new_thread.thread.codex.session.session_id(), + new_thread.thread_id, + parent_thread_id, + thread_config, + subagent_source.clone(), + ); + } + + 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, initial_operation) + .await?; + if multi_agent_version != MultiAgentVersion::V2 { + let child_reference = agent_metadata + .agent_path + .as_ref() + .map(ToString::to_string) + .unwrap_or_else(|| new_thread.thread_id.to_string()); + self.maybe_start_completion_watcher( + new_thread.thread_id, + notification_source, + child_reference, + ); + } + + Ok(LiveAgent { + thread_id: new_thread.thread_id, + metadata: agent_metadata, + status: self.get_status(new_thread.thread_id).await, + }) + } + + async fn spawn_forked_thread( + &self, + state: &Arc, + config: Config, + session_source: SessionSource, + options: &SpawnAgentOptions, + inheritance: SpawnAgentThreadInheritance, + multi_agent_version: MultiAgentVersion, + ) -> CodexResult { + let SpawnAgentThreadInheritance { + shell_snapshot: inherited_shell_snapshot, + exec_policy: inherited_exec_policy, + } = inheritance; + if options.fork_parent_spawn_call_id.is_none() { + return Err(CodexErr::Fatal( + "spawn_agent fork requires a parent spawn call id".to_string(), + )); + } + let Some(fork_mode) = options.fork_mode.as_ref() else { + return Err(CodexErr::Fatal( + "spawn_agent fork requires a fork mode".to_string(), + )); + }; + let SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id, .. + }) = &session_source + else { + return Err(CodexErr::Fatal( + "spawn_agent fork requires a thread-spawn session source".to_string(), + )); + }; + + let parent_thread_id = *parent_thread_id; + let parent_thread = state.get_thread(parent_thread_id).await.ok(); + if let Some(parent_thread) = parent_thread.as_ref() { + parent_thread.ensure_rollout_materialized().await; + parent_thread.flush_rollout().await?; + } + + let parent_history = state + .read_stored_thread(ReadThreadParams { + thread_id: parent_thread_id, + include_archived: true, + include_history: true, + }) + .await? + .history + .ok_or_else(|| { + CodexErr::Fatal(format!( + "parent thread history unavailable for fork: {parent_thread_id}" + )) + })?; + + let mut forked_rollout_items = parent_history.items; + if let SpawnAgentForkMode::LastNTurns(last_n_turns) = fork_mode { + forked_rollout_items = + truncate_rollout_to_last_n_fork_turns(&forked_rollout_items, *last_n_turns); + } + let multi_agent_v2_usage_hint_texts_to_filter: Vec = + if let Some(parent_thread) = parent_thread.as_ref() { + if multi_agent_version == MultiAgentVersion::V2 { + let parent_config = parent_thread.codex.session.get_config().await; + [ + parent_config + .multi_agent_v2 + .root_agent_usage_hint_text + .clone(), + parent_config + .multi_agent_v2 + .subagent_usage_hint_text + .clone(), + ] + .into_iter() + .flatten() + .collect() + } else { + Vec::new() + } + } else if multi_agent_version == MultiAgentVersion::V2 { + [ + config.multi_agent_v2.root_agent_usage_hint_text.clone(), + config.multi_agent_v2.subagent_usage_hint_text.clone(), + ] + .into_iter() + .flatten() + .collect() + } else { + Vec::new() + }; + let preserve_reference_context_item = matches!(fork_mode, SpawnAgentForkMode::FullHistory); + forked_rollout_items.retain(|item| { + keep_forked_rollout_item(item, preserve_reference_context_item) + && !matches!( + item, + RolloutItem::ResponseItem(response_item) + if is_multi_agent_v2_usage_hint_message( + response_item, + &multi_agent_v2_usage_hint_texts_to_filter, + ) + ) + }); + for item in &mut forked_rollout_items { + if let RolloutItem::Compacted(compacted) = item + && let Some(replacement_history) = compacted.replacement_history.as_mut() + { + replacement_history.retain(|response_item| { + !is_multi_agent_v2_usage_hint_message( + response_item, + &multi_agent_v2_usage_hint_texts_to_filter, + ) + }); + } + } + if preserve_reference_context_item + && multi_agent_version == MultiAgentVersion::V2 + && config.multi_agent_v2.usage_hint_enabled + && let Some(subagent_usage_hint_text) = + config.multi_agent_v2.subagent_usage_hint_text.clone() + && let Some(subagent_usage_hint_message) = + crate::context_manager::updates::build_developer_update_item(vec![ + subagent_usage_hint_text, + ]) + { + forked_rollout_items.push(RolloutItem::ResponseItem(subagent_usage_hint_message)); + } + + state + .fork_thread_with_source( + config.clone(), + InitialHistory::Forked(forked_rollout_items), + self.clone(), + session_source, + /*thread_source*/ Some(ThreadSource::Subagent), + /*parent_thread_id*/ Some(parent_thread_id), + /*forked_from_thread_id*/ Some(parent_thread_id), + inherited_shell_snapshot, + inherited_exec_policy, + options.environments.clone(), + ) + .await + } + + /// Resume an existing agent thread from a recorded rollout file. + pub(crate) async fn resume_agent_from_rollout( + &self, + config: Config, + thread_id: ThreadId, + session_source: SessionSource, + ) -> CodexResult { + let root_depth = thread_spawn_depth(&session_source).unwrap_or(0); + let resumed_thread_id = Box::pin(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_path: None, + agent_nickname: None, + agent_role: None, + }); + match Box::pin(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) + } + + pub(super) fn resume_single_agent_from_rollout( + &self, + config: Config, + thread_id: ThreadId, + session_source: SessionSource, + ) -> BoxFuture<'_, CodexResult> { + Box::pin(async move { + let state = self.upgrade()?; + let stored_thread = state + .read_stored_thread(ReadThreadParams { + thread_id, + include_archived: true, + include_history: true, + }) + .await?; + let stored_session_source = stored_thread.source.clone(); + let history = stored_thread + .history + .ok_or_else(|| CodexErr::ThreadNotFound(thread_id))? + .items; + let initial_history = InitialHistory::Resumed(ResumedHistory { + conversation_id: thread_id, + history, + rollout_path: stored_thread.rollout_path, + }); + let parent_thread_id = stored_thread.parent_thread_id; + let multi_agent_version = state + .effective_multi_agent_version_for_spawn( + &initial_history, + Some(&session_source), + parent_thread_id, + /*forked_from_thread_id*/ None, + &config, + ) + .await; + let reservation = + if multi_agent_version == MultiAgentVersion::V2 { + None + } else { + Some(self.state.reserve_spawn_slot( + config.effective_agent_max_threads(multi_agent_version), + )?) + }; + let session_source = + resume_session_source(session_source, stored_session_source, multi_agent_version); + let agent_metadata = agent_metadata_from_session_source(thread_id, &session_source); + let notification_source = session_source.clone(); + let inherited_shell_snapshot = self + .inherited_shell_snapshot_for_source(&state, Some(&session_source)) + .await; + let inherited_exec_policy = self + .inherited_exec_policy_for_source(&state, Some(&session_source), &config) + .await; + + let resumed_thread = state + .resume_thread_with_history_with_source(ResumeThreadWithHistoryOptions { + config: config.clone(), + initial_history, + agent_control: self.clone(), + session_source, + parent_thread_id, + inherited_shell_snapshot, + inherited_exec_policy, + }) + .await?; + let mut agent_metadata = agent_metadata; + agent_metadata.agent_id = Some(resumed_thread.thread_id); + if let Some(reservation) = reservation { + reservation.commit(agent_metadata.clone()); + } else { + self.state.register_agent(agent_metadata.clone()); + } + state.notify_thread_created(resumed_thread.thread_id); + if multi_agent_version != MultiAgentVersion::V2 { + let child_reference = agent_metadata + .agent_path + .as_ref() + .map(ToString::to_string) + .unwrap_or_else(|| resumed_thread.thread_id.to_string()); + self.maybe_start_completion_watcher( + resumed_thread.thread_id, + Some(notification_source.clone()), + child_reference, + ); + } + 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) + }) + } + + #[allow(clippy::too_many_arguments)] + fn prepare_thread_spawn( + &self, + reservation: &mut crate::agent::registry::SpawnReservation, + config: &Config, + parent_thread_id: ThreadId, + depth: i32, + agent_path: Option, + agent_role: Option, + preferred_agent_nickname: Option, + multi_agent_version: MultiAgentVersion, + ) -> CodexResult<(SessionSource, AgentMetadata)> { + if depth == 1 { + self.state.register_root_thread(parent_thread_id); + } + if let Some(agent_path) = agent_path.as_ref() { + reservation.reserve_agent_path(agent_path)?; + } + let agent_nickname = if multi_agent_version == MultiAgentVersion::V2 { + None + } else { + let candidate_names = agent_nickname_candidates(config, agent_role.as_deref()); + let candidate_name_refs: Vec<&str> = + candidate_names.iter().map(String::as_str).collect(); + Some(reservation.reserve_agent_nickname_with_preference( + &candidate_name_refs, + preferred_agent_nickname.as_deref(), + )?) + }; + let session_source = SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id, + depth, + agent_path: agent_path.clone(), + agent_nickname: agent_nickname.clone(), + agent_role: agent_role.clone(), + }); + let agent_metadata = AgentMetadata { + agent_id: None, + agent_path, + agent_nickname, + agent_role, + last_task_message: None, + }; + Ok((session_source, agent_metadata)) + } + + async fn ensure_agent_path_available( + &self, + state: &Arc, + parent_thread_id: ThreadId, + depth: i32, + agent_path: &AgentPath, + ) -> CodexResult<()> { + let root_thread_id = if depth == 1 { + parent_thread_id + } else { + self.state + .agent_id_for_path(&AgentPath::root()) + .ok_or_else(|| { + CodexErr::UnsupportedOperation("root agent is unavailable".to_string()) + })? + }; + if let Some(state_db_ctx) = state.state_db() + && state_db_ctx + .find_open_thread_spawn_descendant_by_path(root_thread_id, agent_path.as_str()) + .await + .map_err(|err| { + CodexErr::Fatal(format!("failed to check agent path `{agent_path}`: {err}")) + })? + .is_some() + { + return Err(CodexErr::UnsupportedOperation(format!( + "agent path `{agent_path}` already exists" + ))); + } + Ok(()) + } + + async fn inherited_shell_snapshot_for_source( + &self, + state: &Arc, + session_source: Option<&SessionSource>, + ) -> Option> { + let Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id, .. + })) = session_source + else { + return None; + }; + + let parent_thread = state.get_thread(*parent_thread_id).await.ok()?; + parent_thread.codex.session.user_shell().shell_snapshot() + } + + async fn inherited_exec_policy_for_source( + &self, + state: &Arc, + session_source: Option<&SessionSource>, + child_config: &Config, + ) -> Option> { + let Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id, .. + })) = session_source + else { + return None; + }; + + let parent_thread = state.get_thread(*parent_thread_id).await.ok()?; + let parent_config = parent_thread.codex.session.get_config().await; + if !crate::exec_policy::child_uses_parent_exec_policy(&parent_config, child_config) { + return None; + } + + Some(Arc::clone( + &parent_thread.codex.session.services.exec_policy, + )) + } + + 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(SessionSource::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}"); + } + } +} diff --git a/codex-rs/core/src/agent/control_tests.rs b/codex-rs/core/src/agent/control_tests.rs index 3866d22013d0..014bc8120f59 100644 --- a/codex-rs/core/src/agent/control_tests.rs +++ b/codex-rs/core/src/agent/control_tests.rs @@ -132,6 +132,46 @@ impl AgentControlHarness { } } +async fn spawn_v2_test_agent( + harness: &AgentControlHarness, + config: &Config, + parent_thread_id: ThreadId, + agent_path: AgentPath, +) -> ThreadId { + harness + .control + .spawn_agent( + config.clone(), + Op::Interrupt, + Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id, + depth: 1, + agent_path: Some(agent_path), + agent_nickname: None, + agent_role: None, + })), + ) + .await + .expect("test agent should spawn") +} + +async fn terminate_test_agent( + harness: &AgentControlHarness, + agent_id: ThreadId, +) -> Arc { + let agent_thread = harness + .manager + .get_thread(agent_id) + .await + .expect("test agent should exist"); + agent_thread + .submit(Op::Shutdown {}) + .await + .expect("test agent shutdown should submit"); + agent_thread.wait_until_terminated().await; + agent_thread +} + fn has_subagent_notification(history_items: &[ResponseItem]) -> bool { history_items.iter().any(|item| { let ResponseItem::Message { role, content, .. } = item else { @@ -1750,6 +1790,261 @@ async fn multi_agent_v2_completion_ignores_dead_direct_parent() { assert!(!has_subagent_notification(&root_history_items)); } +#[tokio::test] +async fn multi_agent_v2_direct_trigger_turn_respects_execution_limit() { + let harness = AgentControlHarness::new().await; + let mut config = harness.config.clone(); + let _ = config.features.enable(Feature::MultiAgentV2); + config.multi_agent_v2.max_concurrent_threads_per_session = 2; + let root = harness + .manager + .start_thread(config.clone()) + .await + .expect("root thread should start"); + let target_path = AgentPath::root().join("target").expect("target path"); + let target_thread_id = + spawn_v2_test_agent(&harness, &config, root.thread_id, target_path.clone()).await; + harness + .control + .state + .release_execution_slot(target_thread_id); + + let blocker_path = AgentPath::root().join("blocker").expect("blocker path"); + let blocker_thread_id = + spawn_v2_test_agent(&harness, &config, root.thread_id, blocker_path).await; + let target_thread = harness + .manager + .get_thread(target_thread_id) + .await + .expect("target thread should exist"); + target_thread + .submit(Op::InterAgentCommunication { + communication: InterAgentCommunication::new( + AgentPath::root(), + target_path, + Vec::new(), + "run when capacity is available".to_string(), + /*trigger_turn*/ true, + ), + }) + .await + .expect("direct communication should submit"); + + timeout(Duration::from_secs(5), async { + while !target_thread + .codex + .session + .input_queue + .has_trigger_turn_mailbox_items() + .await + { + sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("target communication should reach the mailbox"); + assert!(!harness.control.state.is_execution_active(target_thread_id)); + assert!( + target_thread + .codex + .session + .active_turn + .lock() + .await + .is_none() + ); + + harness + .control + .state + .release_execution_slot(blocker_thread_id); + target_thread + .codex + .session + .maybe_start_turn_for_pending_work() + .await; + assert!(harness.control.state.is_execution_active(target_thread_id)); + + harness + .control + .shutdown_and_forget_agent(target_thread_id) + .await + .expect("target shutdown should succeed"); + harness + .control + .shutdown_and_forget_agent(blocker_thread_id) + .await + .expect("blocker shutdown should succeed"); +} + +#[tokio::test] +async fn multi_agent_v2_resident_cache_evicts_and_reloads_dead_lru_agent() { + let harness = AgentControlHarness::new().await; + let mut config = harness.config.clone(); + let _ = config.features.enable(Feature::MultiAgentV2); + let root = harness + .manager + .start_thread(config.clone()) + .await + .expect("root thread should start"); + let protected_path = AgentPath::root().join("protected").expect("protected path"); + let protected_thread_id = + spawn_v2_test_agent(&harness, &config, root.thread_id, protected_path.clone()).await; + harness + .control + .state + .release_execution_slot(protected_thread_id); + let protected_thread = harness + .manager + .get_thread(protected_thread_id) + .await + .expect("protected thread should exist"); + protected_thread + .codex + .session + .input_queue + .enqueue_mailbox_communication(InterAgentCommunication::new( + AgentPath::root(), + protected_path, + Vec::new(), + "keep this resident".to_string(), + /*trigger_turn*/ false, + )) + .await; + let target_path = AgentPath::root().join("target").expect("target path"); + let target_thread_id = + spawn_v2_test_agent(&harness, &config, root.thread_id, target_path.clone()).await; + harness + .control + .state + .release_execution_slot(target_thread_id); + let target_thread = terminate_test_agent(&harness, target_thread_id).await; + // Simulate a stale nonterminal status after the session channel has died. + let synthetic_turn = target_thread.codex.session.new_default_turn().await; + target_thread + .codex + .session + .send_event( + synthetic_turn.as_ref(), + EventMsg::TurnStarted(TurnStartedEvent { + turn_id: synthetic_turn.sub_id.clone(), + trace_id: None, + started_at: None, + model_context_window: None, + collaboration_mode_kind: ModeKind::Default, + }), + ) + .await; + + let resident_limit = crate::agent::residency::MIN_RESIDENT_SUBAGENTS.max( + config + .multi_agent_v2 + .max_concurrent_threads_per_session + .saturating_add(4), + ); + { + let mut residency = harness.control.residency.lock().await; + for _ in 2..resident_limit { + residency.touch(ThreadId::new()); + } + assert_eq!(residency.len(), resident_limit); + } + + let newcomer_path = AgentPath::root().join("newcomer").expect("newcomer path"); + let newcomer_thread_id = + spawn_v2_test_agent(&harness, &config, root.thread_id, newcomer_path).await; + assert!(harness.manager.get_thread(target_thread_id).await.is_err()); + assert!( + harness + .manager + .get_thread(protected_thread_id) + .await + .is_ok() + ); + assert!( + harness + .control + .residency + .lock() + .await + .contains(protected_thread_id) + ); + assert!( + !harness + .control + .residency + .lock() + .await + .contains(target_thread_id) + ); + + harness + .control + .send_followup_to_agent( + &config, + target_thread_id, + InterAgentCommunication::new( + AgentPath::root(), + target_path, + Vec::new(), + "continue after eviction".to_string(), + /*trigger_turn*/ true, + ), + ) + .await + .expect("follow-up should reload the evicted agent"); + assert!(harness.manager.get_thread(target_thread_id).await.is_ok()); + assert!( + harness + .control + .residency + .lock() + .await + .contains(target_thread_id) + ); + + harness + .control + .shutdown_and_forget_agent(target_thread_id) + .await + .expect("target shutdown should succeed"); + harness + .control + .shutdown_and_forget_agent(newcomer_thread_id) + .await + .expect("newcomer shutdown should succeed"); + harness + .control + .shutdown_and_forget_agent(protected_thread_id) + .await + .expect("protected shutdown should succeed"); +} + +#[tokio::test] +async fn interrupt_agent_cleans_up_dead_resident_thread() { + let harness = AgentControlHarness::new().await; + let mut config = harness.config.clone(); + let _ = config.features.enable(Feature::MultiAgentV2); + let root = harness + .manager + .start_thread(config.clone()) + .await + .expect("root thread should start"); + let agent_path = AgentPath::root().join("worker").expect("worker path"); + let agent_id = spawn_v2_test_agent(&harness, &config, root.thread_id, agent_path).await; + terminate_test_agent(&harness, agent_id).await; + + let err = harness + .control + .interrupt_agent(agent_id) + .await + .expect_err("interrupting a dead worker should report the dead channel"); + assert_matches!(err, CodexErr::InternalAgentDied); + assert!(harness.manager.get_thread(agent_id).await.is_err()); + assert!(!harness.control.residency.lock().await.contains(agent_id)); + assert!(!harness.control.state.is_execution_active(agent_id)); +} + #[tokio::test] async fn multi_agent_v2_completion_queues_message_for_direct_parent() { let harness = AgentControlHarness::new().await; diff --git a/codex-rs/core/src/agent/mod.rs b/codex-rs/core/src/agent/mod.rs index 350962dc0811..a3016486d290 100644 --- a/codex-rs/core/src/agent/mod.rs +++ b/codex-rs/core/src/agent/mod.rs @@ -1,6 +1,7 @@ pub(crate) mod agent_resolver; pub(crate) mod control; mod registry; +mod residency; pub(crate) mod role; pub(crate) mod status; diff --git a/codex-rs/core/src/agent/registry.rs b/codex-rs/core/src/agent/registry.rs index 43aca201bfa2..28d34aad96a4 100644 --- a/codex-rs/core/src/agent/registry.rs +++ b/codex-rs/core/src/agent/registry.rs @@ -13,12 +13,7 @@ use std::sync::Mutex; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; -/// This structure is used to add some limits on the multi-agent capabilities for Codex. In -/// the current implementation, it limits: -/// * Total number of sub-agents (i.e. threads) per user session -/// -/// This structure is shared by all agents in the same user session (because the `AgentControl` -/// is). +/// Tracks logical agents and active execution slots for one multi-agent session tree. #[derive(Default)] pub(crate) struct AgentRegistry { active_agents: Mutex, @@ -28,6 +23,7 @@ pub(crate) struct AgentRegistry { #[derive(Default)] struct ActiveAgents { agent_tree: HashMap, + active_thread_ids: HashSet, used_agent_nicknames: HashSet, nickname_reset_count: usize, } @@ -97,25 +93,68 @@ impl AgentRegistry { } pub(crate) fn release_spawned_thread(&self, thread_id: ThreadId) { - let removed_counted_agent = { - let mut active_agents = self - .active_agents - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); - let removed_key = active_agents - .agent_tree - .iter() - .find_map(|(key, metadata)| (metadata.agent_id == Some(thread_id)).then_some(key)) - .cloned(); - removed_key - .and_then(|key| active_agents.agent_tree.remove(key.as_str())) - .is_some_and(|metadata| { - !metadata.agent_path.as_ref().is_some_and(AgentPath::is_root) - }) - }; - if removed_counted_agent { + let mut active_agents = self + .active_agents + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let removed_key = active_agents + .agent_tree + .iter() + .find_map(|(key, metadata)| (metadata.agent_id == Some(thread_id)).then_some(key)) + .cloned(); + if let Some(removed_key) = removed_key { + active_agents.agent_tree.remove(removed_key.as_str()); + } + if active_agents.active_thread_ids.remove(&thread_id) { + self.total_count.fetch_sub(1, Ordering::AcqRel); + } + } + + pub(crate) fn release_execution_slot(&self, thread_id: ThreadId) -> bool { + let mut active_agents = self + .active_agents + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let removed = active_agents.active_thread_ids.remove(&thread_id); + if removed { self.total_count.fetch_sub(1, Ordering::AcqRel); } + removed + } + + pub(crate) fn reserve_execution_slot( + self: &Arc, + thread_id: ThreadId, + max_threads: usize, + ) -> Result { + let mut active_agents = self + .active_agents + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if active_agents.active_thread_ids.contains(&thread_id) { + return Ok(ExecutionReservation { + state: Arc::clone(self), + thread_id, + release_on_drop: false, + }); + } + if !self.try_increment_spawned(max_threads) { + return Err(CodexErr::AgentLimitReached { max_threads }); + } + active_agents.active_thread_ids.insert(thread_id); + Ok(ExecutionReservation { + state: Arc::clone(self), + thread_id, + release_on_drop: true, + }) + } + + pub(crate) fn is_execution_active(&self, thread_id: ThreadId) -> bool { + self.active_agents + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .active_thread_ids + .contains(&thread_id) } pub(crate) fn register_root_thread(&self, thread_id: ThreadId) { @@ -152,20 +191,6 @@ impl AgentRegistry { .cloned() } - pub(crate) fn live_agents(&self) -> Vec { - self.active_agents - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .agent_tree - .values() - .filter(|metadata| { - metadata.agent_id.is_some() - && !metadata.agent_path.as_ref().is_some_and(AgentPath::is_root) - }) - .cloned() - .collect() - } - pub(crate) fn update_last_task_message(&self, thread_id: ThreadId, last_task_message: String) { let mut active_agents = self .active_agents @@ -194,7 +219,7 @@ impl AgentRegistry { } } - fn register_spawned_thread(&self, agent_metadata: AgentMetadata) { + pub(crate) fn register_agent(&self, mut agent_metadata: AgentMetadata) { let Some(thread_id) = agent_metadata.agent_id else { return; }; @@ -210,9 +235,25 @@ impl AgentRegistry { if let Some(agent_nickname) = agent_metadata.agent_nickname.clone() { active_agents.used_agent_nicknames.insert(agent_nickname); } + if let Some(existing) = active_agents.agent_tree.get(key.as_str()) + && agent_metadata.last_task_message.is_none() + { + agent_metadata.last_task_message = existing.last_task_message.clone(); + } active_agents.agent_tree.insert(key, agent_metadata); } + fn commit_spawned_thread(&self, agent_metadata: AgentMetadata) { + if let Some(thread_id) = agent_metadata.agent_id { + self.active_agents + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .active_thread_ids + .insert(thread_id); + } + self.register_agent(agent_metadata); + } + fn reserve_agent_nickname(&self, names: &[&str], preferred: Option<&str>) -> Option { let mut active_agents = self .active_agents @@ -337,7 +378,7 @@ impl SpawnReservation { pub(crate) fn commit(mut self, agent_metadata: AgentMetadata) { self.reserved_agent_nickname = None; self.reserved_agent_path = None; - self.state.register_spawned_thread(agent_metadata); + self.state.commit_spawned_thread(agent_metadata); self.active = false; } } @@ -353,6 +394,26 @@ impl Drop for SpawnReservation { } } +pub(crate) struct ExecutionReservation { + state: Arc, + thread_id: ThreadId, + release_on_drop: bool, +} + +impl ExecutionReservation { + pub(crate) fn commit(mut self) { + self.release_on_drop = false; + } +} + +impl Drop for ExecutionReservation { + fn drop(&mut self) { + if self.release_on_drop { + self.state.release_execution_slot(self.thread_id); + } + } +} + #[cfg(test)] #[path = "registry_tests.rs"] mod tests; diff --git a/codex-rs/core/src/agent/registry_tests.rs b/codex-rs/core/src/agent/registry_tests.rs index fc172fb336df..44349de5c184 100644 --- a/codex-rs/core/src/agent/registry_tests.rs +++ b/codex-rs/core/src/agent/registry_tests.rs @@ -81,7 +81,7 @@ fn reservation_drop_releases_slot() { } #[test] -fn commit_holds_slot_until_release() { +fn completion_releases_slot_without_removing_logical_agent() { let registry = Arc::new(AgentRegistry::default()); let reservation = registry.reserve_spawn_slot(Some(1)).expect("reserve slot"); let thread_id = ThreadId::new(); @@ -96,10 +96,21 @@ fn commit_holds_slot_until_release() { }; assert_eq!(max_threads, 1); - registry.release_spawned_thread(thread_id); + let followup = registry + .reserve_execution_slot(thread_id, 1) + .expect("running agent should retain its slot"); + followup.commit(); + assert!(registry.reserve_spawn_slot(Some(1)).is_err()); + registry.release_execution_slot(thread_id); + assert_eq!( + registry + .agent_metadata_for_thread(thread_id) + .and_then(|metadata| metadata.agent_id), + Some(thread_id) + ); let reservation = registry .reserve_spawn_slot(Some(1)) - .expect("slot released after thread removal"); + .expect("slot released after agent completion"); drop(reservation); } diff --git a/codex-rs/core/src/agent/residency.rs b/codex-rs/core/src/agent/residency.rs new file mode 100644 index 000000000000..7d398c9d2f7a --- /dev/null +++ b/codex-rs/core/src/agent/residency.rs @@ -0,0 +1,32 @@ +use codex_protocol::ThreadId; +use std::collections::VecDeque; + +pub(crate) const MIN_RESIDENT_SUBAGENTS: usize = 64; + +#[derive(Default)] +pub(crate) struct ResidentAgents { + lru: VecDeque, +} + +impl ResidentAgents { + pub(crate) fn len(&self) -> usize { + self.lru.len() + } + + pub(crate) fn contains(&self, thread_id: ThreadId) -> bool { + self.lru.contains(&thread_id) + } + + pub(crate) fn touch(&mut self, thread_id: ThreadId) { + self.remove(thread_id); + self.lru.push_back(thread_id); + } + + pub(crate) fn remove(&mut self, thread_id: ThreadId) { + self.lru.retain(|resident_id| *resident_id != thread_id); + } + + pub(crate) fn snapshot(&self) -> Vec { + self.lru.iter().copied().collect() + } +} diff --git a/codex-rs/core/src/session/handlers.rs b/codex-rs/core/src/session/handlers.rs index 7730a30ad77e..6a032f8166f0 100644 --- a/codex-rs/core/src/session/handlers.rs +++ b/codex-rs/core/src/session/handlers.rs @@ -22,6 +22,7 @@ use crate::realtime_conversation::prefix_realtime_v2_text; use crate::review_prompts::resolve_review_request; use crate::session::spawn_review_thread; use crate::tasks::CompactTask; +use crate::tasks::StartTaskOutcome; use crate::tasks::UserShellCommandMode; use crate::tasks::UserShellCommandTask; use crate::tasks::execute_user_shell_command; @@ -273,13 +274,24 @@ pub(super) async fn user_input_or_turn_inner( client_id: client_user_message_id, }); } - sess.spawn_task( - Arc::clone(¤t_context), - task_input, - crate::tasks::RegularTask::new(), - ) - .await; - Some(accepted_items) + match sess + .spawn_task( + Arc::clone(¤t_context), + task_input, + crate::tasks::RegularTask::new(), + ) + .await + { + StartTaskOutcome::Started => Some(accepted_items), + StartTaskOutcome::Rejected(err) => { + sess.send_event_raw(Event { + id: sub_id, + msg: EventMsg::Error(err.to_error_event(/*message_prefix*/ None)), + }) + .await; + None + } + } } Err(err) => { sess.send_event_raw(Event { diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 1b40bd765697..611a52627c2c 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -1760,7 +1760,11 @@ impl Session { if let Err(err) = self .services .agent_control - .send_inter_agent_communication(parent_thread_id, communication) + .send_message_to_agent( + turn_context.config.as_ref(), + parent_thread_id, + communication, + ) .await { debug!("failed to notify parent thread {parent_thread_id}: {err}"); diff --git a/codex-rs/core/src/tasks/lifecycle.rs b/codex-rs/core/src/tasks/lifecycle.rs index 782cc8c7deda..59c46e617da6 100644 --- a/codex-rs/core/src/tasks/lifecycle.rs +++ b/codex-rs/core/src/tasks/lifecycle.rs @@ -1,5 +1,6 @@ use codex_extension_api::ExtensionData; use codex_protocol::protocol::CodexErrorInfo; +use codex_protocol::protocol::MultiAgentVersion; use codex_protocol::protocol::TokenUsage; use codex_protocol::protocol::TurnAbortReason; @@ -53,6 +54,13 @@ impl Session { }) .await; } + + if self.multi_agent_version() == Some(MultiAgentVersion::V2) { + self.services + .agent_control + .release_execution_slot_if_idle(self) + .await; + } } pub(super) async fn emit_turn_abort_lifecycle( diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index ef2ba0f16d1a..ac0687cd1acb 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -33,6 +33,7 @@ use crate::session::turn_context::TurnContext; use crate::state::ActiveTurn; use crate::state::RunningTask; use crate::state::TaskKind; +use crate::state::TurnState; use codex_analytics::TurnTokenUsageFact; use codex_login::AuthManager; use codex_models_manager::manager::SharedModelsManager; @@ -42,6 +43,7 @@ use codex_otel::TURN_MEMORY_METRIC; use codex_otel::TURN_NETWORK_PROXY_METRIC; use codex_otel::TURN_TOKEN_USAGE_METRIC; use codex_otel::TURN_TOOL_CALL_METRIC; +use codex_protocol::error::CodexErr; use codex_protocol::models::ResponseItem; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::MultiAgentVersion; @@ -70,6 +72,12 @@ pub(crate) enum InterruptedTurnHistoryMarker { Developer, } +#[derive(Debug)] +pub(crate) enum StartTaskOutcome { + Started, + Rejected(CodexErr), +} + impl InterruptedTurnHistoryMarker { pub(crate) fn from_config_and_version( config: &Config, @@ -307,7 +315,7 @@ impl Session { turn_context: Arc, input: Vec, task: T, - ) { + ) -> StartTaskOutcome { self.abort_all_tasks(TurnAbortReason::Replaced).await; self.clear_connector_selection().await; self.start_task(turn_context, input, task).await; @@ -318,7 +326,37 @@ impl Session { turn_context: Arc, input: Vec, task: T, - ) { + ) -> StartTaskOutcome { + let turn_state = { + let mut active = self.active_turn.lock().await; + let turn = active.get_or_insert_with(ActiveTurn::default); + if turn.task.is_some() { + return StartTaskOutcome::Rejected(CodexErr::InvalidRequest( + "thread already has an active turn".to_string(), + )); + } + Arc::clone(&turn.turn_state) + }; + let execution_reservation = match self + .services + .agent_control + .reserve_execution_slot_for_turn_start(self) + .await + { + Ok(reservation) => reservation, + Err(err) => { + self.clear_reserved_active_turn(&turn_state).await; + warn!( + thread_id = %self.thread_id, + "agent turn could not reserve an execution slot: {err}" + ); + return StartTaskOutcome::Rejected(err); + } + }; + if let Some(execution_reservation) = execution_reservation { + execution_reservation.commit(); + } + let task: Arc = Arc::new(task); let task_kind = task.kind(); let span_name = task.span_name(); @@ -351,12 +389,6 @@ impl Session { warn!("failed to apply goal runtime turn-start event: {err}"); } let pending_items = self.input_queue.get_pending_input(&self.active_turn).await; - let turn_state = { - let mut active = self.active_turn.lock().await; - let turn = active.get_or_insert_with(ActiveTurn::default); - debug_assert!(turn.task.is_none()); - Arc::clone(&turn.turn_state) - }; turn_state.lock().await.token_usage_at_turn_start = token_usage_at_turn_start.clone(); self.input_queue .extend_pending_input_for_turn_state(turn_state.as_ref(), pending_items) @@ -366,7 +398,12 @@ impl Session { let turn_extension_data = Arc::clone(&turn_context.extension_data); let mut active = self.active_turn.lock().await; - let turn = active.get_or_insert_with(ActiveTurn::default); + let Some(turn) = active.as_mut() else { + return StartTaskOutcome::Rejected(CodexErr::Fatal( + "active turn reservation was lost before task start".to_string(), + )); + }; + debug_assert!(Arc::ptr_eq(&turn.turn_state, &turn_state)); debug_assert!(turn.task.is_none()); let done_clone = Arc::clone(&done); let session_ctx = Arc::new(SessionTaskContext::new( @@ -442,6 +479,7 @@ impl Session { _timer: timer, }; turn.task = Some(running_task); + StartTaskOutcome::Started } /// Starts a regular turn when the session is idle and pending work is waiting. @@ -460,27 +498,34 @@ impl Session { /// /// The turn is created only when there is mailbox mail marked with `trigger_turn`, and only /// if the session is currently idle. - pub(crate) async fn maybe_start_turn_for_pending_work_with_sub_id( + pub(crate) fn maybe_start_turn_for_pending_work_with_sub_id( self: &Arc, sub_id: String, - ) { - if !self.input_queue.has_trigger_turn_mailbox_items().await { - return; - } - - { - let mut active_turn = self.active_turn.lock().await; - if active_turn.is_some() { + ) -> BoxFuture<'_, ()> { + Box::pin(async move { + if !self.input_queue.has_trigger_turn_mailbox_items().await { return; } - *active_turn = Some(ActiveTurn::default()); - } - let turn_context = self.new_default_turn_with_sub_id(sub_id).await; - self.maybe_emit_unknown_model_warning_for_turn(turn_context.as_ref()) - .await; - self.start_task(turn_context, Vec::new(), RegularTask::new()) - .await; + let turn_state = { + let mut active_turn = self.active_turn.lock().await; + if active_turn.is_some() { + return; + } + let active_turn = active_turn.get_or_insert_with(ActiveTurn::default); + Arc::clone(&active_turn.turn_state) + }; + + let turn_context = self.new_default_turn_with_sub_id(sub_id).await; + self.maybe_emit_unknown_model_warning_for_turn(turn_context.as_ref()) + .await; + if !self.input_queue.has_trigger_turn_mailbox_items().await { + self.clear_reserved_active_turn(&turn_state).await; + return; + } + self.start_task(turn_context, Vec::new(), RegularTask::new()) + .await; + }) } pub async fn abort_all_tasks(self: &Arc, reason: TurnAbortReason) { @@ -519,6 +564,7 @@ impl Session { } if reason == TurnAbortReason::Interrupted && aborted_turn { self.maybe_start_turn_for_pending_work().await; + self.emit_thread_idle_lifecycle_if_idle().await; } } @@ -566,6 +612,7 @@ impl Session { if reason == TurnAbortReason::Interrupted { self.maybe_start_turn_for_pending_work().await; + self.emit_thread_idle_lifecycle_if_idle().await; } true @@ -797,6 +844,7 @@ impl Session { { warn!("failed to apply goal runtime maybe-continue event: {err}"); } + self.maybe_start_turn_for_pending_work().await; self.emit_thread_idle_lifecycle_if_idle().await; } @@ -805,6 +853,16 @@ impl Session { active.take() } + async fn clear_reserved_active_turn(&self, turn_state: &Arc>) { + let mut active_turn_guard = self.active_turn.lock().await; + if let Some(active_turn) = active_turn_guard.as_ref() + && active_turn.task.is_none() + && Arc::ptr_eq(&active_turn.turn_state, turn_state) + { + *active_turn_guard = None; + } + } + pub(crate) async fn close_unified_exec_processes(&self) { self.services .unified_exec_manager diff --git a/codex-rs/core/src/thread_manager.rs b/codex-rs/core/src/thread_manager.rs index d130a358158d..e0312fa7767f 100644 --- a/codex-rs/core/src/thread_manager.rs +++ b/codex-rs/core/src/thread_manager.rs @@ -1,5 +1,6 @@ use crate::SkillsManager; use crate::agent::AgentControl; +use crate::agent::AgentStatus; use crate::attestation::AttestationProvider; use crate::codex_thread::CodexThread; use crate::config::Config; @@ -1032,6 +1033,35 @@ impl ThreadManagerState { thread.submit(op).await } + pub(crate) async fn enqueue_inter_agent_communication( + &self, + thread_id: ThreadId, + communication: codex_protocol::protocol::InterAgentCommunication, + ) -> CodexResult { + let thread = self.get_thread(thread_id).await?; + if matches!(thread.agent_status().await, AgentStatus::Shutdown) { + return Err(CodexErr::InternalAgentDied); + } + if let Some(ops_log) = &self.ops_log + && let Ok(mut log) = ops_log.lock() + { + log.push(( + thread_id, + Op::InterAgentCommunication { + communication: communication.clone(), + }, + )); + } + let sub_id = uuid::Uuid::now_v7().to_string(); + thread + .codex + .session + .input_queue + .enqueue_mailbox_communication(communication) + .await; + Ok(sub_id) + } + /// Remove a thread from the manager by ID, returning it when present. pub(crate) async fn remove_thread(&self, thread_id: &ThreadId) -> Option> { self.threads.write().await.remove(thread_id) diff --git a/codex-rs/core/src/tools/handlers/agent_jobs.rs b/codex-rs/core/src/tools/handlers/agent_jobs.rs index 360744973c56..5608f352dde0 100644 --- a/codex-rs/core/src/tools/handlers/agent_jobs.rs +++ b/codex-rs/core/src/tools/handlers/agent_jobs.rs @@ -248,7 +248,7 @@ async fn run_agent_job_loop( let _ = session .services .agent_control - .shutdown_live_agent(thread_id) + .shutdown_and_forget_agent(thread_id) .await; continue; } @@ -369,7 +369,7 @@ async fn recover_running_items( let _ = session .services .agent_control - .shutdown_live_agent(thread_id) + .shutdown_and_forget_agent(thread_id) .await; } continue; @@ -491,7 +491,7 @@ async fn reap_stale_active_items( let _ = session .services .agent_control - .shutdown_live_agent(thread_id) + .shutdown_and_forget_agent(thread_id) .await; active_items.remove(&thread_id); } @@ -527,7 +527,7 @@ async fn finalize_finished_item( let _ = session .services .agent_control - .shutdown_live_agent(thread_id) + .shutdown_and_forget_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 0d4e17447dc8..e47d217d8c70 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 @@ -89,7 +89,7 @@ async fn handle_close_agent( return Err(collab_agent_error(agent_id, err)); } }; - let result = Box::pin(session.services.agent_control.close_agent(agent_id)) + let result = Box::pin(session.services.agent_control.close_agent_v1(agent_id)) .await .map_err(|err| collab_agent_error(agent_id, err)) .map(|_| ()); diff --git a/codex-rs/core/src/tools/handlers/multi_agents_spec.rs b/codex-rs/core/src/tools/handlers/multi_agents_spec.rs index 1dcb55f436fd..a1f25996545c 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_spec.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_spec.rs @@ -111,9 +111,7 @@ pub fn create_spawn_agent_tool_v2(options: SpawnAgentToolOptions) -> ToolSpec { Some(vec!["task_name".to_string(), "message".to_string()]), Some(false.into()), ), - output_schema: Some(spawn_agent_output_schema_v2( - options.hide_agent_type_model_reasoning, - )), + output_schema: Some(spawn_agent_output_schema_v2()), }) } @@ -207,7 +205,7 @@ pub fn create_followup_task_tool() -> ToolSpec { ToolSpec::Function(ResponsesApiTool { name: "followup_task".to_string(), - description: "Send a follow-up task to an existing non-root target agent and trigger a turn if it is idle. If the target is already running, deliver the task promptly at message boundaries while sampling, or after the pending tool call completes." + description: "Send a follow-up task to an existing non-root target agent and trigger a turn if it is idle. If the target is already running, deliver the task promptly at message boundaries while sampling, or after the pending tool call completes. Starting an idle agent can fail when the session concurrency limit is full." .to_string(), strict: false, defer_loading: None, @@ -270,16 +268,15 @@ pub fn create_list_agents_tool() -> ToolSpec { let properties = BTreeMap::from([( "path_prefix".to_string(), JsonSchema::string(Some( - "Task-path prefix filter without a trailing slash. Omit to list all live agents." + "Task-path prefix filter without a trailing slash. Omit to list all resident agents." .to_string(), )), )]); ToolSpec::Function(ResponsesApiTool { name: "list_agents".to_string(), - description: - "List live agents in the current root thread tree. Optionally filter by task-path prefix." - .to_string(), + description: "List currently resident agents in the root thread tree. Optionally filter by task-path prefix." + .to_string(), strict: false, defer_loading: None, parameters: JsonSchema::object(properties, /*required*/ None, Some(false.into())), @@ -317,7 +314,7 @@ pub fn create_close_agent_tool_v2() -> ToolSpec { ToolSpec::Function(ResponsesApiTool { name: "close_agent".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. Completed agents remain open and count toward the concurrency limit until closed. 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. Idle agents do not count toward the concurrency limit.".to_string(), strict: false, defer_loading: None, parameters: JsonSchema::object(properties, Some(vec!["target".to_string()]), Some(false.into())), @@ -374,34 +371,16 @@ fn spawn_agent_output_schema_v1() -> Value { }) } -fn spawn_agent_output_schema_v2(hide_agent_metadata: bool) -> Value { - if hide_agent_metadata { - return json!({ - "type": "object", - "properties": { - "task_name": { - "type": "string", - "description": "Canonical task name for the spawned agent." - } - }, - "required": ["task_name"], - "additionalProperties": false - }); - } - +fn spawn_agent_output_schema_v2() -> Value { json!({ "type": "object", "properties": { "task_name": { "type": "string", "description": "Canonical task name for the spawned agent." - }, - "nickname": { - "type": ["string", "null"], - "description": "User-facing nickname for the spawned agent when available." } }, - "required": ["task_name", "nickname"], + "required": ["task_name"], "additionalProperties": false }) } @@ -445,7 +424,7 @@ fn list_agents_output_schema() -> Value { "required": ["agent_name", "agent_status", "last_task_message"], "additionalProperties": false }, - "description": "Live agents visible in the current root thread tree." + "description": "Currently resident agents visible in the root thread tree." } }, "required": ["agents"], @@ -724,8 +703,9 @@ fn spawn_agent_tool_description_v2( let inherited_model_guidance = inherited_model_guidance.unwrap_or_default(); let concurrency_guidance = max_concurrent_threads_per_session .map(|limit| { + let active_spawned_agent_limit = limit.saturating_sub(1); format!( - "This session is configured with `max_concurrent_threads_per_session = {limit}` for concurrently open agent threads." + "This session is configured with `max_concurrent_threads_per_session = {limit}`, including the root thread, so at most {active_spawned_agent_limit} spawned agent threads can run concurrently. Idle spawned agents do not count toward this limit." ) }) .unwrap_or_default(); diff --git a/codex-rs/core/src/tools/handlers/multi_agents_spec_tests.rs b/codex-rs/core/src/tools/handlers/multi_agents_spec_tests.rs index 512e318a8225..88c08f1e6d72 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_spec_tests.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_spec_tests.rs @@ -112,7 +112,7 @@ fn spawn_agent_tool_v2_requires_task_name_and_lists_visible_models() { ); assert_eq!( output_schema.expect("spawn_agent output schema")["required"], - json!(["task_name", "nickname"]) + json!(["task_name"]) ); } @@ -307,7 +307,7 @@ fn followup_task_tool_requires_message_and_has_no_output_schema() { assert_eq!(name, "followup_task"); assert_eq!( description, - "Send a follow-up task to an existing non-root target agent and trigger a turn if it is idle. If the target is already running, deliver the task promptly at message boundaries while sampling, or after the pending tool call completes." + "Send a follow-up task to an existing non-root target agent and trigger a turn if it is idle. If the target is already running, deliver the task promptly at message boundaries while sampling, or after the pending tool call completes. Starting an idle agent can fail when the session concurrency limit is full." ); assert_eq!( parameters.schema_type, @@ -397,7 +397,7 @@ fn list_agents_tool_includes_path_prefix_and_agent_fields() { properties .get("path_prefix") .and_then(|schema| schema.description.as_deref()), - Some("Task-path prefix filter without a trailing slash. Omit to list all live agents.") + Some("Task-path prefix filter without a trailing slash. Omit to list all resident agents.") ); assert_eq!( output_schema.expect("list_agents output schema")["properties"]["agents"]["items"]["required"], 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 67ccf36f916c..10d4957c90cd 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_tests.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_tests.rs @@ -1680,70 +1680,6 @@ async fn multi_agent_v2_list_agents_filters_by_relative_path_prefix() { assert_eq!(result.agents[0].last_task_message.as_deref(), Some("build")); } -#[tokio::test] -async fn multi_agent_v2_list_agents_omits_closed_agents() { - let (mut session, mut turn) = make_session_and_context().await; - let manager = thread_manager(); - let root = manager - .start_thread((*turn.config).clone()) - .await - .expect("root thread should start"); - session.services.agent_control = manager.agent_control(); - session.thread_id = root.thread_id; - let mut config = (*turn.config).clone(); - let _ = config.features.enable(Feature::MultiAgentV2); - set_turn_config(&mut turn, config); - - let session = Arc::new(session); - let turn = Arc::new(turn); - let spawn_output = SpawnAgentHandlerV2::default() - .handle(invocation( - session.clone(), - turn.clone(), - "spawn_agent", - function_payload(json!({ - "message": "inspect this repo", - "task_name": "worker" - })), - )) - .await - .expect("spawn_agent should succeed"); - let _ = expect_text_output(spawn_output); - - let agent_id = session - .services - .agent_control - .resolve_agent_reference(session.thread_id, &turn.session_source, "worker") - .await - .expect("worker path should resolve"); - session - .services - .agent_control - .close_agent(agent_id) - .await - .expect("close_agent should succeed"); - - let output = ListAgentsHandlerV2 - .handle(invocation( - session, - turn, - "list_agents", - function_payload(json!({})), - )) - .await - .expect("list_agents should succeed"); - let (content, _) = expect_text_output(output); - let result: ListAgentsResult = - serde_json::from_str(&content).expect("list_agents result should be json"); - - assert_eq!(result.agents.len(), 1); - assert_eq!(result.agents[0].agent_name, "/root"); - assert_eq!( - result.agents[0].last_task_message.as_deref(), - Some("Main thread") - ); -} - #[tokio::test] async fn multi_agent_v2_send_message_rejects_legacy_items_field() { let (mut session, mut turn) = make_session_and_context().await; @@ -2722,7 +2658,7 @@ async fn resume_agent_restores_closed_agent_and_accepts_send_input() { let agent_id = thread.thread_id; let _ = manager .agent_control() - .shutdown_live_agent(agent_id) + .shutdown_and_forget_agent(agent_id) .await .expect("shutdown agent"); assert_eq!( @@ -2770,7 +2706,7 @@ async fn resume_agent_restores_closed_agent_and_accepts_send_input() { let _ = manager .agent_control() - .shutdown_live_agent(agent_id) + .shutdown_and_forget_agent(agent_id) .await .expect("shutdown resumed agent"); } diff --git a/codex-rs/core/src/tools/handlers/multi_agents_v2/message_tool.rs b/codex-rs/core/src/tools/handlers/multi_agents_v2/message_tool.rs index f23d700e9e49..d2c61a841084 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_v2/message_tool.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_v2/message_tool.rs @@ -74,8 +74,9 @@ pub(crate) async fn handle_message_string_tool( let receiver_agent = session .services .agent_control - .get_agent_metadata(receiver_thread_id) - .unwrap_or_default(); + .ensure_agent_known(receiver_thread_id) + .await + .map_err(|err| collab_agent_error(receiver_thread_id, err))?; if mode == MessageDeliveryMode::TriggerTurn && receiver_agent .agent_path @@ -107,12 +108,24 @@ pub(crate) async fn handle_message_string_tool( .get_agent_path() .unwrap_or_else(AgentPath::root); let communication = communication_from_tool_message(author, receiver_agent_path, message); - let result = session - .services - .agent_control - .send_inter_agent_communication(receiver_thread_id, mode.apply(communication)) - .await - .map_err(|err| collab_agent_error(receiver_thread_id, err)); + let communication = mode.apply(communication); + let result = match mode { + MessageDeliveryMode::QueueOnly => { + session + .services + .agent_control + .send_message_to_agent(turn.config.as_ref(), receiver_thread_id, communication) + .await + } + MessageDeliveryMode::TriggerTurn => { + session + .services + .agent_control + .send_followup_to_agent(turn.config.as_ref(), receiver_thread_id, communication) + .await + } + } + .map_err(|err| collab_agent_error(receiver_thread_id, err)); let status = session .services .agent_control diff --git a/codex-rs/core/src/tools/handlers/multi_agents_v2/spawn.rs b/codex-rs/core/src/tools/handlers/multi_agents_v2/spawn.rs index b5b7d210f590..8bdd6453a17d 100644 --- a/codex-rs/core/src/tools/handlers/multi_agents_v2/spawn.rs +++ b/codex-rs/core/src/tools/handlers/multi_agents_v2/spawn.rs @@ -188,7 +188,6 @@ async fn handle_spawn_agent( .as_ref() .and_then(|snapshot| snapshot.reasoning_effort.clone()) .unwrap_or(args.reasoning_effort.unwrap_or_default()); - let nickname = new_agent_nickname.clone(); session .send_event( &turn, @@ -220,15 +219,7 @@ async fn handle_spawn_agent( ) })?; - let hide_agent_metadata = turn.config.multi_agent_v2.hide_spawn_agent_metadata; - if hide_agent_metadata { - Ok(SpawnAgentResult::HiddenMetadata { task_name }) - } else { - Ok(SpawnAgentResult::WithNickname { - task_name, - nickname, - }) - } + Ok(SpawnAgentResult { task_name }) } impl CoreToolRuntime for Handler { @@ -288,15 +279,8 @@ impl SpawnAgentArgs { } #[derive(Debug, Serialize)] -#[serde(untagged)] -pub(crate) enum SpawnAgentResult { - WithNickname { - task_name: String, - nickname: Option, - }, - HiddenMetadata { - task_name: String, - }, +pub(crate) struct SpawnAgentResult { + task_name: String, } impl ToolOutput for SpawnAgentResult { diff --git a/codex-rs/state/src/runtime/threads.rs b/codex-rs/state/src/runtime/threads.rs index 4cc7ee4dc10c..a4bee6e63f60 100644 --- a/codex-rs/state/src/runtime/threads.rs +++ b/codex-rs/state/src/runtime/threads.rs @@ -115,6 +115,20 @@ ON CONFLICT(child_thread_id) DO UPDATE SET Ok(()) } + /// Return whether `child_thread_id` is still an open spawned agent. + pub async fn is_thread_spawn_edge_open( + &self, + child_thread_id: ThreadId, + ) -> anyhow::Result { + let status = sqlx::query_scalar::<_, String>( + "SELECT status FROM thread_spawn_edges WHERE child_thread_id = ?", + ) + .bind(child_thread_id.to_string()) + .fetch_optional(self.pool.as_ref()) + .await?; + Ok(status.as_deref() == Some(crate::DirectionalThreadSpawnEdgeStatus::Open.as_ref())) + } + /// List direct spawned children of `parent_thread_id` whose edge matches `status`. pub async fn list_thread_spawn_children_with_status( &self, @@ -181,8 +195,8 @@ LIMIT 2 one_thread_id_from_rows(rows, agent_path) } - /// Find a spawned descendant of `root_thread_id` by canonical agent path. - pub async fn find_thread_spawn_descendant_by_path( + /// Find an open spawned descendant of `root_thread_id` by canonical agent path. + pub async fn find_open_thread_spawn_descendant_by_path( &self, root_thread_id: ThreadId, agent_path: &str, @@ -193,10 +207,12 @@ WITH RECURSIVE subtree(child_thread_id) AS ( SELECT child_thread_id FROM thread_spawn_edges WHERE parent_thread_id = ? + AND status = ? UNION ALL SELECT edge.child_thread_id FROM thread_spawn_edges AS edge JOIN subtree ON edge.parent_thread_id = subtree.child_thread_id + WHERE edge.status = ? ) SELECT threads.id FROM subtree @@ -207,6 +223,8 @@ LIMIT 2 "#, ) .bind(root_thread_id.to_string()) + .bind(crate::DirectionalThreadSpawnEdgeStatus::Open.as_ref()) + .bind(crate::DirectionalThreadSpawnEdgeStatus::Open.as_ref()) .bind(agent_path) .fetch_all(self.pool.as_ref()) .await?;