From b4e16db47b9a39e00406f61ac677f79133f7e08b Mon Sep 17 00:00:00 2001 From: Dave Aitel Date: Fri, 17 Apr 2026 08:26:00 -0400 Subject: [PATCH] Reduce SQLite lock contention in agent jobs High-fanout agent-job runs were opening many StateRuntime instances against the same state_5.sqlite during subagent and session startup. That amplified SQLite writer contention and caused spawn_agents_on_csv to fail immediately with database-locked errors before workers could run. Fix this by caching rollout StateRuntime handles per (sqlite_home, provider) so subagents reuse one SQLite runtime instead of opening a new pool. Also retry transient SQLite lock errors around agent-job DB operations, replace per-thread reclaim lookups with one running-items query, and map cancelled background tool tasks to normal tool failures instead of panicking the turn drain path. Add rollout tests covering cache reuse and concurrent initialization. --- .../core/src/tools/handlers/agent_jobs.rs | 297 +++++++++++------- .../core/src/tools/handlers/agent_jobs_db.rs | 95 ++++++ .../src/tools/handlers/agent_jobs_slots.rs | 35 ++- .../src/tools/handlers/agent_jobs_startup.rs | 100 +++--- codex-rs/core/src/tools/parallel.rs | 87 ++++- codex-rs/rollout/src/state_db.rs | 73 ++++- codex-rs/rollout/src/state_db_tests.rs | 83 +++++ 7 files changed, 601 insertions(+), 169 deletions(-) create mode 100644 codex-rs/core/src/tools/handlers/agent_jobs_db.rs diff --git a/codex-rs/core/src/tools/handlers/agent_jobs.rs b/codex-rs/core/src/tools/handlers/agent_jobs.rs index dad74b532b..89b6665c67 100644 --- a/codex-rs/core/src/tools/handlers/agent_jobs.rs +++ b/codex-rs/core/src/tools/handlers/agent_jobs.rs @@ -193,6 +193,8 @@ impl JobProgressEmitter { } } +#[path = "agent_jobs_db.rs"] +mod db_ops; #[path = "agent_jobs_slots.rs"] mod slots; #[path = "agent_jobs_startup.rs"] @@ -339,44 +341,50 @@ mod spawn_agents_on_csv { args.max_runtime_seconds .or(turn.config.agent_job_max_runtime_seconds), )?; - let _job = db - .create_agent_job( + let _job = db_ops::retry_locked("create_agent_job", || async { + db.create_agent_job( &codex_state::AgentJobCreateParams { id: job_id.clone(), - name: job_name, - instruction: args.instruction, + name: job_name.clone(), + instruction: args.instruction.clone(), auto_export: true, max_runtime_seconds, - output_schema_json: args.output_schema, - input_headers: headers, + output_schema_json: args.output_schema.clone(), + input_headers: headers.clone(), input_csv_path: input_path.display().to_string(), output_csv_path: output_csv_path.display().to_string(), }, items.as_slice(), ) .await - .map_err(|err| { - FunctionCallError::RespondToModel(format!("failed to create agent job: {err}")) - })?; + }) + .await + .map_err(|err| { + FunctionCallError::RespondToModel(format!("failed to create agent job: {err}")) + })?; let requested_concurrency = args.max_concurrency.or(args.max_workers); let options = match build_runner_options(&session, &turn, requested_concurrency).await { Ok(options) => options, Err(err) => { let error_message = err.to_string(); - let _ = db - .mark_agent_job_failed(job_id.as_str(), error_message.as_str()) - .await; + let _ = db_ops::retry_locked("mark_agent_job_failed_after_options", || async { + db.mark_agent_job_failed(job_id.as_str(), error_message.as_str()) + .await + }) + .await; return Err(err); } }; - db.mark_agent_job_running(job_id.as_str()) - .await - .map_err(|err| { - FunctionCallError::RespondToModel(format!( - "failed to transition agent job {job_id} to running: {err}" - )) - })?; + db_ops::retry_locked("mark_agent_job_running", || async { + db.mark_agent_job_running(job_id.as_str()).await + }) + .await + .map_err(|err| { + FunctionCallError::RespondToModel(format!( + "failed to transition agent job {job_id} to running: {err}" + )) + })?; let max_threads = turn.config.agent_max_threads; let effective_concurrency = options.max_concurrency; let message = format!( @@ -393,25 +401,26 @@ mod spawn_agents_on_csv { .await { let error_message = format!("job runner failed: {err}"); - let _ = db - .mark_agent_job_failed(job_id.as_str(), error_message.as_str()) - .await; + let _ = db_ops::retry_locked("mark_agent_job_failed_after_runner_error", || async { + db.mark_agent_job_failed(job_id.as_str(), error_message.as_str()) + .await + }) + .await; return Err(FunctionCallError::RespondToModel(format!( "agent job {job_id} failed: {err}" ))); } - let job = db - .get_agent_job(job_id.as_str()) - .await - .map_err(|err| { - FunctionCallError::RespondToModel(format!( - "failed to load agent job {job_id}: {err}" - )) - })? - .ok_or_else(|| { - FunctionCallError::RespondToModel(format!("agent job {job_id} not found")) - })?; + let job = db_ops::retry_locked("get_agent_job_after_runner", || async { + db.get_agent_job(job_id.as_str()).await + }) + .await + .map_err(|err| { + FunctionCallError::RespondToModel(format!("failed to load agent job {job_id}: {err}")) + })? + .ok_or_else(|| { + FunctionCallError::RespondToModel(format!("agent job {job_id} not found")) + })?; let output_path = PathBuf::from(job.output_csv_path.clone()); if !tokio::fs::try_exists(&output_path).await.unwrap_or(false) { export_job_csv_snapshot(db.clone(), &job) @@ -422,24 +431,27 @@ mod spawn_agents_on_csv { )) })?; } - let progress = db - .get_agent_job_progress(job_id.as_str()) - .await - .map_err(|err| { - FunctionCallError::RespondToModel(format!( - "failed to load agent job progress {job_id}: {err}" - )) - })?; + let progress = db_ops::retry_locked("get_agent_job_progress_after_runner", || async { + db.get_agent_job_progress(job_id.as_str()).await + }) + .await + .map_err(|err| { + FunctionCallError::RespondToModel(format!( + "failed to load agent job progress {job_id}: {err}" + )) + })?; let mut job_error = job.last_error.clone().filter(|err| !err.trim().is_empty()); let failed_item_errors = if progress.failed_items > 0 { - let items = db - .list_agent_job_items( + let items = db_ops::retry_locked("list_failed_agent_job_items", || async { + db.list_agent_job_items( job_id.as_str(), Some(codex_state::AgentJobItemStatus::Failed), Some(5), ) .await - .unwrap_or_default(); + }) + .await + .unwrap_or_default(); let summaries: Vec<_> = items .into_iter() .filter_map(|item| { @@ -503,21 +515,27 @@ mod report_agent_job_result { let reporting_thread_id = session.conversation_id; let reporting_thread_id_str = reporting_thread_id.to_string(); 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", - ) + db_ops::retry_locked("report_agent_job_item_result_and_cancel_job", || async { + 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 + }) .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, - ) + db_ops::retry_locked("report_agent_job_item_result", || async { + 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 + }) .await } .map_err(|err| { @@ -609,10 +627,11 @@ async fn run_agent_job_loop( job_id: String, options: JobRunnerOptions, ) -> anyhow::Result<()> { - let job = db - .get_agent_job(job_id.as_str()) - .await? - .ok_or_else(|| anyhow::anyhow!("agent job {job_id} was not found"))?; + let job = db_ops::retry_locked("get_agent_job_for_runner", || async { + db.get_agent_job(job_id.as_str()).await + }) + .await? + .ok_or_else(|| anyhow::anyhow!("agent job {job_id} was not found"))?; let runtime_timeout = job_runtime_timeout(&job); let mut active_items: HashMap = HashMap::new(); let mut starting_items = startup::StartupTasks::default(); @@ -625,7 +644,10 @@ async fn run_agent_job_loop( runtime_timeout, ) .await?; - let initial_progress = db.get_agent_job_progress(job_id.as_str()).await?; + let initial_progress = db_ops::retry_locked("get_initial_agent_job_progress", || async { + db.get_agent_job_progress(job_id.as_str()).await + }) + .await?; progress_emitter .maybe_emit( &session, @@ -636,12 +658,20 @@ async fn run_agent_job_loop( ) .await?; - let mut cancel_requested = db.is_agent_job_cancelled(job_id.as_str()).await?; + let mut cancel_requested = db_ops::retry_locked("is_agent_job_cancelled_initial", || async { + db.is_agent_job_cancelled(job_id.as_str()).await + }) + .await?; loop { let mut progressed = false; let mut agent_limit_reached = false; - if !cancel_requested && db.is_agent_job_cancelled(job_id.as_str()).await? { + if !cancel_requested + && db_ops::retry_locked("is_agent_job_cancelled_pre_launch", || async { + db.is_agent_job_cancelled(job_id.as_str()).await + }) + .await? + { cancel_requested = true; let _ = session .notify_background_event( @@ -662,7 +692,11 @@ async fn run_agent_job_loop( progressed |= startup_result.progressed; agent_limit_reached |= startup_result.agent_limit_reached; - let scheduler_progress = db.get_agent_job_progress(job_id.as_str()).await?; + let scheduler_progress = + db_ops::retry_locked("get_scheduler_agent_job_progress", || async { + db.get_agent_job_progress(job_id.as_str()).await + }) + .await?; if slots::reclaim_inactive_active_items( session.clone(), db.clone(), @@ -675,7 +709,12 @@ async fn run_agent_job_loop( progressed = true; } - if !cancel_requested && db.is_agent_job_cancelled(job_id.as_str()).await? { + if !cancel_requested + && db_ops::retry_locked("is_agent_job_cancelled_post_reclaim", || async { + db.is_agent_job_cancelled(job_id.as_str()).await + }) + .await? + { cancel_requested = true; progressed = true; let _ = session @@ -765,7 +804,10 @@ async fn run_agent_job_loop( let finished = find_finished_threads(session.clone(), &active_items).await; if finished.is_empty() { let progress = if progressed { - db.get_agent_job_progress(job_id.as_str()).await? + db_ops::retry_locked("get_agent_job_progress_after_progress", || async { + db.get_agent_job_progress(job_id.as_str()).await + }) + .await? } else { scheduler_progress }; @@ -812,7 +854,11 @@ async fn run_agent_job_loop( ) .await?; active_items.remove(&thread_id); - let progress = db.get_agent_job_progress(job_id.as_str()).await?; + let progress = + db_ops::retry_locked("get_agent_job_progress_after_finalize", || async { + db.get_agent_job_progress(job_id.as_str()).await + }) + .await?; progress_emitter .maybe_emit( &session, @@ -825,14 +871,24 @@ async fn run_agent_job_loop( } } - let progress = db.get_agent_job_progress(job_id.as_str()).await?; + let progress = db_ops::retry_locked("get_agent_job_progress_before_export", || async { + db.get_agent_job_progress(job_id.as_str()).await + }) + .await?; if let Err(err) = export_job_csv_snapshot(db.clone(), &job).await { let message = format!("auto-export failed: {err}"); - db.mark_agent_job_failed(job_id.as_str(), message.as_str()) - .await?; + db_ops::retry_locked("mark_agent_job_failed_after_export", || async { + db.mark_agent_job_failed(job_id.as_str(), message.as_str()) + .await + }) + .await?; return Ok(()); } - let cancelled = cancel_requested || db.is_agent_job_cancelled(job_id.as_str()).await?; + let cancelled = cancel_requested + || db_ops::retry_locked("is_agent_job_cancelled_before_complete", || async { + db.is_agent_job_cancelled(job_id.as_str()).await + }) + .await?; if cancelled { let pending_items = progress.pending_items; let message = @@ -854,8 +910,14 @@ async fn run_agent_job_loop( let message = format!("agent job completed with {failed_items} failed items"); let _ = session.notify_background_event(&turn, message).await; } - db.mark_agent_job_completed(job_id.as_str()).await?; - let progress = db.get_agent_job_progress(job_id.as_str()).await?; + db_ops::retry_locked("mark_agent_job_completed", || async { + db.mark_agent_job_completed(job_id.as_str()).await + }) + .await?; + let progress = db_ops::retry_locked("get_agent_job_progress_after_complete", || async { + db.get_agent_job_progress(job_id.as_str()).await + }) + .await?; progress_emitter .maybe_emit( &session, @@ -872,9 +934,11 @@ async fn export_job_csv_snapshot( db: Arc, job: &codex_state::AgentJob, ) -> anyhow::Result<()> { - let items = db - .list_agent_job_items(job.id.as_str(), /*status*/ None, /*limit*/ None) - .await?; + let items = db_ops::retry_locked("list_agent_job_items_for_export", || async { + db.list_agent_job_items(job.id.as_str(), /*status*/ None, /*limit*/ None) + .await + }) + .await?; let csv_content = render_job_csv(job.input_headers.as_slice(), items.as_slice()) .map_err(|err| anyhow::anyhow!("failed to render job csv for auto-export: {err}"))?; let output_path = PathBuf::from(job.output_csv_path.clone()); @@ -892,18 +956,24 @@ async fn recover_running_items( active_items: &mut HashMap, runtime_timeout: Duration, ) -> anyhow::Result<()> { - let running_items = db - .list_agent_job_items( - job_id, - Some(codex_state::AgentJobItemStatus::Running), - /*limit*/ None, - ) + let running_items = + db_ops::retry_locked("list_running_agent_job_items_for_recovery", || async { + db.list_agent_job_items( + job_id, + Some(codex_state::AgentJobItemStatus::Running), + /*limit*/ None, + ) + .await + }) .await?; for item in running_items { if is_item_stale(&item, runtime_timeout) { let error_message = format!("worker exceeded max runtime of {runtime_timeout:?}"); - db.mark_agent_job_item_failed(job_id, item.item_id.as_str(), error_message.as_str()) - .await?; + db_ops::retry_locked("mark_stale_agent_job_item_failed_on_recovery", || async { + db.mark_agent_job_item_failed(job_id, item.item_id.as_str(), error_message.as_str()) + .await + }) + .await?; if let Some(assigned_thread_id) = item.assigned_thread_id.as_ref() && let Ok(thread_id) = ThreadId::from_string(assigned_thread_id.as_str()) { @@ -912,11 +982,14 @@ async fn recover_running_items( continue; } let Some(assigned_thread_id) = item.assigned_thread_id.clone() else { - db.mark_agent_job_item_pending( - job_id, - item.item_id.as_str(), - Some("worker startup was interrupted before a thread was assigned"), - ) + db_ops::retry_locked("mark_agent_job_item_pending_on_recovery", || async { + db.mark_agent_job_item_pending( + job_id, + item.item_id.as_str(), + Some("worker startup was interrupted before a thread was assigned"), + ) + .await + }) .await?; continue; }; @@ -924,11 +997,14 @@ async fn recover_running_items( Ok(thread_id) => thread_id, Err(err) => { let error_message = format!("invalid assigned_thread_id: {err:?}"); - db.mark_agent_job_item_failed( - job_id, - item.item_id.as_str(), - error_message.as_str(), - ) + db_ops::retry_locked("mark_agent_job_item_failed_invalid_thread", || async { + db.mark_agent_job_item_failed( + job_id, + item.item_id.as_str(), + error_message.as_str(), + ) + .await + }) .await?; continue; } @@ -1023,8 +1099,11 @@ async fn reap_stale_active_items( } for (thread_id, item_id) in stale { let error_message = format!("worker exceeded max runtime of {runtime_timeout:?}"); - db.mark_agent_job_item_failed(job_id, item_id.as_str(), error_message.as_str()) - .await?; + db_ops::retry_locked("mark_stale_active_agent_job_item_failed", || async { + db.mark_agent_job_item_failed(job_id, item_id.as_str(), error_message.as_str()) + .await + }) + .await?; let _ = session .services .agent_control @@ -1042,23 +1121,27 @@ async fn finalize_finished_item( item_id: &str, thread_id: ThreadId, ) -> anyhow::Result<()> { - let item = db - .get_agent_job_item(job_id, item_id) - .await? - .ok_or_else(|| { - anyhow::anyhow!("job item not found for finalization: {job_id}/{item_id}") - })?; + let item = db_ops::retry_locked("get_agent_job_item_for_finalization", || async { + db.get_agent_job_item(job_id, item_id).await + }) + .await? + .ok_or_else(|| anyhow::anyhow!("job item not found for finalization: {job_id}/{item_id}"))?; if matches!(item.status, codex_state::AgentJobItemStatus::Running) { if item.result_json.is_some() { - let _ = db.mark_agent_job_item_completed(job_id, item_id).await?; + let _ = db_ops::retry_locked("mark_agent_job_item_completed", || async { + db.mark_agent_job_item_completed(job_id, item_id).await + }) + .await?; } else { - let _ = db - .mark_agent_job_item_failed( + let _ = db_ops::retry_locked("mark_agent_job_item_failed_missing_report", || async { + db.mark_agent_job_item_failed( job_id, item_id, "worker finished without calling report_agent_job_result", ) - .await?; + .await + }) + .await?; } } let _ = session diff --git a/codex-rs/core/src/tools/handlers/agent_jobs_db.rs b/codex-rs/core/src/tools/handlers/agent_jobs_db.rs new file mode 100644 index 0000000000..03a607822f --- /dev/null +++ b/codex-rs/core/src/tools/handlers/agent_jobs_db.rs @@ -0,0 +1,95 @@ +use std::future::Future; +use std::time::Duration; + +const SQLITE_LOCK_RETRY_LIMIT: usize = 2; + +fn is_sqlite_lock_error(err: &anyhow::Error) -> bool { + err.chain().any(|cause| { + let message = cause.to_string(); + message.contains("database is locked") || message.contains("database table is locked") + }) +} + +fn retry_delay(attempt: usize) -> Duration { + Duration::from_millis(250 * (attempt as u64 + 1)) +} + +pub(super) async fn retry_locked(operation: &'static str, mut op: F) -> anyhow::Result +where + F: FnMut() -> Fut, + Fut: Future>, +{ + let mut attempt = 0usize; + loop { + match op().await { + Ok(value) => return Ok(value), + Err(err) if is_sqlite_lock_error(&err) && attempt < SQLITE_LOCK_RETRY_LIMIT => { + let retry_in = retry_delay(attempt); + tracing::warn!( + operation, + attempt = attempt + 1, + max_attempts = SQLITE_LOCK_RETRY_LIMIT + 1, + retry_delay_ms = retry_in.as_millis() as u64, + error = %err, + "agent job DB operation hit sqlite lock; retrying" + ); + tokio::time::sleep(retry_in).await; + attempt += 1; + } + Err(err) => return Err(err), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use pretty_assertions::assert_eq; + use std::sync::Arc; + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; + + #[tokio::test] + async fn retry_locked_retries_transient_sqlite_locks() -> anyhow::Result<()> { + let attempts = Arc::new(AtomicUsize::new(0)); + let attempts_for_op = Arc::clone(&attempts); + + let result = retry_locked("test_sqlite_retry", move || { + let attempts = Arc::clone(&attempts_for_op); + async move { + let attempt = attempts.fetch_add(1, Ordering::SeqCst); + if attempt < 2 { + Err(anyhow::anyhow!( + "error returned from database: (code: 5) database is locked" + )) + } else { + Ok("ok") + } + } + }) + .await?; + + assert_eq!(result, "ok"); + assert_eq!(attempts.load(Ordering::SeqCst), 3); + Ok(()) + } + + #[tokio::test] + async fn retry_locked_does_not_retry_non_lock_errors() { + let attempts = Arc::new(AtomicUsize::new(0)); + let attempts_for_op = Arc::clone(&attempts); + + let err = retry_locked("test_non_sqlite_retry", move || { + let attempts = Arc::clone(&attempts_for_op); + async move { + attempts.fetch_add(1, Ordering::SeqCst); + Err::<(), anyhow::Error>(anyhow::anyhow!("boom")) + } + }) + .await + .expect_err("non-lock errors should not be retried"); + + assert_eq!(err.to_string(), "boom"); + assert_eq!(attempts.load(Ordering::SeqCst), 1); + } +} 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 52cc70645a..5134ad7e82 100644 --- a/codex-rs/core/src/tools/handlers/agent_jobs_slots.rs +++ b/codex-rs/core/src/tools/handlers/agent_jobs_slots.rs @@ -5,7 +5,7 @@ use std::collections::HashMap; struct InactiveActiveItem { thread_id: ThreadId, item_id: String, - db_status: Option, + db_state: &'static str, db_assigned_thread_id: Option, } @@ -16,29 +16,42 @@ pub(super) async fn reclaim_inactive_active_items( active_items: &mut HashMap, db_running_items: usize, ) -> anyhow::Result { + let running_items = + db_ops::retry_locked("list_running_agent_job_items_for_reclaim", || async { + db.list_agent_job_items( + job_id, + Some(codex_state::AgentJobItemStatus::Running), + /*limit*/ None, + ) + .await + }) + .await?; + let running_by_item_id: HashMap<_, _> = running_items + .into_iter() + .map(|item| (item.item_id, item.assigned_thread_id)) + .collect(); + let mut inactive_items = Vec::new(); for (thread_id, item) in active_items.iter() { let thread_id_str = thread_id.to_string(); - let db_item = db.get_agent_job_item(job_id, item.item_id.as_str()).await?; - let Some(db_item) = db_item else { + let Some(db_assigned_thread_id) = running_by_item_id.get(item.item_id.as_str()) else { inactive_items.push(InactiveActiveItem { thread_id: *thread_id, item_id: item.item_id.clone(), - db_status: None, + db_state: "missing", db_assigned_thread_id: None, }); continue; }; - let still_running = matches!(db_item.status, codex_state::AgentJobItemStatus::Running) - && db_item.assigned_thread_id.as_deref() == Some(thread_id_str.as_str()); + let still_running = db_assigned_thread_id.as_deref() == Some(thread_id_str.as_str()); if still_running { continue; } inactive_items.push(InactiveActiveItem { thread_id: *thread_id, item_id: item.item_id.clone(), - db_status: Some(db_item.status), - db_assigned_thread_id: db_item.assigned_thread_id, + db_state: "running", + db_assigned_thread_id: db_assigned_thread_id.clone(), }); } if inactive_items.is_empty() { @@ -63,11 +76,7 @@ pub(super) async fn reclaim_inactive_active_items( job_id, item_id = inactive_item.item_id, thread_id = %inactive_item.thread_id, - db_status = inactive_item - .db_status - .as_ref() - .map(|status| status.as_str()) - .unwrap_or("missing"), + db_status = inactive_item.db_state, db_assigned_thread_id = inactive_item.db_assigned_thread_id.as_deref().unwrap_or(""), active_items = active_items.len(), "agent job reclaimed scheduler slot" 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 12c071bf53..ada8c42f00 100644 --- a/codex-rs/core/src/tools/handlers/agent_jobs_startup.rs +++ b/codex-rs/core/src/tools/handlers/agent_jobs_startup.rs @@ -92,20 +92,24 @@ pub(super) async fn launch_pending_items( return Ok(false); } - let pending_items = db - .list_agent_job_items( + let pending_items = db_ops::retry_locked("list_pending_agent_job_items_for_launch", || async { + db.list_agent_job_items( job_id, Some(codex_state::AgentJobItemStatus::Pending), Some(slots), ) - .await?; + .await + }) + .await?; let mut launched = 0usize; let mut progressed = false; for item in pending_items { - let claimed = db - .mark_agent_job_item_running(job_id, item.item_id.as_str()) - .await?; + let claimed = db_ops::retry_locked("mark_agent_job_item_running_for_launch", || async { + db.mark_agent_job_item_running(job_id, item.item_id.as_str()) + .await + }) + .await?; if !claimed { continue; } @@ -114,11 +118,14 @@ pub(super) async fn launch_pending_items( Ok(prompt) => prompt, Err(err) => { let error_message = format!("failed to build worker prompt: {err}"); - db.mark_agent_job_item_failed( - job_id, - item.item_id.as_str(), - error_message.as_str(), - ) + db_ops::retry_locked("mark_agent_job_item_failed_for_prompt_build", || async { + db.mark_agent_job_item_failed( + job_id, + item.item_id.as_str(), + error_message.as_str(), + ) + .await + }) .await?; progressed = true; continue; @@ -283,8 +290,11 @@ pub(super) async fn reap_stale_startups( item.abort_handle.abort(); let error_message = format!("worker exceeded max runtime of {runtime_timeout:?} before startup completed"); - db.mark_agent_job_item_failed(job_id, item.item_id.as_str(), error_message.as_str()) - .await?; + db_ops::retry_locked("mark_agent_job_item_failed_for_stale_startup", || async { + db.mark_agent_job_item_failed(job_id, item.item_id.as_str(), error_message.as_str()) + .await + }) + .await?; tracing::warn!( job_id, item_id = item.item_id, @@ -310,12 +320,15 @@ async fn handle_worker_startup_result( match startup.result { Ok(thread_id) => { let thread_id_str = thread_id.to_string(); - let assigned = db - .set_agent_job_item_thread( - job_id, - startup.item_id.as_str(), - thread_id_str.as_str(), - ) + let assigned = + db_ops::retry_locked("set_agent_job_item_thread_after_startup", || async { + db.set_agent_job_item_thread( + job_id, + startup.item_id.as_str(), + thread_id_str.as_str(), + ) + .await + }) .await?; if !assigned { let _ = session @@ -364,12 +377,15 @@ async fn handle_worker_startup_result( }) } Err(CodexErr::AgentLimitReached { .. }) => { - let _ = db - .mark_agent_job_item_pending( - job_id, - startup.item_id.as_str(), - /*error_message*/ None, - ) + let _ = + db_ops::retry_locked("mark_agent_job_item_pending_after_limit", || async { + db.mark_agent_job_item_pending( + job_id, + startup.item_id.as_str(), + /*error_message*/ None, + ) + .await + }) .await?; tracing::debug!( job_id, @@ -384,13 +400,18 @@ async fn handle_worker_startup_result( } Err(err) => { let error_message = format!("failed to spawn worker: {err}"); - let _ = db - .mark_agent_job_item_failed( - job_id, - startup.item_id.as_str(), - error_message.as_str(), - ) - .await?; + let _ = db_ops::retry_locked( + "mark_agent_job_item_failed_after_spawn_error", + || async { + db.mark_agent_job_item_failed( + job_id, + startup.item_id.as_str(), + error_message.as_str(), + ) + .await + }, + ) + .await?; tracing::warn!( job_id, item_id = startup.item_id, @@ -410,9 +431,18 @@ async fn handle_worker_startup_result( return Ok(StartupDrainResult::default()); }; let error_message = format!("worker startup task failed: {join_error}"); - let _ = db - .mark_agent_job_item_failed(job_id, item.item_id.as_str(), error_message.as_str()) - .await?; + let _ = db_ops::retry_locked( + "mark_agent_job_item_failed_after_startup_join_error", + || async { + db.mark_agent_job_item_failed( + job_id, + item.item_id.as_str(), + error_message.as_str(), + ) + .await + }, + ) + .await?; tracing::warn!( job_id, item_id = item.item_id, diff --git a/codex-rs/core/src/tools/parallel.rs b/codex-rs/core/src/tools/parallel.rs index 0e649333df..59ccb34100 100644 --- a/codex-rs/core/src/tools/parallel.rs +++ b/codex-rs/core/src/tools/parallel.rs @@ -2,6 +2,7 @@ use std::sync::Arc; use std::time::Instant; use tokio::sync::RwLock; +use tokio::task::JoinError; use tokio_util::either::Either; use tokio_util::sync::CancellationToken; use tokio_util::task::AbortOnDropHandle; @@ -85,6 +86,7 @@ impl ToolCallRuntime { let tracker = Arc::clone(&self.tracker); let lock = Arc::clone(&self.parallel_execution); let started = Instant::now(); + let join_error_call = call.clone(); let dispatch_span = trace_span!( "dispatch_tool_call_with_code_mode_result", @@ -125,13 +127,38 @@ impl ToolCallRuntime { async move { handle.await.map_err(|err| { - FunctionCallError::Fatal(format!("tool task failed to receive: {err:?}")) + tool_task_join_error_to_function_call_error( + err, + &join_error_call, + started.elapsed(), + ) })? } .in_current_span() } } +fn tool_task_join_error_to_function_call_error( + err: JoinError, + call: &ToolCall, + elapsed: std::time::Duration, +) -> FunctionCallError { + if err.is_cancelled() { + let secs = elapsed.as_secs_f32().max(0.1); + tracing::warn!( + tool_name = call.tool_name, + call_id = call.call_id, + elapsed_seconds = secs, + "tool task was cancelled before delivering a response" + ); + return FunctionCallError::RespondToModel(format!( + "tool execution interrupted after {secs:.1}s" + )); + } + + FunctionCallError::Fatal(format!("tool task failed to receive: {err:?}")) +} + impl ToolCallRuntime { fn failure_response(call: ToolCall, err: FunctionCallError) -> ResponseInputItem { let message = err.to_string(); @@ -179,3 +206,61 @@ impl ToolCallRuntime { } } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::tools::context::ToolPayload; + use pretty_assertions::assert_eq; + use serde_json::json; + + fn test_tool_call() -> ToolCall { + ToolCall { + tool_name: "spawn_agents_on_csv".to_string(), + call_id: "call-1".to_string(), + tool_namespace: None, + payload: ToolPayload::Function { + arguments: json!({"csv_path":"in.csv","instruction":"test"}).to_string(), + }, + } + } + + #[tokio::test] + async fn cancelled_tool_task_join_error_becomes_model_response() { + let handle = tokio::spawn(async { + tokio::time::sleep(std::time::Duration::from_secs(60)).await; + }); + handle.abort(); + let err = handle.await.expect_err("task should be cancelled"); + + let mapped = tool_task_join_error_to_function_call_error( + err, + &test_tool_call(), + std::time::Duration::from_millis(200), + ); + + assert_eq!( + mapped, + FunctionCallError::RespondToModel("tool execution interrupted after 0.2s".to_string()) + ); + } + + #[tokio::test] + async fn panicked_tool_task_join_error_stays_fatal() { + let handle = tokio::spawn(async { + panic!("boom"); + }); + let err = handle.await.expect_err("task should panic"); + + let mapped = tool_task_join_error_to_function_call_error( + err, + &test_tool_call(), + std::time::Duration::from_millis(200), + ); + + let FunctionCallError::Fatal(message) = mapped else { + panic!("panic join errors should remain fatal"); + }; + assert!(message.contains("tool task failed to receive")); + } +} diff --git a/codex-rs/rollout/src/state_db.rs b/codex-rs/rollout/src/state_db.rs index 6367a27ca9..cba86194b8 100644 --- a/codex-rs/rollout/src/state_db.rs +++ b/codex-rs/rollout/src/state_db.rs @@ -15,21 +15,72 @@ pub use codex_state::LogEntry; use codex_state::ThreadMetadataBuilder; use codex_utils_path::normalize_for_path_comparison; use serde_json::Value; +use std::collections::HashMap; use std::path::Path; use std::path::PathBuf; use std::sync::Arc; +use std::sync::LazyLock; +use std::sync::Mutex; +use tokio::sync::OnceCell; use tracing::warn; use uuid::Uuid; /// Core-facing handle to the SQLite-backed state runtime. pub type StateDbHandle = Arc; +#[derive(Clone, Debug, Eq, Hash, PartialEq)] +struct StateDbCacheKey { + sqlite_home: PathBuf, + default_provider: String, +} + +type StateDbRuntimeCell = Arc>; + +static STATE_DB_RUNTIME_CACHE: LazyLock>> = + LazyLock::new(|| Mutex::new(HashMap::new())); + +fn cache_key(sqlite_home: &Path, default_provider: &str) -> StateDbCacheKey { + StateDbCacheKey { + sqlite_home: normalize_for_path_comparison(sqlite_home) + .unwrap_or_else(|_| sqlite_home.to_path_buf()), + default_provider: default_provider.to_string(), + } +} + +fn cached_runtime_cell(sqlite_home: &Path, default_provider: &str) -> StateDbRuntimeCell { + let key = cache_key(sqlite_home, default_provider); + let mut cache = STATE_DB_RUNTIME_CACHE + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if let Some(cell) = cache.get(&key) { + return Arc::clone(cell); + } + + let cell = Arc::new(OnceCell::new()); + cache.insert(key, Arc::clone(&cell)); + cell +} + +async fn shared_state_db_runtime( + sqlite_home: &Path, + default_provider: &str, +) -> anyhow::Result { + let cell = cached_runtime_cell(sqlite_home, default_provider); + let runtime = cell + .get_or_try_init(|| async { + codex_state::StateRuntime::init(sqlite_home.to_path_buf(), default_provider.to_string()) + .await + }) + .await?; + Ok(Arc::clone(runtime)) +} + /// Initialize the state runtime for thread state persistence and backfill checks. pub async fn init(config: &impl RolloutConfigView) -> Option { let config = RolloutConfig::from_view(config); - let runtime = match codex_state::StateRuntime::init( - config.sqlite_home.clone(), - config.model_provider_id.clone(), + let runtime = match shared_state_db_runtime( + config.sqlite_home.as_path(), + config.model_provider_id.as_str(), ) .await { @@ -68,12 +119,9 @@ pub async fn get_state_db(config: &impl RolloutConfigView) -> Option Optio if !tokio::fs::try_exists(&db_path).await.unwrap_or(false) { return None; } - let runtime = - codex_state::StateRuntime::init(codex_home.to_path_buf(), default_provider.to_string()) - .await - .ok()?; + let runtime = shared_state_db_runtime(codex_home, default_provider) + .await + .ok()?; require_backfill_complete(runtime, codex_home).await } diff --git a/codex-rs/rollout/src/state_db_tests.rs b/codex-rs/rollout/src/state_db_tests.rs index e5e3211768..6cee6aa473 100644 --- a/codex-rs/rollout/src/state_db_tests.rs +++ b/codex-rs/rollout/src/state_db_tests.rs @@ -1,12 +1,16 @@ #![allow(warnings, clippy::all)] use super::*; +use crate::config::RolloutConfig; use crate::list::parse_cursor; use chrono::DateTime; use chrono::NaiveDateTime; use chrono::Timelike; use chrono::Utc; use pretty_assertions::assert_eq; +use std::path::Path; +use std::sync::Arc; +use tempfile::TempDir; use uuid::Uuid; #[test] @@ -26,3 +30,82 @@ fn cursor_to_anchor_normalizes_timestamp_format() { assert_eq!(anchor.id, uuid); assert_eq!(anchor.ts, expected_ts); } + +fn test_config(home: &Path) -> RolloutConfig { + RolloutConfig { + codex_home: home.to_path_buf(), + sqlite_home: home.to_path_buf(), + cwd: home.to_path_buf(), + model_provider_id: "test-provider".to_string(), + generate_memories: false, + } +} + +#[tokio::test] +async fn init_reuses_cached_runtime_for_same_home() { + let temp = TempDir::new().expect("temp dir"); + let config = test_config(temp.path()); + + let first = init(&config).await.expect("state db init should succeed"); + first + .mark_backfill_complete(/*last_watermark*/ None) + .await + .expect("backfill should be marked complete"); + let second = init(&config) + .await + .expect("cached state db init should succeed"); + + assert!(Arc::ptr_eq(&first, &second)); +} + +#[tokio::test] +async fn get_state_db_reuses_cached_runtime() { + let temp = TempDir::new().expect("temp dir"); + let config = test_config(temp.path()); + + let first = init(&config).await.expect("state db init should succeed"); + first + .mark_backfill_complete(/*last_watermark*/ None) + .await + .expect("backfill should be marked complete"); + + let reopened = get_state_db(&config) + .await + .expect("cached state db should be returned"); + + assert!(Arc::ptr_eq(&first, &reopened)); + assert_eq!(reopened.codex_home(), config.sqlite_home.as_path()); +} + +#[tokio::test] +async fn concurrent_init_reuses_single_cached_runtime() { + let temp = TempDir::new().expect("temp dir"); + let config = Arc::new(test_config(temp.path())); + + let mut handles = Vec::new(); + for _ in 0..8 { + let config = Arc::clone(&config); + handles.push(tokio::spawn(async move { + init(config.as_ref()) + .await + .expect("state db init should succeed") + })); + } + + let mut runtimes = Vec::new(); + for handle in handles { + runtimes.push(handle.await.expect("task should join")); + } + let first = runtimes + .first() + .cloned() + .expect("at least one runtime should exist"); + first + .mark_backfill_complete(/*last_watermark*/ None) + .await + .expect("backfill should be marked complete"); + + for runtime in runtimes.iter().skip(1) { + assert!(Arc::ptr_eq(&first, runtime)); + } +}