feat: add agent execution slot accounting

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<Option<crate::agent::registry::ExecutionReservation>> {
+        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,

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<ActiveAgents>,
@@ -28,6 +23,7 @@ pub(crate) struct AgentRegistry {
 #[derive(Default)]
 struct ActiveAgents {
     agent_tree: HashMap<String, AgentMetadata>,
+    active_thread_ids: HashSet<ThreadId>,
     used_agent_nicknames: HashSet<String>,
     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<Self>,
+        thread_id: ThreadId,
+        max_threads: usize,
+    ) -> Result<ExecutionReservation> {
+        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<String> {
         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<AgentRegistry>,
+    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;

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);
 }
 

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(

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<Self>,
         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<Self>, 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;
     }