From f576ea4dfed47c54dadf5f32c81afcfc01256685 Mon Sep 17 00:00:00 2001 From: jif-oai Date: Wed, 22 Apr 2026 15:15:19 +0100 Subject: [PATCH] Fix agent job shutdown cleanup Keep agent job startup waits bounded so stale startup reaping can run while no active workers are present. Keep preserve-mode shutdown threads tracked and accounted until their session loop exits, then remove the thread and release the spawned-agent slot. Co-authored-by: Codex --- codex-rs/core/src/agent/control.rs | 46 +++++++++---- codex-rs/core/src/agent/control_tests.rs | 21 +++--- .../src/tools/handlers/agent_jobs_startup.rs | 68 ++++++++++++++++++- 3 files changed, 109 insertions(+), 26 deletions(-) diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index 6ae9f06e73..aa94a7f54b 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/codex-rs/core/src/agent/control.rs @@ -691,30 +691,50 @@ impl AgentControl { /// Submit a shutdown request for a live agent while keeping the thread tracked until the /// session loop actually terminates. /// - /// This releases the spawned-agent slot immediately so callers can refill concurrency, but - /// it deliberately does not remove the thread from [`ThreadManagerState`]. Keeping the thread - /// tracked ensures later global shutdown paths can still await or abort the underlying session - /// loop instead of leaving an orphaned task alive on the runtime. + /// The spawned-agent slot is released only after a background waiter observes shutdown + /// completion and removes the thread from [`ThreadManagerState`]. Keeping the thread tracked + /// and counted until then ensures later global shutdown paths can still await or abort the + /// underlying session loop instead of leaving an orphaned task alive on the runtime. pub(crate) async fn request_live_agent_shutdown_preserving_thread( &self, agent_id: ThreadId, ) -> CodexResult { let state = self.upgrade()?; - let result = if let Ok(thread) = state.get_thread(agent_id).await { - thread.codex.session.ensure_rollout_materialized().await; - let _ = thread.codex.session.flush_rollout().await; - if matches!(thread.agent_status().await, AgentStatus::Shutdown) { - Ok(String::new()) - } else { - state.send_op(agent_id, Op::Shutdown {}).await + let thread = match state.get_thread(agent_id).await { + Ok(thread) => thread, + Err(_) => { + self.state.release_spawned_thread(agent_id); + return Ok(String::new()); } - } else { + }; + thread.codex.session.ensure_rollout_materialized().await; + let _ = thread.codex.session.flush_rollout().await; + let result = if matches!(thread.agent_status().await, AgentStatus::Shutdown) { Ok(String::new()) + } else { + state.send_op(agent_id, Op::Shutdown {}).await }; if matches!(result, Err(CodexErr::InternalAgentDied)) { let _ = state.remove_thread(&agent_id).await; + self.state.release_spawned_thread(agent_id); + } else if result.is_ok() { + let registry = self.state.clone(); + tokio::spawn(async move { + match thread.shutdown_and_wait().await { + Ok(()) | Err(CodexErr::InternalAgentDied) => { + let _ = state.remove_thread(&agent_id).await; + registry.release_spawned_thread(agent_id); + } + Err(err) => { + warn!( + thread_id = %agent_id, + error = %err, + "failed to wait for live agent shutdown; keeping thread tracked" + ); + } + } + }); } - self.state.release_spawned_thread(agent_id); result } diff --git a/codex-rs/core/src/agent/control_tests.rs b/codex-rs/core/src/agent/control_tests.rs index cf5da5331d..e4d2aad6f6 100644 --- a/codex-rs/core/src/agent/control_tests.rs +++ b/codex-rs/core/src/agent/control_tests.rs @@ -997,7 +997,7 @@ async fn spawn_agent_releases_slot_after_shutdown() { } #[tokio::test] -async fn request_live_agent_shutdown_preserving_thread_releases_slot_without_untracking() { +async fn request_live_agent_shutdown_preserving_thread_cleans_up_after_shutdown() { let max_threads = 1usize; let (_home, config) = test_config_with_cli_overrides(vec![( "agents.max_threads".to_string(), @@ -1008,9 +1008,7 @@ async fn request_live_agent_shutdown_preserving_thread_releases_slot_without_unt CodexAuth::from_api_key("dummy"), config.model_provider.clone(), config.codex_home.clone().to_path_buf(), - std::sync::Arc::new(codex_exec_server::EnvironmentManager::new( - /*exec_server_url*/ None, - )), + std::sync::Arc::new(codex_exec_server::EnvironmentManager::default_for_tests()), ); let control = manager.agent_control(); @@ -1027,10 +1025,13 @@ async fn request_live_agent_shutdown_preserving_thread_releases_slot_without_unt .await .expect("shutdown request should succeed"); - manager - .get_thread(first_agent_id) - .await - .expect("thread should remain tracked after requested shutdown"); + timeout(Duration::from_secs(5), async { + while manager.get_thread(first_agent_id).await.is_ok() { + sleep(Duration::from_millis(25)).await; + } + }) + .await + .expect("thread should be removed after shutdown completes"); let second_agent_id = control .spawn_agent( @@ -1044,9 +1045,7 @@ async fn request_live_agent_shutdown_preserving_thread_releases_slot_without_unt let report = manager .shutdown_all_threads_bounded(Duration::from_secs(10)) .await; - let mut expected_completed = vec![first_agent_id, second_agent_id]; - expected_completed.sort_by_key(std::string::ToString::to_string); - assert_eq!(report.completed, expected_completed); + assert_eq!(report.completed, vec![second_agent_id]); assert_eq!(report.submit_failed, Vec::::new()); assert_eq!(report.timed_out, Vec::::new()); assert!(manager.list_thread_ids().await.is_empty()); diff --git a/codex-rs/core/src/tools/handlers/agent_jobs_startup.rs b/codex-rs/core/src/tools/handlers/agent_jobs_startup.rs index 905617a602..8aa3491d62 100644 --- a/codex-rs/core/src/tools/handlers/agent_jobs_startup.rs +++ b/codex-rs/core/src/tools/handlers/agent_jobs_startup.rs @@ -213,7 +213,12 @@ pub(super) async fn wait_for_startup_or_status_change( let active_items_ref = &*active_items; if active_items_ref.is_empty() { - if let Some(result) = startup_tasks.starting_items.join_next_with_id().await { + if let Ok(Some(result)) = timeout( + STATUS_POLL_INTERVAL, + startup_tasks.starting_items.join_next_with_id(), + ) + .await + { let starting_items_len = startup_tasks.starting_items.len(); handle_worker_startup_result( session, @@ -523,6 +528,65 @@ mod tests { assert_eq!(outputs, vec!["item-0", "item-1", "item-2"]); } + #[tokio::test] + async fn wait_for_startup_or_status_change_returns_when_only_startups_are_pending() + -> anyhow::Result<()> { + let home = TempDir::new()?; + let config = ConfigBuilder::without_managed_config_for_tests() + .codex_home(home.path().to_path_buf()) + .build() + .await?; + + let manager = ThreadManager::with_models_provider_and_home_for_tests( + CodexAuth::from_api_key("dummy"), + config.model_provider.clone(), + config.codex_home.clone().to_path_buf(), + Arc::new(EnvironmentManager::default_for_tests()), + ); + let root = manager.start_thread(config.clone()).await?; + let session = root.thread.codex.session.clone(); + let db = codex_state::StateRuntime::init( + config.codex_home.clone().to_path_buf(), + "test-provider".to_string(), + ) + .await?; + + let mut startup_tasks = StartupTasks::default(); + spawn_tracked_startup_task( + &mut startup_tasks, + "item-1".to_string(), + Instant::now(), + std::future::pending(), + ); + + let mut active_items = HashMap::new(); + timeout( + Duration::from_secs(1), + wait_for_startup_or_status_change( + session, + db, + "job-1", + &mut active_items, + &mut startup_tasks, + ), + ) + .await + .expect("wait should return so stale startup reaping can run")?; + + assert!(active_items.is_empty()); + assert_eq!(startup_tasks.len(), 1); + assert_eq!(startup_tasks.launching_items.len(), 1); + + let aborted = abort_all_startups(&mut startup_tasks).await; + assert_eq!(aborted, 1); + let report = manager + .shutdown_all_threads_bounded(Duration::from_secs(10)) + .await; + assert_eq!(report.submit_failed, Vec::new()); + assert_eq!(report.timed_out, Vec::new()); + Ok(()) + } + #[tokio::test] async fn drain_ready_startups_reports_agent_limit_and_requeues_item() -> anyhow::Result<()> { let home = TempDir::new()?; @@ -535,7 +599,7 @@ mod tests { CodexAuth::from_api_key("dummy"), config.model_provider.clone(), config.codex_home.clone().to_path_buf(), - Arc::new(EnvironmentManager::new(/*exec_server_url*/ None)), + Arc::new(EnvironmentManager::default_for_tests()), ); let root = manager.start_thread(config.clone()).await?; let session = root.thread.codex.session.clone();