From fa83369441b59e203f6016956f89d4968d6f9c57 Mon Sep 17 00:00:00 2001 From: Dave Aitel Date: Thu, 9 Apr 2026 21:27:57 -0400 Subject: [PATCH] Finish agent-job fanout completion and exec shutdown --- codex-rs/core/src/agent/control.rs | 30 +++ codex-rs/core/src/agent/control_tests.rs | 56 +++++ .../core/src/tools/handlers/agent_jobs.rs | 68 +++++- .../agent_jobs/report_agent_job_result.rs | 41 ++-- .../src/tools/handlers/agent_jobs_slots.rs | 47 +++- .../src/tools/handlers/agent_jobs_startup.rs | 20 +- codex-rs/core/tests/suite/agent_jobs.rs | 214 +++++++++++++++++- codex-rs/exec/src/lib.rs | 57 ++++- codex-rs/exec/src/lib_tests.rs | 43 ++++ codex-rs/state/src/runtime/agent_jobs.rs | 101 +++++++++ 10 files changed, 625 insertions(+), 52 deletions(-) diff --git a/codex-rs/core/src/agent/control.rs b/codex-rs/core/src/agent/control.rs index 42981a5d0a..32795b43c1 100644 --- a/codex-rs/core/src/agent/control.rs +++ b/codex-rs/core/src/agent/control.rs @@ -780,6 +780,36 @@ impl AgentControl { result } + /// 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. + 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; + 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 + } + } else { + Ok(String::new()) + }; + if matches!(result, Err(CodexErr::InternalAgentDied)) { + let _ = state.remove_thread(&agent_id).await; + } + self.state.release_spawned_thread(agent_id); + result + } + /// Mark `agent_id` as explicitly closed in persisted spawn-edge state, then shut down the /// agent and any live descendants reached from the in-memory tree. pub(crate) async fn close_agent(&self, agent_id: ThreadId) -> CodexResult { diff --git a/codex-rs/core/src/agent/control_tests.rs b/codex-rs/core/src/agent/control_tests.rs index 7ff50ffef8..9d420e9018 100644 --- a/codex-rs/core/src/agent/control_tests.rs +++ b/codex-rs/core/src/agent/control_tests.rs @@ -1305,6 +1305,62 @@ async fn spawn_agent_releases_slot_after_shutdown() { .expect("shutdown agent"); } +#[tokio::test] +async fn request_live_agent_shutdown_preserving_thread_releases_slot_without_untracking() { + let max_threads = 1usize; + let (_home, config) = test_config_with_cli_overrides(vec![( + "agents.max_threads".to_string(), + TomlValue::Integer(max_threads as i64), + )]) + .await; + let manager = ThreadManager::with_models_provider_and_home_for_tests( + CodexAuth::from_api_key("dummy"), + config.model_provider.clone(), + config.codex_home.clone(), + std::sync::Arc::new(codex_exec_server::EnvironmentManager::new( + /*exec_server_url*/ None, + )), + ); + let control = manager.agent_control(); + + let first_agent_id = control + .spawn_agent( + config.clone(), + text_input("hello"), + /*session_source*/ None, + ) + .await + .expect("spawn_agent should succeed"); + let _ = control + .request_live_agent_shutdown_preserving_thread(first_agent_id) + .await + .expect("shutdown request should succeed"); + + manager + .get_thread(first_agent_id) + .await + .expect("thread should remain tracked after requested shutdown"); + + let second_agent_id = control + .spawn_agent( + config.clone(), + text_input("hello again"), + /*session_source*/ None, + ) + .await + .expect("spawn_agent should succeed after preserving-thread shutdown"); + + 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.submit_failed, Vec::::new()); + assert_eq!(report.timed_out, Vec::::new()); + assert!(manager.list_thread_ids().await.is_empty()); +} + #[tokio::test] async fn spawn_agent_limit_shared_across_clones() { let max_threads = 1usize; diff --git a/codex-rs/core/src/tools/handlers/agent_jobs.rs b/codex-rs/core/src/tools/handlers/agent_jobs.rs index ff196c6f67..7e304bd813 100644 --- a/codex-rs/core/src/tools/handlers/agent_jobs.rs +++ b/codex-rs/core/src/tools/handlers/agent_jobs.rs @@ -104,6 +104,18 @@ struct ActiveJobItem { status_rx: Option>, } +fn request_live_agent_shutdown( + agent_control: crate::agent::control::AgentControl, + thread_id: ThreadId, +) { + tokio::spawn(async move { + tokio::task::yield_now().await; + let _ = agent_control + .request_live_agent_shutdown_preserving_thread(thread_id) + .await; + }); +} + fn required_state_db( session: &Arc, ) -> Result, FunctionCallError> { @@ -224,6 +236,27 @@ async fn run_agent_job_loop( progressed = true; } + let terminal_in_db = if cancel_requested { + scheduler_progress.running_items == 0 + } else { + scheduler_progress.pending_items == 0 && scheduler_progress.running_items == 0 + }; + if terminal_in_db + && slots::reconcile_terminal_scheduler_state( + session.clone(), + job_id.as_str(), + &scheduler_progress, + &mut active_items, + &mut starting_items, + ) + .await? + { + progressed = true; + } + if terminal_in_db && active_items.is_empty() && starting_items.is_empty() { + break; + } + if !cancel_requested && active_items.len() + starting_items.len() < options.max_concurrency && startup::launch_pending_items( @@ -234,6 +267,7 @@ async fn run_agent_job_loop( &options, startup::SchedulerOccupancy { active_items: active_items.len(), + db_pending_items: scheduler_progress.pending_items, db_running_items: scheduler_progress.running_items, }, &mut starting_items, @@ -280,19 +314,29 @@ async fn run_agent_job_loop( let finished = find_finished_threads(session.clone(), &active_items).await; if finished.is_empty() { - let progress = db.get_agent_job_progress(job_id.as_str()).await?; - if cancel_requested { - if progress.running_items == 0 - && active_items.is_empty() - && starting_items.is_empty() - { - break; - } - } else if progress.pending_items == 0 - && progress.running_items == 0 - && active_items.is_empty() - && starting_items.is_empty() + let progress = if progressed { + db.get_agent_job_progress(job_id.as_str()).await? + } else { + scheduler_progress + }; + let terminal_in_db = if cancel_requested { + progress.running_items == 0 + } else { + progress.pending_items == 0 && progress.running_items == 0 + }; + if terminal_in_db + && slots::reconcile_terminal_scheduler_state( + session.clone(), + job_id.as_str(), + &progress, + &mut active_items, + &mut starting_items, + ) + .await? { + progressed = true; + } + if terminal_in_db && active_items.is_empty() && starting_items.is_empty() { break; } if !progressed { diff --git a/codex-rs/core/src/tools/handlers/agent_jobs/report_agent_job_result.rs b/codex-rs/core/src/tools/handlers/agent_jobs/report_agent_job_result.rs index aa0836a18d..8e92d18981 100644 --- a/codex-rs/core/src/tools/handlers/agent_jobs/report_agent_job_result.rs +++ b/codex-rs/core/src/tools/handlers/agent_jobs/report_agent_job_result.rs @@ -63,39 +63,42 @@ pub async fn handle( let db = required_state_db(&session)?; let reporting_thread_id = session.conversation_id; let reporting_thread_id_str = reporting_thread_id.to_string(); - let accepted = db - .report_agent_job_item_result( + let accepted = if args.stop.unwrap_or(false) { + db.report_agent_job_item_result_and_cancel_job( + args.job_id.as_str(), + args.item_id.as_str(), + reporting_thread_id_str.as_str(), + &args.result, + "cancelled by worker request", + ) + .await + } else { + db.report_agent_job_item_result( args.job_id.as_str(), args.item_id.as_str(), reporting_thread_id_str.as_str(), &args.result, ) .await - .map_err(|err| { - let job_id = args.job_id.as_str(); - let item_id = args.item_id.as_str(); - FunctionCallError::RespondToModel(format!( - "failed to record agent job result for {job_id} / {item_id}: {err}" - )) - })?; - if accepted && args.stop.unwrap_or(false) { - let message = "cancelled by worker request"; - let _ = db - .mark_agent_job_cancelled(args.job_id.as_str(), message) - .await; } + .map_err(|err| { + let job_id = args.job_id.as_str(); + let item_id = args.item_id.as_str(); + FunctionCallError::RespondToModel(format!( + "failed to record agent job result for {job_id} / {item_id}: {err}" + )) + })?; if accepted { - let agent_control = session.services.agent_control.clone(); tracing::debug!( job_id = args.job_id, item_id = args.item_id, thread_id = %reporting_thread_id, "agent job accepted worker result; scheduling worker shutdown" ); - tokio::spawn(async move { - tokio::task::yield_now().await; - let _ = agent_control.shutdown_live_agent(reporting_thread_id).await; - }); + request_live_agent_shutdown( + session.services.agent_control.clone(), + reporting_thread_id, + ); } let content = serde_json::to_string(&ReportAgentJobResultToolResult { accepted }).map_err(|err| { diff --git a/codex-rs/core/src/tools/handlers/agent_jobs_slots.rs b/codex-rs/core/src/tools/handlers/agent_jobs_slots.rs index 1552d6c473..52cc70645a 100644 --- a/codex-rs/core/src/tools/handlers/agent_jobs_slots.rs +++ b/codex-rs/core/src/tools/handlers/agent_jobs_slots.rs @@ -54,12 +54,11 @@ pub(super) async fn reclaim_inactive_active_items( ); for inactive_item in inactive_items { - let _ = session - .services - .agent_control - .shutdown_live_agent(inactive_item.thread_id) - .await; active_items.remove(&inactive_item.thread_id); + request_live_agent_shutdown( + session.services.agent_control.clone(), + inactive_item.thread_id, + ); tracing::debug!( job_id, item_id = inactive_item.item_id, @@ -77,3 +76,41 @@ pub(super) async fn reclaim_inactive_active_items( Ok(true) } + +pub(super) async fn reconcile_terminal_scheduler_state( + session: Arc, + job_id: &str, + progress: &codex_state::AgentJobProgress, + active_items: &mut HashMap, + startup_tasks: &mut startup::StartupTasks, +) -> anyhow::Result { + if active_items.is_empty() && startup_tasks.is_empty() { + return Ok(false); + } + + let active_count = active_items.len(); + let starting_count = startup_tasks.len(); + tracing::info!( + job_id, + pending_items = progress.pending_items, + db_running_items = progress.running_items, + active_items = active_count, + starting_items = starting_count, + "agent job state is terminal in DB; forcing scheduler teardown" + ); + + let thread_ids: Vec<_> = active_items.keys().copied().collect(); + for thread_id in thread_ids { + active_items.remove(&thread_id); + request_live_agent_shutdown(session.services.agent_control.clone(), thread_id); + } + + let aborted_startups = startup::abort_all_startups(startup_tasks).await; + tracing::debug!( + job_id, + active_items_reclaimed = active_count, + starting_items_aborted = aborted_startups, + "agent job terminal scheduler teardown completed" + ); + Ok(active_count > 0 || aborted_startups > 0) +} 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 1fb1445232..944ab474e1 100644 --- a/codex-rs/core/src/tools/handlers/agent_jobs_startup.rs +++ b/codex-rs/core/src/tools/handlers/agent_jobs_startup.rs @@ -30,6 +30,7 @@ pub(super) struct StartupTasks { #[derive(Clone, Copy, Debug)] pub(super) struct SchedulerOccupancy { pub(super) active_items: usize, + pub(super) db_pending_items: usize, pub(super) db_running_items: usize, } @@ -152,6 +153,7 @@ pub(super) async fn launch_pending_items( tracing::info!( job_id, launched, + db_pending_items = occupancy.db_pending_items, db_running_items = occupancy.db_running_items, active_items = occupancy.active_items, starting_items = startup_tasks.len(), @@ -238,6 +240,22 @@ pub(super) async fn wait_for_startup_or_status_change( Ok(()) } +pub(super) async fn abort_all_startups(startup_tasks: &mut StartupTasks) -> usize { + let startup_count = startup_tasks.starting_items.len(); + if startup_count == 0 { + startup_tasks.launching_items.clear(); + return 0; + } + + for launching_item in startup_tasks.launching_items.values() { + launching_item.abort_handle.abort(); + } + startup_tasks.launching_items.clear(); + + while startup_tasks.starting_items.join_next().await.is_some() {} + startup_count +} + pub(super) async fn reap_stale_startups( db: Arc, job_id: &str, @@ -300,7 +318,7 @@ async fn handle_worker_startup_result( let _ = session .services .agent_control - .shutdown_live_agent(thread_id) + .request_live_agent_shutdown_preserving_thread(thread_id) .await; tracing::debug!( job_id, diff --git a/codex-rs/core/tests/suite/agent_jobs.rs b/codex-rs/core/tests/suite/agent_jobs.rs index 04008b4025..02ff18b508 100644 --- a/codex-rs/core/tests/suite/agent_jobs.rs +++ b/codex-rs/core/tests/suite/agent_jobs.rs @@ -1,5 +1,6 @@ use anyhow::Result; use codex_features::Feature; +use core_test_support::responses::ev_assistant_message; use core_test_support::responses::ev_completed; use core_test_support::responses::ev_function_call; use core_test_support::responses::ev_response_created; @@ -61,19 +62,96 @@ struct DelayedWorkerAfterReportResponder { seen_main: AtomicBool, worker_calls: Arc, delayed_worker_output: AtomicBool, + worker_output_delay: Duration, +} + +struct WorkerNeverCompletesAfterReportResponder { + spawn_args_json: String, + seen_main: AtomicBool, + worker_calls: Arc, } impl DelayedWorkerAfterReportResponder { - fn new(spawn_args_json: String, worker_calls: Arc) -> Self { + fn new( + spawn_args_json: String, + worker_calls: Arc, + worker_output_delay: Duration, + ) -> Self { Self { spawn_args_json, seen_main: AtomicBool::new(false), worker_calls, delayed_worker_output: AtomicBool::new(false), + worker_output_delay, } } } +impl WorkerNeverCompletesAfterReportResponder { + fn new(spawn_args_json: String, worker_calls: Arc) -> Self { + Self { + spawn_args_json, + seen_main: AtomicBool::new(false), + worker_calls, + } + } +} + +impl Respond for WorkerNeverCompletesAfterReportResponder { + fn respond(&self, request: &wiremock::Request) -> ResponseTemplate { + let body_bytes = decode_body_bytes(request); + let body: Value = serde_json::from_slice(&body_bytes).unwrap_or(Value::Null); + + let call_output_ids = function_call_output_call_ids(&body); + if !call_output_ids.is_empty() { + if call_output_ids + .iter() + .any(|call_id| call_id.starts_with("call-worker-")) + { + return sse_response(sse(vec![ + ev_response_created("resp-worker-post-report"), + ev_assistant_message("msg-worker-post-report", "no bugs found"), + ])); + } + return sse_response(sse(vec![ + ev_response_created("resp-tool"), + ev_completed("resp-tool"), + ])); + } + + if let Some((job_id, item_id)) = extract_job_and_item(&body) { + let call_index = self.worker_calls.fetch_add(1, Ordering::SeqCst); + let call_id = format!("call-worker-{call_index}"); + let args = json!({ + "job_id": job_id, + "item_id": item_id, + "result": { "item_id": item_id } + }); + let args_json = serde_json::to_string(&args).unwrap_or_else(|err| { + panic!("worker args serialize: {err}"); + }); + return sse_response(sse(vec![ + ev_response_created("resp-worker"), + ev_function_call(&call_id, "report_agent_job_result", &args_json), + ev_completed("resp-worker"), + ])); + } + + if !self.seen_main.swap(true, Ordering::SeqCst) { + return sse_response(sse(vec![ + ev_response_created("resp-main"), + ev_function_call("call-spawn", "spawn_agents_on_csv", &self.spawn_args_json), + ev_completed("resp-main"), + ])); + } + + sse_response(sse(vec![ + ev_response_created("resp-default"), + ev_completed("resp-default"), + ])) + } +} + impl Respond for DelayedWorkerAfterReportResponder { fn respond(&self, request: &wiremock::Request) -> ResponseTemplate { let body_bytes = decode_body_bytes(request); @@ -90,7 +168,7 @@ impl Respond for DelayedWorkerAfterReportResponder { .any(|call_id| call_id == "call-worker-0") && !self.delayed_worker_output.swap(true, Ordering::SeqCst) { - return response.set_delay(Duration::from_secs(2)); + return response.set_delay(self.worker_output_delay); } return response; } @@ -561,7 +639,11 @@ async fn spawn_agents_on_csv_reclaims_slot_after_report_before_worker_completes( let args_json = serde_json::to_string(&args)?; let worker_calls = Arc::new(AtomicUsize::new(0)); - let responder = DelayedWorkerAfterReportResponder::new(args_json, worker_calls.clone()); + let responder = DelayedWorkerAfterReportResponder::new( + args_json, + worker_calls.clone(), + Duration::from_secs(2), + ); Mock::given(method("POST")) .and(path_regex(".*/responses$")) .respond_with(responder) @@ -604,3 +686,129 @@ async fn spawn_agents_on_csv_reclaims_slot_after_report_before_worker_completes( Ok(()) } + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn spawn_agents_on_csv_finishes_after_rows_complete_even_if_worker_exit_lags() -> Result<()> { + let server = start_mock_server().await; + let mut builder = test_codex().with_config(|config| { + config + .features + .enable(Feature::SpawnCsv) + .expect("test config should allow feature update"); + config + .features + .enable(Feature::Sqlite) + .expect("test config should allow feature update"); + }); + let test = builder.build(&server).await?; + + let input_path = test.cwd_path().join("agent_jobs_terminal.csv"); + let output_path = test.cwd_path().join("agent_jobs_terminal_out.csv"); + fs::write(&input_path, "path\nfile-1\n")?; + + let args = json!({ + "csv_path": input_path.display().to_string(), + "instruction": "Return {path}", + "output_csv_path": output_path.display().to_string(), + "max_concurrency": 1, + }); + let args_json = serde_json::to_string(&args)?; + + let worker_calls = Arc::new(AtomicUsize::new(0)); + let responder = DelayedWorkerAfterReportResponder::new( + args_json, + worker_calls.clone(), + Duration::from_secs(30), + ); + Mock::given(method("POST")) + .and(path_regex(".*/responses$")) + .respond_with(responder) + .mount(&server) + .await; + + timeout(Duration::from_secs(5), test.submit_turn("run batch job")) + .await + .expect("root turn should finalize without waiting for the delayed worker exit")?; + + let output = fs::read_to_string(&output_path)?; + let rows: Vec<&str> = output.lines().skip(1).collect(); + assert_eq!(rows.len(), 1); + + let job_id = rows + .first() + .and_then(|line| { + parse_simple_csv_line(line) + .iter() + .find(|value| value.len() == 36) + .cloned() + }) + .expect("job_id from csv"); + let db = test.codex.state_db().expect("state db"); + let job = db.get_agent_job(job_id.as_str()).await?.expect("job"); + assert_eq!(job.status, codex_state::AgentJobStatus::Completed); + assert!(job.completed_at.is_some()); + assert_eq!(worker_calls.load(Ordering::SeqCst), 1); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn spawn_agents_on_csv_finishes_when_worker_reports_but_never_completes_turn() -> Result<()> { + let server = start_mock_server().await; + let mut builder = test_codex().with_config(|config| { + config + .features + .enable(Feature::SpawnCsv) + .expect("test config should allow feature update"); + config + .features + .enable(Feature::Sqlite) + .expect("test config should allow feature update"); + }); + let test = builder.build(&server).await?; + + let input_path = test.cwd_path().join("agent_jobs_incomplete_turn.csv"); + let output_path = test.cwd_path().join("agent_jobs_incomplete_turn_out.csv"); + fs::write(&input_path, "path\nfile-1\n")?; + + let args = json!({ + "csv_path": input_path.display().to_string(), + "instruction": "Return {path}", + "output_csv_path": output_path.display().to_string(), + "max_concurrency": 1, + }); + let args_json = serde_json::to_string(&args)?; + + let worker_calls = Arc::new(AtomicUsize::new(0)); + let responder = WorkerNeverCompletesAfterReportResponder::new(args_json, worker_calls.clone()); + Mock::given(method("POST")) + .and(path_regex(".*/responses$")) + .respond_with(responder) + .mount(&server) + .await; + + timeout(Duration::from_secs(5), test.submit_turn("run batch job")) + .await + .expect("root turn should finalize even if a worker never reaches task_complete")?; + + let output = fs::read_to_string(&output_path)?; + let rows: Vec<&str> = output.lines().skip(1).collect(); + assert_eq!(rows.len(), 1); + + let job_id = rows + .first() + .and_then(|line| { + parse_simple_csv_line(line) + .iter() + .find(|value| value.len() == 36) + .cloned() + }) + .expect("job_id from csv"); + let db = test.codex.state_db().expect("state db"); + let job = db.get_agent_job(job_id.as_str()).await?.expect("job"); + assert_eq!(job.status, codex_state::AgentJobStatus::Completed); + assert!(job.completed_at.is_some()); + assert_eq!(worker_calls.load(Ordering::SeqCst), 1); + + Ok(()) +} diff --git a/codex-rs/exec/src/lib.rs b/codex-rs/exec/src/lib.rs index 3db7a51576..84d7fb6ead 100644 --- a/codex-rs/exec/src/lib.rs +++ b/codex-rs/exec/src/lib.rs @@ -135,12 +135,15 @@ pub use exec_events::Usage; pub use exec_events::WebSearchItem; use serde_json::Value; use std::collections::HashMap; +use std::future::Future; use std::io::IsTerminal; use std::io::Read; use std::path::Path; use std::path::PathBuf; +use std::time::Duration; use supports_color::Stream; use tokio::sync::mpsc; +use tokio::time::timeout; use tracing::Instrument; use tracing::error; use tracing::field; @@ -156,6 +159,8 @@ use crate::event_processor::EventProcessor; const DEFAULT_ANALYTICS_ENABLED: bool = true; const EXEC_DEFAULT_LOG_FILTER: &str = "error,opentelemetry_sdk=off,opentelemetry_otlp=off"; +const THREAD_UNSUBSCRIBE_TIMEOUT: Duration = Duration::from_secs(1); +const TURN_COMPLETED_BACKFILL_TIMEOUT: Duration = Duration::from_secs(2); enum InitialOperation { UserTurn { @@ -1251,16 +1256,23 @@ async fn maybe_backfill_turn_completed_items( return; }; - let response = send_request_with_response::( - client, - ClientRequest::ThreadRead { - request_id: request_ids.next(), - params: ThreadReadParams { - thread_id: payload.thread_id.clone(), - include_turns: true, - }, - }, + // This runs inline on exec's event loop immediately after `TurnCompleted`. + // Bound the request so a backpressured in-process event queue cannot deadlock + // shutdown by blocking `thread/read` forever behind unrelated lossless events. + let response = await_request_with_timeout( "thread/read", + TURN_COMPLETED_BACKFILL_TIMEOUT, + send_request_with_response::( + client, + ClientRequest::ThreadRead { + request_id: request_ids.next(), + params: ThreadReadParams { + thread_id: payload.thread_id.clone(), + include_turns: true, + }, + }, + "thread/read", + ), ) .await; @@ -1489,9 +1501,30 @@ async fn request_shutdown( thread_id: thread_id.to_string(), }, }; - send_request_with_response::(client, request, "thread/unsubscribe") - .await - .map(|_| ()) + await_request_with_timeout( + "thread/unsubscribe", + THREAD_UNSUBSCRIBE_TIMEOUT, + send_request_with_response::( + client, + request, + "thread/unsubscribe", + ), + ) + .await + .map(|_| ()) +} + +async fn await_request_with_timeout( + request_name: &str, + timeout_duration: Duration, + request: impl Future>, +) -> Result { + match timeout(timeout_duration, request).await { + Ok(result) => result, + Err(_) => Err(format!( + "{request_name} timed out after {timeout_duration:?}" + )), + } } async fn resolve_server_request( diff --git a/codex-rs/exec/src/lib_tests.rs b/codex-rs/exec/src/lib_tests.rs index 79d07a0b76..66ceac127a 100644 --- a/codex-rs/exec/src/lib_tests.rs +++ b/codex-rs/exec/src/lib_tests.rs @@ -15,6 +15,7 @@ use std::io::Write; use std::sync::Arc; use std::sync::Mutex; use tempfile::tempdir; +use tokio::time::sleep; use tracing_opentelemetry::OpenTelemetrySpanExt; fn test_tracing_subscriber() -> impl tracing::Subscriber + Send + Sync { @@ -626,3 +627,45 @@ fn sample_thread_start_response() -> ThreadStartResponse { reasoning_effort: None, } } + +#[tokio::test] +async fn await_request_with_timeout_returns_ready_result() { + let result = await_request_with_timeout( + "test/request", + std::time::Duration::from_millis(50), + async { Ok::<_, String>(123usize) }, + ) + .await + .expect("ready request should succeed"); + + assert_eq!(result, 123); +} + +#[tokio::test] +async fn await_request_with_timeout_errors_when_request_stalls() { + let err = await_request_with_timeout( + "test/request", + std::time::Duration::from_millis(10), + async { + sleep(std::time::Duration::from_millis(50)).await; + Ok::<_, String>(()) + }, + ) + .await + .expect_err("stalled request should time out"); + + assert_eq!(err, "test/request timed out after 10ms"); +} + +#[tokio::test] +async fn turn_completed_backfill_times_out_instead_of_hanging_exec() { + let err = + await_request_with_timeout("thread/read", std::time::Duration::from_millis(10), async { + sleep(std::time::Duration::from_millis(50)).await; + Ok::<_, String>(()) + }) + .await + .expect_err("stalled turn/read backfill should time out"); + + assert_eq!(err, "thread/read timed out after 10ms"); +} diff --git a/codex-rs/state/src/runtime/agent_jobs.rs b/codex-rs/state/src/runtime/agent_jobs.rs index bbc7dc4d1e..de7d152274 100644 --- a/codex-rs/state/src/runtime/agent_jobs.rs +++ b/codex-rs/state/src/runtime/agent_jobs.rs @@ -462,6 +462,68 @@ WHERE Ok(result.rows_affected() > 0) } + pub async fn report_agent_job_item_result_and_cancel_job( + &self, + job_id: &str, + item_id: &str, + reporting_thread_id: &str, + result_json: &Value, + cancel_reason: &str, + ) -> anyhow::Result { + let now = Utc::now().timestamp(); + let serialized = serde_json::to_string(result_json)?; + let mut tx = self.pool.begin().await?; + let result = sqlx::query( + r#" +UPDATE agent_job_items +SET + status = ?, + result_json = ?, + reported_at = ?, + completed_at = ?, + updated_at = ?, + last_error = NULL, + assigned_thread_id = NULL +WHERE + job_id = ? + AND item_id = ? + AND status = ? + AND assigned_thread_id = ? + "#, + ) + .bind(AgentJobItemStatus::Completed.as_str()) + .bind(serialized) + .bind(now) + .bind(now) + .bind(now) + .bind(job_id) + .bind(item_id) + .bind(AgentJobItemStatus::Running.as_str()) + .bind(reporting_thread_id) + .execute(&mut *tx) + .await?; + if result.rows_affected() > 0 { + sqlx::query( + r#" +UPDATE agent_jobs +SET status = ?, updated_at = ?, completed_at = ?, last_error = ? +WHERE id = ? AND status IN (?, ?) + "#, + ) + .bind(AgentJobStatus::Cancelled.as_str()) + .bind(now) + .bind(now) + .bind(cancel_reason) + .bind(job_id) + .bind(AgentJobStatus::Pending.as_str()) + .bind(AgentJobStatus::Running.as_str()) + .execute(&mut *tx) + .await?; + } + tx.commit().await?; + Ok(result.rows_affected() > 0) + } + pub async fn mark_agent_job_item_completed( &self, job_id: &str, @@ -681,6 +743,45 @@ mod tests { Ok(()) } + #[tokio::test] + async fn report_agent_job_item_result_and_cancel_job_is_atomic() -> anyhow::Result<()> { + let codex_home = unique_temp_dir(); + let runtime = StateRuntime::init(codex_home, "test-provider".to_string()).await?; + let (job_id, item_id, thread_id) = create_running_single_item_job(runtime.as_ref()).await?; + + let accepted = runtime + .report_agent_job_item_result_and_cancel_job( + job_id.as_str(), + item_id.as_str(), + thread_id.as_str(), + &json!({"ok": true}), + "cancelled by worker request", + ) + .await?; + assert!(accepted); + + let item = runtime + .get_agent_job_item(job_id.as_str(), item_id.as_str()) + .await? + .expect("job item should exist"); + assert_eq!(item.status, AgentJobItemStatus::Completed); + assert_eq!(item.result_json, Some(json!({"ok": true}))); + assert_eq!(item.assigned_thread_id, None); + + let job = runtime + .get_agent_job(job_id.as_str()) + .await? + .expect("job should exist"); + assert_eq!(job.status, AgentJobStatus::Cancelled); + assert_eq!( + job.last_error, + Some("cancelled by worker request".to_string()) + ); + assert!(job.completed_at.is_some()); + + Ok(()) + } + #[tokio::test] async fn set_agent_job_item_thread_increments_attempt_count_after_claim() -> anyhow::Result<()> {