feat: add v2 agent residency lru

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

@@ -42,9 +42,12 @@ use std::sync::Weak;
 use tokio::sync::watch;
 use tracing::warn;
 
+use self::residency::V2Residency;
+
 const ROOT_LAST_TASK_MESSAGE: &str = "Main thread";
 
 mod legacy;
+mod residency;
 mod spawn;
 
 #[derive(Clone, Debug, PartialEq, Eq)]
@@ -91,6 +94,7 @@ pub(crate) struct AgentControl {
     /// `ThreadManagerState -> CodexThread -> Session -> SessionServices -> ThreadManagerState`.
     manager: Weak<ThreadManagerState>,
     state: Arc<AgentRegistry>,
+    v2_residency: Arc<V2Residency>,
 }
 
 impl AgentControl {
@@ -183,6 +187,7 @@ impl AgentControl {
     ) -> CodexResult<String> {
         if matches!(result, Err(CodexErr::InternalAgentDied)) {
             let _ = state.remove_thread(&agent_id).await;
+            self.forget_v2_residency(agent_id);
             self.state.release_spawned_thread(agent_id);
         }
         result

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

@@ -19,6 +19,7 @@ impl AgentControl {
             state.send_op(agent_id, Op::Shutdown {}).await
         };
         let _ = state.remove_thread(&agent_id).await;
+        self.forget_v2_residency(agent_id);
         self.state.release_spawned_thread(agent_id);
         result
     }

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

@@ -0,0 +1,240 @@
+use super::AgentControl;
+use crate::agent::AgentStatus;
+use crate::codex_thread::CodexThread;
+use crate::config::Config;
+use crate::thread_manager::ThreadManagerState;
+use codex_protocol::ThreadId;
+use codex_protocol::error::CodexErr;
+use codex_protocol::error::Result as CodexResult;
+use codex_protocol::protocol::MultiAgentVersion;
+use codex_protocol::protocol::SessionSource;
+use std::collections::VecDeque;
+use std::sync::Arc;
+use std::sync::Mutex;
+use tracing::warn;
+
+#[derive(Default)]
+pub(super) struct V2Residency {
+    state: Mutex<V2ResidencyState>,
+}
+
+#[derive(Default)]
+struct V2ResidencyState {
+    residents: VecDeque<ThreadId>,
+    pending_slots: usize,
+}
+
+pub(super) struct V2ResidencySlot {
+    residency: Arc<V2Residency>,
+    active: bool,
+}
+
+impl V2ResidencySlot {
+    pub(super) fn commit(mut self, thread_id: ThreadId) {
+        self.residency.commit_slot(thread_id);
+        self.active = false;
+    }
+}
+
+impl Drop for V2ResidencySlot {
+    fn drop(&mut self) {
+        if self.active {
+            self.residency.release_pending_slot();
+        }
+    }
+}
+
+impl AgentControl {
+    pub(super) async fn reserve_v2_residency_slot(
+        &self,
+        state: &Arc<ThreadManagerState>,
+        config: &Config,
+        protected_thread_id: Option<ThreadId>,
+    ) -> CodexResult<V2ResidencySlot> {
+        let capacity = config
+            .effective_agent_max_threads(MultiAgentVersion::V2)
+            .unwrap_or(usize::MAX);
+        Arc::clone(&self.v2_residency)
+            .reserve_slot(state, capacity, protected_thread_id)
+            .await
+    }
+
+    pub(super) async fn touch_loaded_v2_residency(
+        &self,
+        state: &Arc<ThreadManagerState>,
+        thread_id: ThreadId,
+    ) {
+        if let Ok(thread) = state.get_thread(thread_id).await
+            && is_resident_candidate(thread.as_ref())
+        {
+            self.v2_residency.touch(thread_id);
+        }
+    }
+
+    pub(super) fn forget_v2_residency(&self, thread_id: ThreadId) {
+        self.v2_residency.remove(thread_id);
+    }
+}
+
+impl V2Residency {
+    async fn reserve_slot(
+        self: Arc<Self>,
+        manager: &Arc<ThreadManagerState>,
+        capacity: usize,
+        protected_thread_id: Option<ThreadId>,
+    ) -> CodexResult<V2ResidencySlot> {
+        loop {
+            if self.try_reserve_pending_slot(capacity) {
+                return Ok(V2ResidencySlot {
+                    residency: self,
+                    active: true,
+                });
+            }
+            if !self
+                .try_unload_one_resident(manager, protected_thread_id)
+                .await
+            {
+                return Err(CodexErr::AgentLimitReached {
+                    max_threads: capacity,
+                });
+            }
+        }
+    }
+
+    fn try_reserve_pending_slot(&self, capacity: usize) -> bool {
+        let mut state = self
+            .state
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner);
+        if state.residents.len().saturating_add(state.pending_slots) >= capacity {
+            return false;
+        }
+        state.pending_slots += 1;
+        true
+    }
+
+    async fn try_unload_one_resident(
+        &self,
+        manager: &Arc<ThreadManagerState>,
+        protected_thread_id: Option<ThreadId>,
+    ) -> bool {
+        let candidates_to_scan = self.resident_count();
+        for _ in 0..candidates_to_scan {
+            let Some(candidate_thread_id) = self.pop_lru_candidate(protected_thread_id) else {
+                return false;
+            };
+            let Some(candidate_thread) = manager
+                .get_thread(candidate_thread_id)
+                .await
+                .ok()
+                .filter(|thread| is_resident_candidate(thread))
+            else {
+                continue;
+            };
+            if !is_unloadable(candidate_thread.as_ref()).await {
+                self.touch(candidate_thread_id);
+                continue;
+            }
+            candidate_thread.ensure_rollout_materialized().await;
+            if let Err(err) = candidate_thread.flush_rollout().await {
+                warn!(
+                    "failed to flush v2 resident thread before unloading {candidate_thread_id}: {err}"
+                );
+                self.touch(candidate_thread_id);
+                continue;
+            }
+            let _ = manager.remove_thread(&candidate_thread_id).await;
+            return true;
+        }
+        false
+    }
+
+    fn resident_count(&self) -> usize {
+        self.state
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner)
+            .residents
+            .len()
+    }
+
+    fn pop_lru_candidate(&self, protected_thread_id: Option<ThreadId>) -> Option<ThreadId> {
+        let mut state = self
+            .state
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner);
+        let candidates_to_scan = state.residents.len();
+        for _ in 0..candidates_to_scan {
+            let candidate_thread_id = state.residents.pop_front()?;
+            if Some(candidate_thread_id) == protected_thread_id {
+                state.residents.push_back(candidate_thread_id);
+                continue;
+            }
+            return Some(candidate_thread_id);
+        }
+        None
+    }
+
+    fn touch(&self, thread_id: ThreadId) {
+        let mut state = self
+            .state
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner);
+        touch_resident(&mut state.residents, thread_id);
+    }
+
+    fn remove(&self, thread_id: ThreadId) {
+        self.state
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner)
+            .residents
+            .retain(|resident_thread_id| *resident_thread_id != thread_id);
+    }
+
+    fn commit_slot(&self, thread_id: ThreadId) {
+        let mut state = self
+            .state
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner);
+        state.pending_slots = state.pending_slots.saturating_sub(1);
+        touch_resident(&mut state.residents, thread_id);
+    }
+
+    fn release_pending_slot(&self) {
+        let mut state = self
+            .state
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner);
+        state.pending_slots = state.pending_slots.saturating_sub(1);
+    }
+}
+
+fn touch_resident(residents: &mut VecDeque<ThreadId>, thread_id: ThreadId) {
+    residents.retain(|resident_thread_id| *resident_thread_id != thread_id);
+    residents.push_back(thread_id);
+}
+
+fn is_resident_candidate(thread: &CodexThread) -> bool {
+    thread.multi_agent_version() == Some(MultiAgentVersion::V2)
+        && is_v2_resident_session_source(&thread.session_source)
+}
+
+pub(super) fn is_v2_resident_session_source(session_source: &SessionSource) -> bool {
+    matches!(session_source, SessionSource::SubAgent(_))
+}
+
+async fn is_unloadable(thread: &CodexThread) -> bool {
+    matches!(
+        thread.agent_status().await,
+        AgentStatus::Completed(_) | AgentStatus::Errored(_)
+    ) && thread.codex.session.active_turn.lock().await.is_none()
+        && !thread
+            .codex
+            .session
+            .input_queue
+            .has_pending_mailbox_items()
+            .await
+}
+
+#[cfg(test)]
+#[path = "residency_tests.rs"]
+mod tests;

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

