Fix completed subagent quota reclamation

Reclaim spawn quota slots when subagents reach final states, retry spawn after opportunistic cleanup, and keep interrupted subagents visible as active quota holders. This prevents completed subagents from leaking thread quota while preserving the ability to resume interrupted agents.

Co-authored-by: Open Codex <hff582580@gmail.com>

Author: LEON-gittech

codex-rs/core/src/agent/control.rs

@@ -200,7 +200,9 @@ impl AgentControl {
         options: SpawnAgentOptions,
     ) -> CodexResult<LiveAgent> {
         let state = self.upgrade()?;
-        let mut reservation = self.state.reserve_spawn_slot(config.agent_max_threads)?;
+        let mut reservation = self
+            .reserve_spawn_slot_with_final_reclaim(&state, &config)
+            .await?;
         let inherited_shell_snapshot = self
             .inherited_shell_snapshot_for_source(&state, session_source.as_ref())
             .await;
@@ -538,7 +540,9 @@ impl AgentControl {
         }
         let state = self.upgrade()?;
         let state_db_ctx = state.state_db();
-        let mut reservation = self.state.reserve_spawn_slot(config.agent_max_threads)?;
+        let mut reservation = self
+            .reserve_spawn_slot_with_final_reclaim(&state, &config)
+            .await?;
         let (session_source, agent_metadata) = match session_source {
             SessionSource::SubAgent(SubAgentSource::ThreadSpawn {
                 parent_thread_id,
@@ -638,6 +642,7 @@ impl AgentControl {
     ) -> CodexResult<String> {
         let last_task_message = render_input_preview(&initial_operation);
         let state = self.upgrade()?;
+        self.ensure_agent_slot_for_input(agent_id, &state).await?;
         let result = self
             .handle_thread_request_result(
                 agent_id,
@@ -675,6 +680,7 @@ impl AgentControl {
     ) -> CodexResult<String> {
         let last_task_message = communication.content.clone();
         let state = self.upgrade()?;
+        self.ensure_agent_slot_for_input(agent_id, &state).await?;
         let result = self
             .handle_thread_request_result(
                 agent_id,
@@ -710,6 +716,66 @@ impl AgentControl {
         result
     }
 
+    async fn reserve_spawn_slot_with_final_reclaim(
+        &self,
+        state: &Arc<ThreadManagerState>,
+        config: &crate::config::Config,
+    ) -> CodexResult<crate::agent::registry::SpawnReservation> {
+        match self.state.reserve_spawn_slot(config.agent_max_threads) {
+            Ok(reservation) => Ok(reservation),
+            Err(CodexErr::AgentLimitReached { .. }) => {
+                self.reclaim_final_agent_slots(state).await;
+                self.state.reserve_spawn_slot(config.agent_max_threads)
+            }
+            Err(err) => Err(err),
+        }
+    }
+
+    async fn ensure_agent_slot_for_input(
+        &self,
+        agent_id: ThreadId,
+        state: &Arc<ThreadManagerState>,
+    ) -> CodexResult<()> {
+        if !self.state.agent_slot_is_reclaimed(agent_id) {
+            return Ok(());
+        }
+        self.reclaim_final_agent_slots(state).await;
+        let thread = state.get_thread(agent_id).await?;
+        let config = thread.codex.session.get_config().await;
+        self.state
+            .occupy_reclaimed_thread_slot(agent_id, config.agent_max_threads)?;
+        Ok(())
+    }
+
+    async fn reclaim_final_agent_slots(&self, state: &Arc<ThreadManagerState>) {
+        for agent_id in self.state.occupied_agent_ids() {
+            let status = match state.get_thread(agent_id).await {
+                Ok(thread) => thread.agent_status().await,
+                Err(_) => AgentStatus::NotFound,
+            };
+            self.reclaim_final_agent_slot(agent_id, state, &status)
+                .await;
+        }
+    }
+
+    async fn reclaim_final_agent_slot(
+        &self,
+        agent_id: ThreadId,
+        state: &Arc<ThreadManagerState>,
+        status: &AgentStatus,
+    ) {
+        if !is_final(status) {
+            return;
+        }
+        if let Ok(thread) = state.get_thread(agent_id).await {
+            thread.codex.session.ensure_rollout_materialized().await;
+            if let Err(err) = thread.codex.session.flush_rollout().await {
+                warn!("failed to flush final subagent rollout before reclaiming slot: {err}");
+            }
+        }
+        self.state.reclaim_spawned_thread_slot(agent_id);
+    }
+
     /// 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<String> {
@@ -978,6 +1044,9 @@ impl AgentControl {
             };
             let child_thread = state.get_thread(child_thread_id).await.ok();
             let message = format_subagent_notification_message(child_reference.as_str(), &status);
+            control
+                .reclaim_final_agent_slot(child_thread_id, &state, &status)
+                .await;
             if child_agent_path.is_some()
                 && child_thread
                     .as_ref()
@@ -1051,6 +1120,7 @@ impl AgentControl {
             agent_nickname,
             agent_role,
             last_task_message: None,
+            slot_occupied: false,
         };
         Ok((session_source, agent_metadata))
     }

codex-rs/core/src/agent/control_tests.rs

@@ -1052,6 +1052,134 @@ async fn spawn_agent_releases_slot_after_shutdown() {
         .expect("shutdown agent");
 }
 
+#[tokio::test]
+async fn spawn_agent_reclaims_completed_slot_before_limit_error() {
+    let max_threads = 1usize;
+    let (_home, config) = test_config_with_cli_overrides(vec![(
+        "agents.max_threads".to_string(),
+        TomlValue::Integer(max_threads as i64),
+    )])
+    .await;
+    let manager = ThreadManager::with_models_provider_and_home_for_tests(
+        CodexAuth::from_api_key("dummy"),
+        config.model_provider.clone(),
+        config.codex_home.to_path_buf(),
+        std::sync::Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()),
+    );
+    let control = manager.agent_control();
+
+    let completed_agent_id = control
+        .spawn_agent(
+            config.clone(),
+            text_input("complete then release"),
+            /*session_source*/ None,
+        )
+        .await
+        .expect("spawn_agent should succeed");
+    let completed_agent = manager
+        .get_thread(completed_agent_id)
+        .await
+        .expect("completed agent should be registered");
+    let turn = completed_agent.codex.session.new_default_turn().await;
+    completed_agent
+        .codex
+        .session
+        .send_event(
+            turn.as_ref(),
+            EventMsg::TurnComplete(TurnCompleteEvent {
+                turn_id: turn.sub_id.clone(),
+                last_agent_message: Some("done".to_string()),
+                completed_at: None,
+                duration_ms: None,
+                time_to_first_token_ms: None,
+            }),
+        )
+        .await;
+
+    let next_agent_id = control
+        .spawn_agent(
+            config.clone(),
+            text_input("new work after completed agent"),
+            /*session_source*/ None,
+        )
+        .await
+        .expect("completed agent should not block the spawn quota");
+
+    let _ = control
+        .shutdown_live_agent(completed_agent_id)
+        .await
+        .expect("shutdown completed agent");
+    let _ = control
+        .shutdown_live_agent(next_agent_id)
+        .await
+        .expect("shutdown next agent");
+}
+
+#[tokio::test]
+async fn spawn_agent_does_not_reclaim_interrupted_slot() {
+    let max_threads = 1usize;
+    let (_home, config) = test_config_with_cli_overrides(vec![(
+        "agents.max_threads".to_string(),
+        TomlValue::Integer(max_threads as i64),
+    )])
+    .await;
+    let manager = ThreadManager::with_models_provider_and_home_for_tests(
+        CodexAuth::from_api_key("dummy"),
+        config.model_provider.clone(),
+        config.codex_home.to_path_buf(),
+        std::sync::Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()),
+    );
+    let control = manager.agent_control();
+
+    let interrupted_agent_id = control
+        .spawn_agent(
+            config.clone(),
+            text_input("interrupt and keep slot"),
+            /*session_source*/ None,
+        )
+        .await
+        .expect("spawn_agent should succeed");
+    let interrupted_agent = manager
+        .get_thread(interrupted_agent_id)
+        .await
+        .expect("interrupted agent should be registered");
+    let turn = interrupted_agent.codex.session.new_default_turn().await;
+    interrupted_agent
+        .codex
+        .session
+        .send_event(
+            turn.as_ref(),
+            EventMsg::TurnAborted(TurnAbortedEvent {
+                turn_id: Some("turn-1".to_string()),
+                reason: TurnAbortReason::Interrupted,
+                completed_at: None,
+                duration_ms: None,
+            }),
+        )
+        .await;
+
+    let err = control
+        .spawn_agent(
+            config,
+            text_input("should still be blocked"),
+            /*session_source*/ None,
+        )
+        .await
+        .expect_err("interrupted agent should still occupy the spawn quota");
+    let CodexErr::AgentLimitReached {
+        max_threads: seen_max_threads,
+    } = err
+    else {
+        panic!("expected CodexErr::AgentLimitReached");
+    };
+    assert_eq!(seen_max_threads, max_threads);
+
+    let _ = control
+        .shutdown_live_agent(interrupted_agent_id)
+        .await
+        .expect("shutdown interrupted agent");
+}
+
 #[tokio::test]
 async fn spawn_agent_limit_shared_across_clones() {
     let max_threads = 1usize;