Codex Handbook
core/src/agent/control/residency.rs 240 lines
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.shutdown_and_wait().await {                warn!(                    "failed to shut down 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(_) | AgentStatus::Interrupted    ) && 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;