@@ -0,0 +1,107 @@
+use crate::ThreadManager;
+use crate::agent::AgentControl;
+use crate::codex_thread::CodexThread;
+use crate::config::Config;
+use crate::config::test_config;
+use crate::thread_manager::ThreadManagerState;
+use codex_features::Feature;
+use codex_login::CodexAuth;
+use codex_protocol::ThreadId;
+use codex_protocol::error::CodexErr;
+use codex_protocol::protocol::EventMsg;
+use codex_protocol::protocol::SessionSource;
+use codex_protocol::protocol::SubAgentSource;
+use codex_protocol::protocol::ThreadSource;
+use codex_protocol::protocol::TurnCompleteEvent;
+use pretty_assertions::assert_eq;
+use std::sync::Arc;
+
+#[tokio::test]
+async fn residency_slot_reservation_unloads_oldest_idle_v2_agent() {
+    let mut config = test_config().await;
+    let _ = config.features.enable(Feature::MultiAgentV2);
+    config.multi_agent_v2.max_concurrent_threads_per_session = 2;
+    let temp_home = tempfile::tempdir().expect("create temp home");
+    config.codex_home = temp_home.path().to_path_buf().try_into().unwrap();
+    config.cwd = temp_home.path().to_path_buf().try_into().unwrap();
+    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(),
+        Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()),
+    );
+    let root = manager
+        .start_thread(config.clone())
+        .await
+        .expect("start root thread");
+    let control = manager.agent_control();
+    let state = control.upgrade().expect("thread manager should be live");
+
+    let first_slot = control
+        .reserve_v2_residency_slot(&state, &config, /*protected_thread_id*/ None)
+        .await
+        .expect("first resident slot");
+    let first =
+        spawn_v2_subagent(&control, &state, config.clone(), root.thread_id, "worker-1").await;
+    first_slot.commit(first.thread_id);
+    mark_thread_completed(first.thread.as_ref()).await;
+
+    let second_slot = control
+        .reserve_v2_residency_slot(&state, &config, /*protected_thread_id*/ None)
+        .await
+        .expect("second resident slot should evict the first idle agent");
+    match manager.get_thread(first.thread_id).await {
+        Err(CodexErr::ThreadNotFound(thread_id)) => assert_eq!(thread_id, first.thread_id),
+        Err(err) => panic!("expected evicted thread to be missing, got {err:?}"),
+        Ok(_) => panic!("expected evicted thread to be missing"),
+    }
+    let second = spawn_v2_subagent(&control, &state, config, root.thread_id, "worker-2").await;
+    second_slot.commit(second.thread_id);
+
+    assert!(manager.get_thread(root.thread_id).await.is_ok());
+    assert!(manager.get_thread(second.thread_id).await.is_ok());
+}
+
+async fn spawn_v2_subagent(
+    control: &AgentControl,
+    state: &Arc<ThreadManagerState>,
+    config: Config,
+    parent_thread_id: ThreadId,
+    label: &str,
+) -> crate::thread_manager::NewThread {
+    state
+        .spawn_new_thread_with_source(
+            config,
+            control.clone(),
+            SessionSource::SubAgent(SubAgentSource::Other(label.to_string())),
+            Some(parent_thread_id),
+            /*forked_from_thread_id*/ None,
+            Some(ThreadSource::Subagent),
+            /*metrics_service_name*/ None,
+            /*inherited_shell_snapshot*/ None,
+            /*inherited_exec_policy*/ None,
+            /*environments*/ None,
+        )
+        .await
+        .expect("spawn v2 subagent")
+}
+
+async fn mark_thread_completed(thread: &CodexThread) {
+    let turn = thread.codex.session.new_default_turn().await;
+    thread
+        .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;
+    // The fixture has no task runner to clear the turn after the terminal event.
+    *thread.codex.session.active_turn.lock().await = None;
+}