diff --git a/codex-rs/core/src/agent/control/residency.rs b/codex-rs/core/src/agent/control/residency.rs index 8056fedeb..901f2c93b 100644 --- a/codex-rs/core/src/agent/control/residency.rs +++ b/codex-rs/core/src/agent/control/residency.rs @@ -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 diff --git a/codex-rs/core/src/agent/control/residency_tests.rs b/codex-rs/core/src/agent/control/residency_tests.rs index a0338eda0..88d4b004d 100644 --- a/codex-rs/core/src/agent/control/residency_tests.rs +++ b/codex-rs/core/src/agent/control/residency_tests.rs @@ -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, @@ -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; } diff --git a/codex-rs/core/src/agent/control/spawn.rs b/codex-rs/core/src/agent/control/spawn.rs index a103fbb80..a0d52ea4d 100644 --- a/codex-rs/core/src/agent/control/spawn.rs +++ b/codex-rs/core/src/agent/control/spawn.rs @@ -527,13 +527,16 @@ impl AgentControl { session_source: SessionSource, ) -> CodexResult { 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 { + ) -> 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)) } } diff --git a/codex-rs/core/src/agent/control_tests.rs b/codex-rs/core/src/agent/control_tests.rs index c9cd44492..3ad1360e5 100644 --- a/codex-rs/core/src/agent/control_tests.rs +++ b/codex-rs/core/src/agent/control_tests.rs @@ -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::::new()); + assert_eq!(report.timed_out, Vec::::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;