openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-21 17:48:44 +00:00
Code Issues Packages Projects Releases Wiki Activity

Reduce SQLite lock contention in agent jobs

Browse Source 
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.
This commit is contained in:
Dave Aitel
2026-04-17 08:26:00 -04:00
parent 47f9e3ed87
commit b4e16db47b
7 changed files with 601 additions and 169 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/tools/handlers/agent_jobs.rs
+190 -107
View File
@@ -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<ThreadId, ActiveJobItem> = 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<codex_state::StateRuntime>,
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<ThreadId, ActiveJobItem>,
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
codex-rs/core/src/tools/handlers/agent_jobs_db.rs
+95
View File
@@ -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<T, F, Fut>(operation: &'static str, mut op: F) -> anyhow::Result<T>
where
F: FnMut() -> Fut,
Fut: Future<Output = anyhow::Result<T>>,
{
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);
}
}
codex-rs/core/src/tools/handlers/agent_jobs_slots.rs
+22 -13
View File
@@ -5,7 +5,7 @@ use std::collections::HashMap;
struct InactiveActiveItem {
thread_id: ThreadId,
item_id: String,
db_status: Option<codex_state::AgentJobItemStatus>,
db_state: &'static str,
db_assigned_thread_id: Option<String>,
}
@@ -16,29 +16,42 @@ pub(super) async fn reclaim_inactive_active_items(
active_items: &mut HashMap<ThreadId, ActiveJobItem>,
db_running_items: usize,
) -> anyhow::Result<bool> {
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"
codex-rs/core/src/tools/handlers/agent_jobs_startup.rs
+65 -35
View File
@@ -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,
codex-rs/core/src/tools/parallel.rs
+86 -1
View File
@@ -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"));
}
}
codex-rs/rollout/src/state_db.rs
+60 -13
View File
@@ -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<codex_state::StateRuntime>;
#[derive(Clone, Debug, Eq, Hash, PartialEq)]
struct StateDbCacheKey {
sqlite_home: PathBuf,
default_provider: String,
}
type StateDbRuntimeCell = Arc<OnceCell<StateDbHandle>>;
static STATE_DB_RUNTIME_CACHE: LazyLock<Mutex<HashMap<StateDbCacheKey, StateDbRuntimeCell>>> =
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<StateDbHandle> {
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<StateDbHandle> {
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<StateDbHand
if !tokio::fs::try_exists(&state_path).await.unwrap_or(false) {
return None;
}
let runtime = codex_state::StateRuntime::init(
config.sqlite_home().to_path_buf(),
config.model_provider_id().to_string(),
)
.await
.ok()?;
let runtime = shared_state_db_runtime(config.sqlite_home(), config.model_provider_id())
.await
.ok()?;
require_backfill_complete(runtime, config.sqlite_home()).await
}
@@ -85,10 +133,9 @@ pub async fn open_if_present(codex_home: &Path, default_provider: &str) -> 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
}
codex-rs/rollout/src/state_db_tests.rs
+83
View File
@@ -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));
}
}