From cfeccd51a0cf79d6957093d4a59fb63d81fb6069 Mon Sep 17 00:00:00 2001 From: jif-oai Date: Fri, 5 Jun 2026 16:03:54 +0200 Subject: [PATCH] feat: add agent execution slot accounting --- codex-rs/core/src/agent/control.rs | 43 ++++++++ codex-rs/core/src/agent/registry.rs | 116 +++++++++++++++++----- codex-rs/core/src/agent/registry_tests.rs | 17 +++- codex-rs/core/src/tasks/lifecycle.rs | 8 ++ codex-rs/core/src/tasks/mod.rs | 56 ++++++++--- 5 files changed, 196 insertions(+), 44 deletions(-) diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index 6c6650a559..60dbaa5456 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/codex-rs/core/src/agent/control.rs @@ -7,6 +7,7 @@ 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; @@ -267,6 +268,48 @@ impl AgentControl { Ok(thread.subscribe_status()) } + 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 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; + } + 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); + } + } + pub(crate) async fn format_environment_context_subagents( &self, parent_thread_id: ThreadId, diff --git a/codex-rs/core/src/agent/registry.rs b/codex-rs/core/src/agent/registry.rs index 43aca201bf..37e457e35c 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,27 +93,61 @@ 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) { + 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); + } + } + + 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 register_root_thread(&self, thread_id: ThreadId) { let mut active_agents = self .active_agents @@ -194,7 +224,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 +240,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 +383,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 +399,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 fc172fb336..44349de5c1 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/tasks/lifecycle.rs b/codex-rs/core/src/tasks/lifecycle.rs index 782cc8c7de..59c46e617d 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 ef2ba0f16d..a3b9936ca9 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -460,27 +460,48 @@ 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 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() { + return; + } + *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()) + .await; + self.start_task(turn_context, Vec::new(), RegularTask::new()) + .await; + }) } pub async fn abort_all_tasks(self: &Arc, reason: TurnAbortReason) { @@ -519,6 +540,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 +588,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 +820,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; }