openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-09-08 15:50:34 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix agent job shutdown cleanup

Browse Source 
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 <noreply@openai.com>
This commit is contained in:
jif-oai
2026-04-22 15:15:19 +01:00
parent 1197683973
commit f576ea4dfe
3 changed files with 109 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files

46
codex-rs/core/src/agent/control.rs
View File

@@ -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<String> {
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
}

21
codex-rs/core/src/agent/control_tests.rs
View File

@@ -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::<ThreadId>::new());
assert_eq!(report.timed_out, Vec::<ThreadId>::new());
assert!(manager.list_thread_ids().await.is_empty());

68
codex-rs/core/src/tools/handlers/agent_jobs_startup.rs
View File

@@ -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();