From 4e803a017c958dd37eb251372ba810232d3e84ba Mon Sep 17 00:00:00 2001 From: jif Date: Mon, 8 Jun 2026 10:24:48 +0200 Subject: [PATCH] feat: add v2 agent residency lru (#26632) ## Why Multi-agent v2 treats agents as durable logical agents, not just live entries in `ThreadManager`. After the reload-on-delivery change, a v2 agent can be addressed even if its thread is not currently loaded. This PR adds the next layer: loaded v2 subagents can be paged out of `ThreadManager` when the session has too many resident agents. That keeps residency separate from logical identity and prepares the stack for making v2 concurrency count active execution instead of existing agents. ## What Changed - Add an `AgentControl`-scoped LRU for resident v2 subagents. - Reserve residency before spawning or reloading a v2 subagent. - If resident capacity is full, unload the least-recently-used idle v2 subagent from `ThreadManager`. - Keep `ThreadManager` as a primitive loaded-thread store; it does not own the LRU policy. - Keep unloaded agents registered and durable so they can be reloaded by the delivery path. - Preserve the existing v2 cap semantics by using the derived non-root v2 cap for residency. Eviction is intentionally conservative. A thread is unloadable only when it is a v2 subagent, has completed or errored, has no active turn, and has no pending mailbox work. Before removal, the rollout is materialized and flushed. ## Assumptions And Non-Goals - PR #26623 provides the reload-on-delivery path for unloaded v2 agents. - `ThreadManager` membership means loaded/resident, not logical agent existence. - `AgentRegistry` remains the logical identity/metadata source for v2 agents that may be unloaded. - `list_agents` remains a recent/resident view for now. - This does not change active execution concurrency; that is the next PR. - This does not change `close_agent` semantics. - This does not change or remove `resume_agent`. - This does not add a new residency config knob. ## Stack 1. V2 durable lookup and reload on delivery (#26623) - reload unloaded v2 agents before delivering follow-up/input. 2. V2 residency LRU (this PR) - unload idle resident v2 agents from `ThreadManager` when resident capacity is full. 3. V2 active-execution concurrency - count running non-root v2 turns instead of logical agents. 4. V2 close/interrupt semantics - make v2 close interrupt the current turn without deleting durable identity. 5. V2 resume cleanup - remove the manual resume surface for v2 while keeping internal reload support. ## Validation - Added focused coverage for the residency LRU eviction path. - Local clippy/check/tests were not run; CI will cover them. --- codex-rs/core/src/agent/control.rs | 5 + codex-rs/core/src/agent/control/legacy.rs | 1 + codex-rs/core/src/agent/control/residency.rs | 240 ++++++++++++++++++ .../core/src/agent/control/residency_tests.rs | 107 ++++++++ codex-rs/core/src/agent/control/spawn.rs | 30 ++- 5 files changed, 382 insertions(+), 1 deletion(-) create mode 100644 codex-rs/core/src/agent/control/residency.rs create mode 100644 codex-rs/core/src/agent/control/residency_tests.rs diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index 6c6650a55..e9994883f 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/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, state: Arc, + v2_residency: Arc, } impl AgentControl { @@ -183,6 +187,7 @@ impl AgentControl { ) -> CodexResult { 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 diff --git a/codex-rs/core/src/agent/control/legacy.rs b/codex-rs/core/src/agent/control/legacy.rs index e85090b90..aa4147268 100644 --- a/codex-rs/core/src/agent/control/legacy.rs +++ b/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 } diff --git a/codex-rs/core/src/agent/control/residency.rs b/codex-rs/core/src/agent/control/residency.rs new file mode 100644 index 000000000..8056fedeb --- /dev/null +++ b/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, +} + +#[derive(Default)] +struct V2ResidencyState { + residents: VecDeque, + pending_slots: usize, +} + +pub(super) struct V2ResidencySlot { + residency: Arc, + 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, + config: &Config, + protected_thread_id: Option, + ) -> CodexResult { + 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, + 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, + manager: &Arc, + capacity: usize, + protected_thread_id: Option, + ) -> CodexResult { + 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, + protected_thread_id: Option, + ) -> 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) -> Option { + 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, 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; diff --git a/codex-rs/core/src/agent/control/residency_tests.rs b/codex-rs/core/src/agent/control/residency_tests.rs new file mode 100644 index 000000000..a0338eda0 --- /dev/null +++ b/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, + 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; +} diff --git a/codex-rs/core/src/agent/control/spawn.rs b/codex-rs/core/src/agent/control/spawn.rs index 8df3b0a07..e3536e89f 100644 --- a/codex-rs/core/src/agent/control/spawn.rs +++ b/codex-rs/core/src/agent/control/spawn.rs @@ -1,3 +1,4 @@ +use super::residency::is_v2_resident_session_source; use super::*; const AGENT_NAMES: &str = include_str!("../agent_names.txt"); @@ -116,6 +117,7 @@ impl AgentControl { ) -> CodexResult<()> { let state = self.upgrade()?; if state.get_thread(thread_id).await.is_ok() { + self.touch_loaded_v2_residency(&state, thread_id).await; return Ok(()); } if self.state.agent_metadata_for_thread(thread_id).is_none() { @@ -143,6 +145,9 @@ impl AgentControl { if initial_history.get_multi_agent_version() != Some(MultiAgentVersion::V2) { return Err(CodexErr::ThreadNotFound(thread_id)); } + let residency_slot = self + .reserve_v2_residency_slot(&state, &config, Some(thread_id)) + .await?; let (session_source, _) = initial_history .get_resumed_session_sources() @@ -170,11 +175,14 @@ impl AgentControl { .await { Ok(reloaded_thread) => { + residency_slot.commit(reloaded_thread.thread_id); state.notify_thread_created(reloaded_thread.thread_id); Ok(()) } Err(err) => { if state.get_thread(thread_id).await.is_ok() { + drop(residency_slot); + self.touch_loaded_v2_residency(&state, thread_id).await; return Ok(()); } Err(err) @@ -200,7 +208,24 @@ impl AgentControl { ) .await; let agent_max_threads = config.effective_agent_max_threads(multi_agent_version); - let mut reservation = self.state.reserve_spawn_slot(agent_max_threads)?; + let spawn_uses_v2_residency = multi_agent_version == MultiAgentVersion::V2 + && session_source + .as_ref() + .is_some_and(is_v2_resident_session_source); + let residency_slot = if spawn_uses_v2_residency { + Some( + self.reserve_v2_residency_slot(&state, &config, /*protected_thread_id*/ None) + .await?, + ) + } else { + None + }; + let reservation_max_threads = if spawn_uses_v2_residency { + None + } else { + agent_max_threads + }; + let mut reservation = self.state.reserve_spawn_slot(reservation_max_threads)?; let inheritance = SpawnAgentThreadInheritance { shell_snapshot: self .inherited_shell_snapshot_for_source(&state, session_source.as_ref()) @@ -264,6 +289,9 @@ impl AgentControl { }; agent_metadata.agent_id = Some(new_thread.thread_id); reservation.commit(agent_metadata.clone()); + if let Some(residency_slot) = residency_slot { + residency_slot.commit(new_thread.thread_id); + } if let Some(SessionSource::SubAgent( subagent_source @ SubAgentSource::ThreadSpawn {