mirror of
https://github.com/pchuan98/codex.git
synced 2026-07-01 00:31:56 +08:00
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.
This commit is contained in:
jif
GitHub
@@ -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
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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;
|
||||
@@ -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;
|
||||
}
|
||||
@@ -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 {
|
||||
|
||||
Reference in New Issue
Block a user