diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index 897d162df5..a6e945b5f6 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/codex-rs/core/src/agent/control.rs @@ -200,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( @@ -245,6 +253,37 @@ impl AgentControl { self.state.agent_metadata_for_thread(agent_id) } + pub(crate) async fn list_live_agent_subtree_thread_ids( + &self, + agent_id: ThreadId, + ) -> CodexResult> { + let mut thread_ids = vec![agent_id]; + thread_ids.extend(self.live_thread_spawn_descendants(agent_id).await?); + Ok(thread_ids) + } + + pub(crate) async fn 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, @@ -309,6 +348,77 @@ impl AgentControl { Ok(thread.subscribe_status()) } + async fn open_thread_spawn_children( + &self, + parent_thread_id: ThreadId, + ) -> CodexResult> { + let mut children_by_parent = self.live_thread_spawn_children().await?; + Ok(children_by_parent + .remove(&parent_thread_id) + .unwrap_or_default()) + } + + async fn live_thread_spawn_children( + &self, + ) -> CodexResult>> { + let state = self.upgrade()?; + let mut children_by_parent = HashMap::>::new(); + + for (parent_thread_id, child_thread_id) in state.list_live_thread_spawn_edges().await { + children_by_parent + .entry(parent_thread_id) + .or_default() + .push(( + child_thread_id, + self.state + .agent_metadata_for_thread(child_thread_id) + .unwrap_or(AgentMetadata { + agent_id: Some(child_thread_id), + ..Default::default() + }), + )); + } + + for children in children_by_parent.values_mut() { + children.sort_by(|left, right| { + left.1 + .agent_path + .as_deref() + .unwrap_or_default() + .cmp(right.1.agent_path.as_deref().unwrap_or_default()) + .then_with(|| left.0.to_string().cmp(&right.0.to_string())) + }); + } + + Ok(children_by_parent) + } + + async fn live_thread_spawn_descendants( + &self, + root_thread_id: ThreadId, + ) -> CodexResult> { + let mut children_by_parent = self.live_thread_spawn_children().await?; + let mut descendants = Vec::new(); + let mut stack = children_by_parent + .remove(&root_thread_id) + .unwrap_or_default() + .into_iter() + .map(|(child_thread_id, _)| child_thread_id) + .rev() + .collect::>(); + + while let Some(thread_id) = stack.pop() { + descendants.push(thread_id); + if let Some(children) = children_by_parent.remove(&thread_id) { + for (child_thread_id, _) in children.into_iter().rev() { + stack.push(child_thread_id); + } + } + } + + Ok(descendants) + } + pub(super) fn upgrade(&self) -> CodexResult> { self.manager .upgrade() diff --git a/codex-rs/core/src/agent/control/legacy.rs b/codex-rs/core/src/agent/control/legacy.rs index 98d9dce3fa..b75fe78594 100644 --- a/codex-rs/core/src/agent/control/legacy.rs +++ b/codex-rs/core/src/agent/control/legacy.rs @@ -82,37 +82,6 @@ impl AgentControl { } } - pub(crate) async fn list_live_agent_subtree_thread_ids( - &self, - agent_id: ThreadId, - ) -> CodexResult> { - let mut thread_ids = vec![agent_id]; - thread_ids.extend(self.live_thread_spawn_descendants(agent_id).await?); - Ok(thread_ids) - } - - pub(crate) async fn 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(super) fn maybe_start_completion_watcher( &self, child_thread_id: ThreadId, @@ -157,75 +126,4 @@ impl AgentControl { .await; }); } - - async fn open_thread_spawn_children( - &self, - parent_thread_id: ThreadId, - ) -> CodexResult> { - let mut children_by_parent = self.live_thread_spawn_children().await?; - Ok(children_by_parent - .remove(&parent_thread_id) - .unwrap_or_default()) - } - - async fn live_thread_spawn_children( - &self, - ) -> CodexResult>> { - let state = self.upgrade()?; - let mut children_by_parent = HashMap::>::new(); - - for (parent_thread_id, child_thread_id) in state.list_live_thread_spawn_edges().await { - children_by_parent - .entry(parent_thread_id) - .or_default() - .push(( - child_thread_id, - self.state - .agent_metadata_for_thread(child_thread_id) - .unwrap_or(AgentMetadata { - agent_id: Some(child_thread_id), - ..Default::default() - }), - )); - } - - for children in children_by_parent.values_mut() { - children.sort_by(|left, right| { - left.1 - .agent_path - .as_deref() - .unwrap_or_default() - .cmp(right.1.agent_path.as_deref().unwrap_or_default()) - .then_with(|| left.0.to_string().cmp(&right.0.to_string())) - }); - } - - Ok(children_by_parent) - } - - async fn live_thread_spawn_descendants( - &self, - root_thread_id: ThreadId, - ) -> CodexResult> { - let mut children_by_parent = self.live_thread_spawn_children().await?; - let mut descendants = Vec::new(); - let mut stack = children_by_parent - .remove(&root_thread_id) - .unwrap_or_default() - .into_iter() - .map(|(child_thread_id, _)| child_thread_id) - .rev() - .collect::>(); - - while let Some(thread_id) = stack.pop() { - descendants.push(thread_id); - if let Some(children) = children_by_parent.remove(&thread_id) { - for (child_thread_id, _) in children.into_iter().rev() { - stack.push(child_thread_id); - } - } - } - - Ok(descendants) - } } diff --git a/codex-rs/core/src/agent/control/resident.rs b/codex-rs/core/src/agent/control/resident.rs index 0b78f86823..7a73a54f4c 100644 --- a/codex-rs/core/src/agent/control/resident.rs +++ b/codex-rs/core/src/agent/control/resident.rs @@ -12,6 +12,34 @@ fn max_resident_subagents(config: &Config) -> usize { } impl AgentControl { + pub(crate) async fn reserve_execution_slot_for_pending_turn( + &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, @@ -192,12 +220,20 @@ impl AgentControl { }; thread.ensure_rollout_materialized().await; thread.flush_rollout().await?; + let mut wait_for_termination = false; if !matches!(thread.agent_status().await, AgentStatus::Shutdown) { - state.send_op(agent_id, Op::Shutdown {}).await?; + 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; } - 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(()) } diff --git a/codex-rs/core/src/agent/control_tests.rs b/codex-rs/core/src/agent/control_tests.rs index 7e8d33acdb..8a8666dc04 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 completion_watcher_notifies_parent_when_child_is_missing() { let harness = AgentControlHarness::new().await; diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index ce03d480e0..a3b9936ca9 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -469,6 +469,22 @@ impl Session { return; } + let execution_reservation = match self + .services + .agent_control + .reserve_execution_slot_for_pending_turn(self) + .await + { + Ok(reservation) => reservation, + Err(err) => { + warn!( + thread_id = %self.thread_id, + "pending agent turn could not reserve an execution slot: {err}" + ); + return; + } + }; + { let mut active_turn = self.active_turn.lock().await; if active_turn.is_some() { @@ -476,6 +492,9 @@ impl Session { } *active_turn = Some(ActiveTurn::default()); } + if let Some(execution_reservation) = execution_reservation { + execution_reservation.commit(); + } 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())