mirror of
https://github.com/pchuan98/codex.git
synced 2026-07-01 00:31:56 +08:00
Avoid reopening v2 descendants on resume (#26997)
## Why Multi-agent v2 residency is intended to keep only the threads that need to be live. The existing rollout resume path still walked persisted open descendants and reopened the entire descendant tree when resuming a v2 root, which turns resume into an eager reload of work that should stay unloaded until it is explicitly needed. The interrupted-agent path has a related residency issue. Interrupted agents remain open by design, so an idle interrupted resident should be eligible for eviction just like an idle completed or errored resident. Otherwise a resident set full of interrupted agents can consume every v2 slot and block later spawns or reloads with `AgentLimitReached`. ## What Changed - Return early from `resume_agent_from_rollout` after resuming a v2 thread so persisted v2 descendants are not reopened eagerly. - Treat idle `Interrupted` v2 residents as unloadable in the LRU residency path. - Add focused coverage for v2 root resume leaving descendants unloaded and for eviction of an idle interrupted v2 resident when a new slot is needed. ## Verification Added targeted `codex-core` tests covering: - v2 root resume with persisted descendants, verifying only the root is loaded after resume. - residency eviction of an idle interrupted v2 agent when the resident set is full.
This commit is contained in:
@@ -136,9 +136,9 @@ impl V2Residency {
|
||||
continue;
|
||||
}
|
||||
candidate_thread.ensure_rollout_materialized().await;
|
||||
if let Err(err) = candidate_thread.flush_rollout().await {
|
||||
if let Err(err) = candidate_thread.shutdown_and_wait().await {
|
||||
warn!(
|
||||
"failed to flush v2 resident thread before unloading {candidate_thread_id}: {err}"
|
||||
"failed to shut down v2 resident thread before unloading {candidate_thread_id}: {err}"
|
||||
);
|
||||
self.touch(candidate_thread_id);
|
||||
continue;
|
||||
@@ -225,7 +225,7 @@ pub(super) fn is_v2_resident_session_source(session_source: &SessionSource) -> b
|
||||
async fn is_unloadable(thread: &CodexThread) -> bool {
|
||||
matches!(
|
||||
thread.agent_status().await,
|
||||
AgentStatus::Completed(_) | AgentStatus::Errored(_)
|
||||
AgentStatus::Completed(_) | AgentStatus::Errored(_) | AgentStatus::Interrupted
|
||||
) && thread.codex.session.active_turn.lock().await.is_none()
|
||||
&& !thread
|
||||
.codex
|
||||
|
||||
@@ -12,6 +12,8 @@ use codex_protocol::protocol::EventMsg;
|
||||
use codex_protocol::protocol::SessionSource;
|
||||
use codex_protocol::protocol::SubAgentSource;
|
||||
use codex_protocol::protocol::ThreadSource;
|
||||
use codex_protocol::protocol::TurnAbortReason;
|
||||
use codex_protocol::protocol::TurnAbortedEvent;
|
||||
use codex_protocol::protocol::TurnCompleteEvent;
|
||||
use pretty_assertions::assert_eq;
|
||||
use std::sync::Arc;
|
||||
@@ -62,6 +64,68 @@ async fn residency_slot_reservation_unloads_oldest_idle_v2_agent() {
|
||||
assert!(manager.get_thread(second.thread_id).await.is_ok());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn interrupted_v2_agent_is_lost_after_residency_eviction() {
|
||||
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_interrupted(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 interrupted 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.clone(), root.thread_id, "worker-2").await;
|
||||
second_slot.commit(second.thread_id);
|
||||
mark_thread_completed(second.thread.as_ref()).await;
|
||||
|
||||
let err = control
|
||||
.ensure_v2_agent_loaded(config, first.thread_id)
|
||||
.await
|
||||
.expect_err("evicted interrupted agent should stay lost");
|
||||
match err {
|
||||
CodexErr::ThreadNotFound(thread_id) => assert_eq!(thread_id, first.thread_id),
|
||||
err => panic!("expected ThreadNotFound, got {err:?}"),
|
||||
}
|
||||
|
||||
assert!(manager.get_thread(root.thread_id).await.is_ok());
|
||||
assert!(manager.get_thread(second.thread_id).await.is_ok());
|
||||
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"),
|
||||
}
|
||||
}
|
||||
|
||||
async fn spawn_v2_subagent(
|
||||
control: &AgentControl,
|
||||
state: &Arc<ThreadManagerState>,
|
||||
@@ -102,6 +166,28 @@ async fn mark_thread_completed(thread: &CodexThread) {
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
clear_active_turn(thread).await;
|
||||
}
|
||||
|
||||
async fn mark_thread_interrupted(thread: &CodexThread) {
|
||||
let turn = thread.codex.session.new_default_turn().await;
|
||||
thread
|
||||
.codex
|
||||
.session
|
||||
.send_event(
|
||||
turn.as_ref(),
|
||||
EventMsg::TurnAborted(TurnAbortedEvent {
|
||||
turn_id: Some(turn.sub_id.clone()),
|
||||
reason: TurnAbortReason::Interrupted,
|
||||
completed_at: None,
|
||||
duration_ms: None,
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
clear_active_turn(thread).await;
|
||||
}
|
||||
|
||||
async fn clear_active_turn(thread: &CodexThread) {
|
||||
// The fixture has no task runner to clear the turn after the terminal event.
|
||||
*thread.codex.session.active_turn.lock().await = None;
|
||||
}
|
||||
|
||||
@@ -527,13 +527,16 @@ impl AgentControl {
|
||||
session_source: SessionSource,
|
||||
) -> CodexResult<ThreadId> {
|
||||
let root_depth = thread_spawn_depth(&session_source).unwrap_or(0);
|
||||
let resumed_thread_id = Box::pin(self.resume_single_agent_from_rollout(
|
||||
config.clone(),
|
||||
thread_id,
|
||||
session_source,
|
||||
))
|
||||
let (resumed_thread_id, resumed_multi_agent_version) = Box::pin(
|
||||
self.resume_single_agent_from_rollout(config.clone(), thread_id, session_source),
|
||||
)
|
||||
.await?;
|
||||
let state = self.upgrade()?;
|
||||
if config.multi_agent_version_from_features() == MultiAgentVersion::V2
|
||||
|| resumed_multi_agent_version == MultiAgentVersion::V2
|
||||
{
|
||||
return Ok(resumed_thread_id);
|
||||
}
|
||||
let Ok(resumed_thread) = state.get_thread(resumed_thread_id).await else {
|
||||
return Ok(resumed_thread_id);
|
||||
};
|
||||
@@ -579,7 +582,7 @@ impl AgentControl {
|
||||
))
|
||||
.await
|
||||
{
|
||||
Ok(_) => true,
|
||||
Ok((_, _)) => true,
|
||||
Err(err) => {
|
||||
warn!("failed to resume descendant thread {child_thread_id}: {err}");
|
||||
false
|
||||
@@ -600,7 +603,7 @@ impl AgentControl {
|
||||
config: Config,
|
||||
thread_id: ThreadId,
|
||||
session_source: SessionSource,
|
||||
) -> CodexResult<ThreadId> {
|
||||
) -> CodexResult<(ThreadId, MultiAgentVersion)> {
|
||||
let state = self.upgrade()?;
|
||||
let state_db_ctx = state.state_db();
|
||||
let stored_thread = state
|
||||
@@ -705,6 +708,6 @@ impl AgentControl {
|
||||
)
|
||||
.await;
|
||||
|
||||
Ok(resumed_thread.thread_id)
|
||||
Ok((resumed_thread.thread_id, multi_agent_version))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -104,6 +104,10 @@ struct AgentControlHarness {
|
||||
impl AgentControlHarness {
|
||||
async fn new() -> Self {
|
||||
let (home, config) = test_config().await;
|
||||
Self::new_with_config(home, config).await
|
||||
}
|
||||
|
||||
async fn new_with_config(home: TempDir, config: Config) -> Self {
|
||||
let state_db = init_state_db(&config).await;
|
||||
let manager = ThreadManager::with_models_provider_home_and_state_for_tests(
|
||||
CodexAuth::from_api_key("dummy"),
|
||||
@@ -249,6 +253,14 @@ async fn wait_for_live_thread_spawn_children(
|
||||
.expect("expected persisted child tree");
|
||||
}
|
||||
|
||||
async fn assert_thread_not_loaded(manager: &ThreadManager, thread_id: ThreadId) {
|
||||
match manager.get_thread(thread_id).await {
|
||||
Err(CodexErr::ThreadNotFound(id)) => assert_eq!(id, thread_id),
|
||||
Err(err) => panic!("expected ThreadNotFound, got {err:?}"),
|
||||
Ok(_) => panic!("expected thread not to be loaded"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn send_input_errors_when_manager_dropped() {
|
||||
let control = AgentControl::default();
|
||||
@@ -532,22 +544,7 @@ async fn ensure_v2_agent_loaded_reloads_registered_unloaded_agent() {
|
||||
let (home, mut config) = test_config().await;
|
||||
let _ = config.features.enable(Feature::MultiAgentV2);
|
||||
let _ = config.features.enable(Feature::Sqlite);
|
||||
let state_db = init_state_db(&config).await;
|
||||
let manager = ThreadManager::with_models_provider_home_and_state_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()),
|
||||
state_db.clone(),
|
||||
);
|
||||
let control = manager.agent_control();
|
||||
let harness = AgentControlHarness {
|
||||
_home: home,
|
||||
config,
|
||||
state_db,
|
||||
manager,
|
||||
control,
|
||||
};
|
||||
let harness = AgentControlHarness::new_with_config(home, config).await;
|
||||
let (parent_thread_id, _parent_thread) = harness.start_thread().await;
|
||||
let agent_path = AgentPath::try_from("/root/worker").expect("agent path");
|
||||
let spawned_agent = harness
|
||||
@@ -634,6 +631,96 @@ async fn ensure_v2_agent_loaded_reloads_registered_unloaded_agent() {
|
||||
assert_eq!(captured, Some(expected));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn resume_agent_from_rollout_does_not_reopen_v2_descendants() {
|
||||
let (home, mut config) = test_config().await;
|
||||
let _ = config.features.enable(Feature::MultiAgentV2);
|
||||
let _ = config.features.enable(Feature::Sqlite);
|
||||
let harness = AgentControlHarness::new_with_config(home, config).await;
|
||||
let (parent_thread_id, parent_thread) = harness.start_thread().await;
|
||||
let worker_path = AgentPath::root().join("worker").expect("worker path");
|
||||
let worker_thread_id = harness
|
||||
.control
|
||||
.spawn_agent(
|
||||
harness.config.clone(),
|
||||
text_input("hello worker"),
|
||||
Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn {
|
||||
parent_thread_id,
|
||||
depth: 1,
|
||||
agent_path: Some(worker_path.clone()),
|
||||
agent_nickname: None,
|
||||
agent_role: Some("worker".to_string()),
|
||||
})),
|
||||
)
|
||||
.await
|
||||
.expect("worker spawn should succeed");
|
||||
let reviewer_path = worker_path.join("reviewer").expect("reviewer path");
|
||||
let reviewer_thread_id = harness
|
||||
.control
|
||||
.spawn_agent(
|
||||
harness.config.clone(),
|
||||
text_input("hello reviewer"),
|
||||
Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn {
|
||||
parent_thread_id: worker_thread_id,
|
||||
depth: 2,
|
||||
agent_path: Some(reviewer_path.clone()),
|
||||
agent_nickname: None,
|
||||
agent_role: Some("reviewer".to_string()),
|
||||
})),
|
||||
)
|
||||
.await
|
||||
.expect("reviewer spawn should succeed");
|
||||
|
||||
let worker_thread = harness
|
||||
.manager
|
||||
.get_thread(worker_thread_id)
|
||||
.await
|
||||
.expect("worker thread should exist");
|
||||
let reviewer_thread = harness
|
||||
.manager
|
||||
.get_thread(reviewer_thread_id)
|
||||
.await
|
||||
.expect("reviewer thread should exist");
|
||||
persist_thread_for_tree_resume(&parent_thread, "parent persisted").await;
|
||||
persist_thread_for_tree_resume(&worker_thread, "worker persisted").await;
|
||||
persist_thread_for_tree_resume(&reviewer_thread, "reviewer persisted").await;
|
||||
wait_for_live_thread_spawn_children(&harness.control, parent_thread_id, &[worker_thread_id])
|
||||
.await;
|
||||
wait_for_live_thread_spawn_children(&harness.control, worker_thread_id, &[reviewer_thread_id])
|
||||
.await;
|
||||
|
||||
let report = harness
|
||||
.manager
|
||||
.shutdown_all_threads_bounded(Duration::from_secs(5))
|
||||
.await;
|
||||
assert_eq!(report.submit_failed, Vec::<ThreadId>::new());
|
||||
assert_eq!(report.timed_out, Vec::<ThreadId>::new());
|
||||
|
||||
let resumed_manager = ThreadManager::with_models_provider_home_and_state_for_tests(
|
||||
CodexAuth::from_api_key("dummy"),
|
||||
harness.config.model_provider.clone(),
|
||||
harness.config.codex_home.to_path_buf(),
|
||||
std::sync::Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()),
|
||||
harness.state_db.clone(),
|
||||
);
|
||||
let resumed_control = resumed_manager.agent_control();
|
||||
let resumed_parent_thread_id = resumed_control
|
||||
.resume_agent_from_rollout(
|
||||
harness.config.clone(),
|
||||
parent_thread_id,
|
||||
SessionSource::Exec,
|
||||
)
|
||||
.await
|
||||
.expect("v2 root resume should succeed");
|
||||
assert_eq!(resumed_parent_thread_id, parent_thread_id);
|
||||
assert_ne!(
|
||||
resumed_control.get_status(parent_thread_id).await,
|
||||
AgentStatus::NotFound
|
||||
);
|
||||
assert_thread_not_loaded(&resumed_manager, worker_thread_id).await;
|
||||
assert_thread_not_loaded(&resumed_manager, reviewer_thread_id).await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn encrypted_inter_agent_communication_clears_existing_last_task_message() {
|
||||
let harness = AgentControlHarness::new().await;
|
||||
|
||||
Reference in New Issue
Block a user