From e112e79b8aa04edcc45514b8c1a1b079725da573 Mon Sep 17 00:00:00 2001 From: Dave Aitel Date: Thu, 5 Feb 2026 12:04:23 -0500 Subject: [PATCH] Add agent job runner for CSV batches --- .../core/src/tools/handlers/agent_jobs.rs | 1139 +++++++++++++++++ codex-rs/core/src/tools/handlers/collab.rs | 11 +- codex-rs/core/src/tools/handlers/mod.rs | 1 + codex-rs/core/src/tools/spec.rs | 232 +++- codex-rs/docs/agent_jobs.md | 70 + codex-rs/state/migrations/0009_agent_jobs.sql | 37 + codex-rs/state/src/lib.rs | 7 + codex-rs/state/src/model/agent_job.rs | 285 +++++ codex-rs/state/src/model/mod.rs | 10 + codex-rs/state/src/runtime.rs | 457 +++++++ 10 files changed, 2239 insertions(+), 10 deletions(-) create mode 100644 codex-rs/core/src/tools/handlers/agent_jobs.rs create mode 100644 codex-rs/docs/agent_jobs.md create mode 100644 codex-rs/state/migrations/0009_agent_jobs.sql create mode 100644 codex-rs/state/src/model/agent_job.rs diff --git a/codex-rs/core/src/tools/handlers/agent_jobs.rs b/codex-rs/core/src/tools/handlers/agent_jobs.rs new file mode 100644 index 0000000000..ac606fc8a7 --- /dev/null +++ b/codex-rs/core/src/tools/handlers/agent_jobs.rs @@ -0,0 +1,1139 @@ +use crate::agent::exceeds_thread_spawn_depth_limit; +use crate::agent::next_thread_spawn_depth; +use crate::agent::status::is_final; +use crate::codex::Session; +use crate::codex::TurnContext; +use crate::config::Config; +use crate::error::CodexErr; +use crate::function_tool::FunctionCallError; +use crate::tools::context::ToolInvocation; +use crate::tools::context::ToolOutput; +use crate::tools::context::ToolPayload; +use crate::tools::handlers::collab::build_agent_spawn_config; +use crate::tools::handlers::parse_arguments; +use crate::tools::registry::ToolHandler; +use crate::tools::registry::ToolKind; +use async_trait::async_trait; +use codex_protocol::ThreadId; +use codex_protocol::models::FunctionCallOutputBody; +use codex_protocol::protocol::SessionSource; +use codex_protocol::protocol::SubAgentSource; +use once_cell::sync::Lazy; +use serde::Deserialize; +use serde::Serialize; +use serde_json::Value; +use serde_json::json; +use std::collections::HashMap; +use std::collections::HashSet; +use std::path::Path; +use std::path::PathBuf; +use std::sync::Arc; +use tokio::sync::Mutex; +use tokio::task::JoinHandle; +use tokio::time::Duration; +use tokio::time::Instant; +use uuid::Uuid; + +pub struct AgentJobsHandler; + +const MIN_WAIT_TIMEOUT_MS: i64 = 1_000; +const DEFAULT_WAIT_TIMEOUT_MS: i64 = 30_000; +const MAX_WAIT_TIMEOUT_MS: i64 = 300_000; +const STATUS_POLL_INTERVAL_MS: u64 = 250; + +static ACTIVE_JOB_RUNNERS: Lazy>>> = + Lazy::new(|| Mutex::new(HashMap::new())); + +#[derive(Debug, Deserialize)] +struct SpawnAgentsOnCsvArgs { + csv_path: String, + instruction: String, + id_column: Option, + job_name: Option, + output_csv_path: Option, + output_schema: Option, + max_concurrency: Option, +} + +#[derive(Debug, Deserialize)] +struct JobIdArgs { + job_id: String, +} + +#[derive(Debug, Deserialize)] +struct RunAgentJobArgs { + job_id: String, + max_concurrency: Option, +} + +#[derive(Debug, Deserialize)] +struct WaitAgentJobArgs { + job_id: String, + timeout_ms: Option, +} + +#[derive(Debug, Deserialize)] +struct ExportAgentJobCsvArgs { + job_id: String, + path: Option, +} + +#[derive(Debug, Deserialize)] +struct ReportAgentJobResultArgs { + job_id: String, + item_id: String, + result: Value, +} + +#[derive(Debug, Serialize)] +struct AgentJobToolResult { + job_id: String, + status: String, + total_items: usize, + pending_items: usize, + running_items: usize, + completed_items: usize, + failed_items: usize, + output_csv_path: String, + runner_active: bool, +} + +#[derive(Debug, Serialize)] +struct SpawnAgentsOnCsvResult { + job_id: String, + started: bool, + output_csv_path: String, + total_items: usize, +} + +#[derive(Debug, Serialize)] +struct WaitAgentJobResult { + status: AgentJobToolResult, + timed_out: bool, +} + +#[derive(Debug, Serialize)] +struct ExportAgentJobCsvResult { + job_id: String, + path: String, + row_count: usize, +} + +#[derive(Debug, Serialize)] +struct ReportAgentJobResultToolResult { + accepted: bool, +} + +#[derive(Debug, Clone)] +struct JobRunnerOptions { + max_concurrency: usize, + spawn_config: Config, + child_depth: i32, +} + +#[async_trait] +impl ToolHandler for AgentJobsHandler { + fn kind(&self) -> ToolKind { + ToolKind::Function + } + + fn matches_kind(&self, payload: &ToolPayload) -> bool { + matches!(payload, ToolPayload::Function { .. }) + } + + async fn handle(&self, invocation: ToolInvocation) -> Result { + let ToolInvocation { + session, + turn, + tool_name, + payload, + .. + } = invocation; + + let arguments = match payload { + ToolPayload::Function { arguments } => arguments, + _ => { + return Err(FunctionCallError::RespondToModel( + "agent jobs handler received unsupported payload".to_string(), + )); + } + }; + + match tool_name.as_str() { + "spawn_agents_on_csv" => spawn_agents_on_csv::handle(session, turn, arguments).await, + "run_agent_job" => run_agent_job::handle(session, turn, arguments).await, + "get_agent_job_status" => get_agent_job_status::handle(session, arguments).await, + "wait_agent_job" => wait_agent_job::handle(session, arguments).await, + "export_agent_job_csv" => export_agent_job_csv::handle(session, turn, arguments).await, + "report_agent_job_result" => report_agent_job_result::handle(session, arguments).await, + other => Err(FunctionCallError::RespondToModel(format!( + "unsupported agent job tool {other}" + ))), + } + } +} + +mod spawn_agents_on_csv { + use super::*; + + pub async fn handle( + session: Arc, + turn: Arc, + arguments: String, + ) -> Result { + let args: SpawnAgentsOnCsvArgs = parse_arguments(arguments.as_str())?; + if args.instruction.trim().is_empty() { + return Err(FunctionCallError::RespondToModel( + "instruction must be non-empty".to_string(), + )); + } + + let db = required_state_db(&session)?; + let input_path = turn.resolve_path(Some(args.csv_path)); + let input_path_display = input_path.display().to_string(); + let csv_content = tokio::fs::read_to_string(&input_path) + .await + .map_err(|err| { + FunctionCallError::RespondToModel(format!( + "failed to read csv input {input_path_display}: {err}" + )) + })?; + let (headers, rows) = parse_csv(csv_content.as_str()).map_err(|err| { + FunctionCallError::RespondToModel(format!("failed to parse csv input: {err}")) + })?; + if headers.is_empty() { + return Err(FunctionCallError::RespondToModel( + "csv input must include a header row".to_string(), + )); + } + + let id_column_index = args.id_column.as_ref().map_or(Ok(None), |column_name| { + headers + .iter() + .position(|header| header == column_name) + .map(Some) + .ok_or_else(|| { + FunctionCallError::RespondToModel(format!( + "id_column {column_name} was not found in csv headers" + )) + }) + })?; + + let mut items = Vec::with_capacity(rows.len()); + let mut seen_ids = HashSet::new(); + for (idx, row) in rows.into_iter().enumerate() { + if row.len() != headers.len() { + let row_index = idx + 2; + let row_len = row.len(); + let header_len = headers.len(); + return Err(FunctionCallError::RespondToModel(format!( + "csv row {row_index} has {row_len} fields but header has {header_len}" + ))); + } + + let source_id = id_column_index + .and_then(|index| row.get(index).cloned()) + .filter(|value| !value.trim().is_empty()); + let row_index = idx + 1; + let mut item_id = source_id + .clone() + .unwrap_or_else(|| format!("row-{row_index}")); + if !seen_ids.insert(item_id.clone()) { + item_id = format!("{item_id}-{row_index}"); + seen_ids.insert(item_id.clone()); + } + + let row_object = headers + .iter() + .zip(row.iter()) + .map(|(header, value)| (header.clone(), Value::String(value.clone()))) + .collect::>(); + items.push(codex_state::AgentJobItemCreateParams { + item_id, + row_index: idx as i64, + source_id, + row_json: Value::Object(row_object), + }); + } + + let job_id = Uuid::new_v4().to_string(); + let output_csv_path = args.output_csv_path.map_or_else( + || default_output_csv_path(input_path.as_path(), job_id.as_str()), + |path| turn.resolve_path(Some(path)), + ); + let job_suffix = &job_id[..8]; + let job_name = args + .job_name + .unwrap_or_else(|| format!("agent-job-{job_suffix}")); + let job = db + .create_agent_job( + &codex_state::AgentJobCreateParams { + id: job_id.clone(), + name: job_name, + instruction: args.instruction, + output_schema_json: args.output_schema, + input_headers: headers, + 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}")) + })?; + + let options = build_runner_options(&session, &turn, args.max_concurrency).await?; + let started = start_job_runner(session, job_id.clone(), options).await?; + + let content = serde_json::to_string(&SpawnAgentsOnCsvResult { + job_id, + started, + output_csv_path: job.output_csv_path, + total_items: items.len(), + }) + .map_err(|err| { + FunctionCallError::Fatal(format!( + "failed to serialize spawn_agents_on_csv result: {err}" + )) + })?; + Ok(ToolOutput::Function { + body: FunctionCallOutputBody::Text(content), + success: Some(true), + }) + } +} + +mod run_agent_job { + use super::*; + + pub async fn handle( + session: Arc, + turn: Arc, + arguments: String, + ) -> Result { + let args: RunAgentJobArgs = parse_arguments(arguments.as_str())?; + let job_id = args.job_id; + let db = required_state_db(&session)?; + if db + .get_agent_job(job_id.as_str()) + .await + .map_err(|err| { + FunctionCallError::RespondToModel(format!( + "failed to look up agent job {job_id}: {err}" + )) + })? + .is_none() + { + return Err(FunctionCallError::RespondToModel(format!( + "agent job {job_id} not found" + ))); + } + let options = build_runner_options(&session, &turn, args.max_concurrency).await?; + let started = start_job_runner(session, job_id.clone(), options).await?; + let status = render_job_status(db, job_id.as_str()).await?; + let content = serde_json::to_string(&json!({ + "started": started, + "status": status, + })) + .map_err(|err| { + FunctionCallError::Fatal(format!("failed to serialize run_agent_job result: {err}")) + })?; + Ok(ToolOutput::Function { + body: FunctionCallOutputBody::Text(content), + success: Some(true), + }) + } +} + +mod get_agent_job_status { + use super::*; + + pub async fn handle( + session: Arc, + arguments: String, + ) -> Result { + let args: JobIdArgs = parse_arguments(arguments.as_str())?; + let job_id = args.job_id; + let db = required_state_db(&session)?; + let status = render_job_status(db, job_id.as_str()).await?; + let content = serde_json::to_string(&status).map_err(|err| { + FunctionCallError::Fatal(format!( + "failed to serialize get_agent_job_status result: {err}" + )) + })?; + Ok(ToolOutput::Function { + body: FunctionCallOutputBody::Text(content), + success: Some(true), + }) + } +} + +mod wait_agent_job { + use super::*; + + pub async fn handle( + session: Arc, + arguments: String, + ) -> Result { + let args: WaitAgentJobArgs = parse_arguments(arguments.as_str())?; + let db = required_state_db(&session)?; + let timeout_ms = args.timeout_ms.unwrap_or(DEFAULT_WAIT_TIMEOUT_MS); + let timeout_ms = match timeout_ms { + ms if ms <= 0 => { + return Err(FunctionCallError::RespondToModel( + "timeout_ms must be greater than zero".to_string(), + )); + } + ms => ms.clamp(MIN_WAIT_TIMEOUT_MS, MAX_WAIT_TIMEOUT_MS), + }; + + let deadline = Instant::now() + Duration::from_millis(timeout_ms as u64); + loop { + let status = render_job_status(db.clone(), args.job_id.as_str()).await?; + if matches!(status.status.as_str(), "completed" | "failed" | "cancelled") { + let content = serde_json::to_string(&WaitAgentJobResult { + status, + timed_out: false, + }) + .map_err(|err| { + FunctionCallError::Fatal(format!( + "failed to serialize wait_agent_job result: {err}" + )) + })?; + return Ok(ToolOutput::Function { + body: FunctionCallOutputBody::Text(content), + success: Some(true), + }); + } + if Instant::now() >= deadline { + let content = serde_json::to_string(&WaitAgentJobResult { + status, + timed_out: true, + }) + .map_err(|err| { + FunctionCallError::Fatal(format!( + "failed to serialize wait_agent_job timeout result: {err}" + )) + })?; + return Ok(ToolOutput::Function { + body: FunctionCallOutputBody::Text(content), + success: Some(true), + }); + } + tokio::time::sleep(Duration::from_millis(STATUS_POLL_INTERVAL_MS)).await; + } + } +} + +mod export_agent_job_csv { + use super::*; + + pub async fn handle( + session: Arc, + turn: Arc, + arguments: String, + ) -> Result { + let args: ExportAgentJobCsvArgs = parse_arguments(arguments.as_str())?; + let job_id = args.job_id; + let db = required_state_db(&session)?; + let job = db + .get_agent_job(job_id.as_str()) + .await + .map_err(|err| { + FunctionCallError::RespondToModel(format!( + "failed to look up agent job {job_id}: {err}" + )) + })? + .ok_or_else(|| { + FunctionCallError::RespondToModel(format!("agent job {job_id} not found")) + })?; + let items = db + .list_agent_job_items(job_id.as_str(), None, None) + .await + .map_err(|err| { + FunctionCallError::RespondToModel(format!( + "failed to load items for agent job {job_id}: {err}" + )) + })?; + let output_path = args.path.map_or_else( + || PathBuf::from(job.output_csv_path.clone()), + |path| turn.resolve_path(Some(path)), + ); + if let Some(parent) = output_path.parent() { + let parent_display = parent.display().to_string(); + tokio::fs::create_dir_all(parent).await.map_err(|err| { + FunctionCallError::RespondToModel(format!( + "failed to create export directory {parent_display}: {err}" + )) + })?; + } + let csv_content = render_job_csv(job.input_headers.as_slice(), items.as_slice())?; + let output_display = output_path.display().to_string(); + tokio::fs::write(&output_path, csv_content) + .await + .map_err(|err| { + FunctionCallError::RespondToModel(format!( + "failed to write csv export {output_display}: {err}" + )) + })?; + let content = serde_json::to_string(&ExportAgentJobCsvResult { + job_id, + path: output_display, + row_count: items.len(), + }) + .map_err(|err| { + FunctionCallError::Fatal(format!( + "failed to serialize export_agent_job_csv result: {err}" + )) + })?; + Ok(ToolOutput::Function { + body: FunctionCallOutputBody::Text(content), + success: Some(true), + }) + } +} + +mod report_agent_job_result { + use super::*; + + pub async fn handle( + session: Arc, + arguments: String, + ) -> Result { + let args: ReportAgentJobResultArgs = parse_arguments(arguments.as_str())?; + if !args.result.is_object() { + return Err(FunctionCallError::RespondToModel( + "result must be a JSON object".to_string(), + )); + } + let db = required_state_db(&session)?; + let reporting_thread_id = session.conversation_id.to_string(); + let accepted = db + .report_agent_job_item_result( + args.job_id.as_str(), + args.item_id.as_str(), + reporting_thread_id.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}" + )) + })?; + let content = + serde_json::to_string(&ReportAgentJobResultToolResult { accepted }).map_err(|err| { + FunctionCallError::Fatal(format!( + "failed to serialize report_agent_job_result result: {err}" + )) + })?; + Ok(ToolOutput::Function { + body: FunctionCallOutputBody::Text(content), + success: Some(true), + }) + } +} + +fn required_state_db( + session: &Arc, +) -> Result, FunctionCallError> { + session.state_db().ok_or_else(|| { + FunctionCallError::RespondToModel( + "sqlite state db is unavailable for this session; enable the sqlite feature" + .to_string(), + ) + }) +} + +async fn build_runner_options( + session: &Arc, + turn: &Arc, + requested_concurrency: Option, +) -> Result { + let session_source = turn.session_source.clone(); + let child_depth = next_thread_spawn_depth(&session_source); + if exceeds_thread_spawn_depth_limit(child_depth) { + return Err(FunctionCallError::RespondToModel( + "agent depth limit reached; this session cannot spawn more subagents".to_string(), + )); + } + let max_concurrency = + normalize_concurrency(requested_concurrency, turn.config.agent_max_threads); + let base_instructions = session.get_base_instructions().await; + let spawn_config = build_agent_spawn_config(&base_instructions, turn.as_ref(), child_depth)?; + Ok(JobRunnerOptions { + max_concurrency, + spawn_config, + child_depth, + }) +} + +fn normalize_concurrency(requested: Option, max_threads: Option) -> usize { + let requested = requested.unwrap_or(4).max(1); + let requested = requested.min(64); + if let Some(max_threads) = max_threads { + requested.min(max_threads.max(1)) + } else { + requested + } +} + +async fn start_job_runner( + session: Arc, + job_id: String, + options: JobRunnerOptions, +) -> Result { + cleanup_finished_runners().await; + let mut runners = ACTIVE_JOB_RUNNERS.lock().await; + if let Some(handle) = runners.get(job_id.as_str()) + && !handle.is_finished() + { + return Ok(false); + } + let db = required_state_db(&session)?; + 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 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}" + )) + })?; + if job.status.is_final() && progress.pending_items == 0 && progress.running_items == 0 { + return Ok(false); + } + 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}" + )) + })?; + let job_id_for_task = job_id.clone(); + let handle = tokio::spawn(async move { + if let Err(err) = + run_agent_job_loop(session, db.clone(), job_id_for_task.clone(), options).await + { + let error_message = format!("job runner failed: {err}"); + let _ = db + .mark_agent_job_failed(job_id_for_task.as_str(), error_message.as_str()) + .await; + } + }); + runners.insert(job_id, handle); + Ok(true) +} + +async fn cleanup_finished_runners() { + let mut runners = ACTIVE_JOB_RUNNERS.lock().await; + runners.retain(|_, handle| !handle.is_finished()); +} + +async fn run_agent_job_loop( + session: Arc, + db: Arc, + job_id: String, + options: JobRunnerOptions, +) -> anyhow::Result<()> { + let mut active_items: HashMap = HashMap::new(); + recover_running_items( + session.clone(), + db.clone(), + job_id.as_str(), + &mut active_items, + ) + .await?; + + let job = db + .get_agent_job(job_id.as_str()) + .await? + .ok_or_else(|| anyhow::anyhow!("agent job {job_id} was not found"))?; + + loop { + let mut progressed = false; + + if active_items.len() < options.max_concurrency { + let slots = options.max_concurrency - active_items.len(); + let pending_items = db + .list_agent_job_items( + job_id.as_str(), + Some(codex_state::AgentJobItemStatus::Pending), + Some(slots), + ) + .await?; + for item in pending_items { + if !db + .mark_agent_job_item_running(job_id.as_str(), item.item_id.as_str()) + .await? + { + continue; + } + let prompt = build_worker_prompt(&job, &item)?; + let thread_id = match session + .services + .agent_control + .spawn_agent( + options.spawn_config.clone(), + prompt, + Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id: session.conversation_id, + depth: options.child_depth, + })), + ) + .await + { + Ok(thread_id) => thread_id, + Err(CodexErr::AgentLimitReached { .. }) => { + break; + } + Err(err) => { + let error_message = format!("failed to spawn worker: {err}"); + db.mark_agent_job_item_failed( + job_id.as_str(), + item.item_id.as_str(), + error_message.as_str(), + ) + .await?; + progressed = true; + continue; + } + }; + let assigned = db + .set_agent_job_item_thread( + job_id.as_str(), + item.item_id.as_str(), + thread_id.to_string().as_str(), + ) + .await?; + if !assigned { + db.mark_agent_job_item_failed( + job_id.as_str(), + item.item_id.as_str(), + "failed to assign worker thread to job item", + ) + .await?; + let _ = session + .services + .agent_control + .shutdown_agent(thread_id) + .await; + progressed = true; + continue; + } + active_items.insert(thread_id, item.item_id.clone()); + progressed = true; + } + } + + 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 progress.pending_items == 0 && progress.running_items == 0 && active_items.is_empty() + { + break; + } + if !progressed { + tokio::time::sleep(Duration::from_millis(STATUS_POLL_INTERVAL_MS)).await; + } + continue; + } + + for (thread_id, item_id) in finished { + finalize_finished_item( + session.clone(), + db.clone(), + job_id.as_str(), + item_id.as_str(), + thread_id, + ) + .await?; + active_items.remove(&thread_id); + } + } + + let progress = db.get_agent_job_progress(job_id.as_str()).await?; + if progress.failed_items > 0 { + let failed_items = progress.failed_items; + let message = format!("job completed with {failed_items} failed items"); + db.mark_agent_job_failed(job_id.as_str(), message.as_str()) + .await?; + } else { + db.mark_agent_job_completed(job_id.as_str()).await?; + } + Ok(()) +} + +async fn recover_running_items( + session: Arc, + db: Arc, + job_id: &str, + active_items: &mut HashMap, +) -> anyhow::Result<()> { + let running_items = db + .list_agent_job_items(job_id, Some(codex_state::AgentJobItemStatus::Running), None) + .await?; + for item in running_items { + let Some(assigned_thread_id) = item.assigned_thread_id.clone() else { + db.mark_agent_job_item_failed( + job_id, + item.item_id.as_str(), + "running item is missing assigned_thread_id", + ) + .await?; + continue; + }; + let thread_id = match ThreadId::from_string(assigned_thread_id.as_str()) { + 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(), + ) + .await?; + continue; + } + }; + if is_final(&session.services.agent_control.get_status(thread_id).await) { + finalize_finished_item( + session.clone(), + db.clone(), + job_id, + item.item_id.as_str(), + thread_id, + ) + .await?; + } else { + active_items.insert(thread_id, item.item_id.clone()); + } + } + Ok(()) +} + +async fn find_finished_threads( + session: Arc, + active_items: &HashMap, +) -> Vec<(ThreadId, String)> { + let mut finished = Vec::new(); + for (thread_id, item_id) in active_items { + if is_final(&session.services.agent_control.get_status(*thread_id).await) { + finished.push((*thread_id, item_id.clone())); + } + } + finished +} + +async fn finalize_finished_item( + session: Arc, + db: Arc, + job_id: &str, + item_id: &str, + thread_id: ThreadId, +) -> anyhow::Result<()> { + let mut 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}") + })?; + if item.result_json.is_none() { + tokio::time::sleep(Duration::from_millis(250)).await; + item = db + .get_agent_job_item(job_id, item_id) + .await? + .ok_or_else(|| { + anyhow::anyhow!("job item not found after grace period: {job_id}/{item_id}") + })?; + } + if item.result_json.is_some() { + if !db.mark_agent_job_item_completed(job_id, item_id).await? { + db.mark_agent_job_item_failed( + job_id, + item_id, + "worker reported result but item could not transition to completed", + ) + .await?; + } + } else { + db.mark_agent_job_item_failed( + job_id, + item_id, + "worker finished without calling report_agent_job_result", + ) + .await?; + } + let _ = session + .services + .agent_control + .shutdown_agent(thread_id) + .await; + Ok(()) +} + +fn build_worker_prompt( + job: &codex_state::AgentJob, + item: &codex_state::AgentJobItem, +) -> anyhow::Result { + let job_id = job.id.as_str(); + let item_id = item.item_id.as_str(); + let instruction = job.instruction.as_str(); + let output_schema = job + .output_schema_json + .as_ref() + .map(serde_json::to_string_pretty) + .transpose()? + .unwrap_or_else(|| "{}".to_string()); + let row_json = serde_json::to_string_pretty(&item.row_json)?; + Ok(format!( + "You are processing one item for a generic agent job.\n\ +Job ID: {job_id}\n\ +Item ID: {item_id}\n\n\ +Task instruction:\n\ +{instruction}\n\n\ +Input row (JSON):\n\ +{row_json}\n\n\ +Expected result schema (JSON Schema or {{}}):\n\ +{output_schema}\n\n\ +You MUST call the `report_agent_job_result` tool exactly once with:\n\ +1. `job_id` = \"{job_id}\"\n\ +2. `item_id` = \"{item_id}\"\n\ +3. `result` = a JSON object that contains your analysis result for this row.\n\n\ +After the tool call succeeds, stop.", + )) +} + +async fn render_job_status( + db: Arc, + job_id: &str, +) -> Result { + let job = db + .get_agent_job(job_id) + .await + .map_err(|err| { + FunctionCallError::RespondToModel(format!("failed to fetch agent job {job_id}: {err}")) + })? + .ok_or_else(|| { + FunctionCallError::RespondToModel(format!("agent job {job_id} not found")) + })?; + let progress = db.get_agent_job_progress(job_id).await.map_err(|err| { + FunctionCallError::RespondToModel(format!("failed to fetch progress for {job_id}: {err}")) + })?; + cleanup_finished_runners().await; + let runners = ACTIVE_JOB_RUNNERS.lock().await; + let runner_active = runners + .get(job_id) + .is_some_and(|handle| !handle.is_finished()); + Ok(AgentJobToolResult { + job_id: job.id, + status: job.status.as_str().to_string(), + total_items: progress.total_items, + pending_items: progress.pending_items, + running_items: progress.running_items, + completed_items: progress.completed_items, + failed_items: progress.failed_items, + output_csv_path: job.output_csv_path, + runner_active, + }) +} + +fn default_output_csv_path(input_csv_path: &Path, job_id: &str) -> PathBuf { + let stem = input_csv_path + .file_stem() + .and_then(|stem| stem.to_str()) + .unwrap_or("agent_job_output"); + let job_suffix = &job_id[..8]; + input_csv_path.with_file_name(format!("{stem}.agent-job-{job_suffix}.csv")) +} + +fn parse_csv(content: &str) -> Result<(Vec, Vec>), String> { + let mut rows: Vec> = Vec::new(); + let mut row: Vec = Vec::new(); + let mut field = String::new(); + let mut in_quotes = false; + let mut chars = content.chars().peekable(); + while let Some(ch) = chars.next() { + match ch { + '"' => { + if in_quotes { + if chars.peek().is_some_and(|next| *next == '"') { + field.push('"'); + let _ = chars.next(); + } else { + in_quotes = false; + } + } else { + in_quotes = true; + } + } + ',' if !in_quotes => { + row.push(std::mem::take(&mut field)); + } + '\n' if !in_quotes => { + row.push(std::mem::take(&mut field)); + rows.push(std::mem::take(&mut row)); + } + '\r' if !in_quotes => { + if chars.peek().is_some_and(|next| *next == '\n') { + continue; + } + row.push(std::mem::take(&mut field)); + rows.push(std::mem::take(&mut row)); + } + other => field.push(other), + } + } + if in_quotes { + return Err("unterminated quoted field".to_string()); + } + if !field.is_empty() || !row.is_empty() { + row.push(field); + rows.push(row); + } + if rows.is_empty() { + return Ok((Vec::new(), Vec::new())); + } + let mut headers = rows.remove(0); + if let Some(first) = headers.first_mut() { + *first = first.trim_start_matches('\u{feff}').to_string(); + } + let data_rows = rows + .into_iter() + .filter(|row| row.iter().any(|value| !value.is_empty())) + .collect(); + Ok((headers, data_rows)) +} + +fn render_job_csv( + headers: &[String], + items: &[codex_state::AgentJobItem], +) -> Result { + let mut csv = String::new(); + let mut output_headers = headers.to_vec(); + output_headers.extend([ + "job_id".to_string(), + "item_id".to_string(), + "row_index".to_string(), + "source_id".to_string(), + "status".to_string(), + "attempt_count".to_string(), + "last_error".to_string(), + "result_json".to_string(), + "reported_at".to_string(), + "completed_at".to_string(), + ]); + csv.push_str( + output_headers + .iter() + .map(|header| csv_escape(header.as_str())) + .collect::>() + .join(",") + .as_str(), + ); + csv.push('\n'); + for item in items { + let row_object = item.row_json.as_object().ok_or_else(|| { + let item_id = item.item_id.as_str(); + FunctionCallError::RespondToModel(format!( + "row_json for item {item_id} is not a JSON object" + )) + })?; + let mut row_values = Vec::new(); + for header in headers { + let value = row_object + .get(header) + .map_or_else(String::new, value_to_csv_string); + row_values.push(csv_escape(value.as_str())); + } + row_values.push(csv_escape(item.job_id.as_str())); + row_values.push(csv_escape(item.item_id.as_str())); + row_values.push(csv_escape(item.row_index.to_string().as_str())); + row_values.push(csv_escape( + item.source_id.clone().unwrap_or_default().as_str(), + )); + row_values.push(csv_escape(item.status.as_str())); + row_values.push(csv_escape(item.attempt_count.to_string().as_str())); + row_values.push(csv_escape( + item.last_error.clone().unwrap_or_default().as_str(), + )); + row_values.push(csv_escape( + item.result_json + .as_ref() + .map_or_else(String::new, std::string::ToString::to_string) + .as_str(), + )); + row_values.push(csv_escape( + item.reported_at + .map(|value| value.to_rfc3339()) + .unwrap_or_default() + .as_str(), + )); + row_values.push(csv_escape( + item.completed_at + .map(|value| value.to_rfc3339()) + .unwrap_or_default() + .as_str(), + )); + csv.push_str(row_values.join(",").as_str()); + csv.push('\n'); + } + Ok(csv) +} + +fn value_to_csv_string(value: &Value) -> String { + match value { + Value::Null => String::new(), + Value::String(s) => s.clone(), + Value::Bool(b) => b.to_string(), + Value::Number(n) => n.to_string(), + Value::Array(_) | Value::Object(_) => value.to_string(), + } +} + +fn csv_escape(value: &str) -> String { + if value.contains(',') || value.contains('\n') || value.contains('\r') || value.contains('"') { + let escaped = value.replace('"', "\"\""); + format!("\"{escaped}\"") + } else { + value.to_string() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use pretty_assertions::assert_eq; + + #[test] + fn parse_csv_supports_quotes_and_commas() { + let input = "id,name\n1,\"alpha, beta\"\n2,gamma\n"; + let (headers, rows) = parse_csv(input).expect("csv parse"); + assert_eq!(headers, vec!["id".to_string(), "name".to_string()]); + assert_eq!( + rows, + vec![ + vec!["1".to_string(), "alpha, beta".to_string()], + vec!["2".to_string(), "gamma".to_string()] + ] + ); + } + + #[test] + fn csv_escape_quotes_when_needed() { + assert_eq!(csv_escape("simple"), "simple"); + assert_eq!(csv_escape("a,b"), "\"a,b\""); + assert_eq!(csv_escape("a\"b"), "\"a\"\"b\""); + } +} diff --git a/codex-rs/core/src/tools/handlers/collab.rs b/codex-rs/core/src/tools/handlers/collab.rs index 3110204d29..8b81ed87d9 100644 --- a/codex-rs/core/src/tools/handlers/collab.rs +++ b/codex-rs/core/src/tools/handlers/collab.rs @@ -1,10 +1,8 @@ use crate::agent::AgentStatus; -use crate::agent::exceeds_thread_spawn_depth_limit; use crate::codex::Session; use crate::codex::TurnContext; use crate::config::Config; use crate::error::CodexErr; -use crate::features::Feature; use crate::function_tool::FunctionCallError; use crate::tools::context::ToolInvocation; use crate::tools::context::ToolOutput; @@ -585,10 +583,10 @@ fn collab_agent_error(agent_id: ThreadId, err: CodexErr) -> FunctionCallError { } } -fn build_agent_spawn_config( +pub(crate) fn build_agent_spawn_config( base_instructions: &BaseInstructions, turn: &TurnContext, - child_depth: i32, + _child_depth: i32, ) -> Result { let base_config = turn.config.clone(); let mut config = (*base_config).clone(); @@ -615,11 +613,6 @@ fn build_agent_spawn_config( FunctionCallError::RespondToModel(format!("sandbox_policy is invalid: {err}")) })?; - // If the new agent will be at max depth: - if exceeds_thread_spawn_depth_limit(child_depth + 1) { - config.features.disable(Feature::Collab); - } - Ok(config) } diff --git a/codex-rs/core/src/tools/handlers/mod.rs b/codex-rs/core/src/tools/handlers/mod.rs index dda4760bd7..b6a850bc94 100644 --- a/codex-rs/core/src/tools/handlers/mod.rs +++ b/codex-rs/core/src/tools/handlers/mod.rs @@ -1,3 +1,4 @@ +pub(crate) mod agent_jobs; pub mod apply_patch; pub(crate) mod collab; mod dynamic; diff --git a/codex-rs/core/src/tools/spec.rs b/codex-rs/core/src/tools/spec.rs index ddac191a9b..4383921bb4 100644 --- a/codex-rs/core/src/tools/spec.rs +++ b/codex-rs/core/src/tools/spec.rs @@ -4,6 +4,7 @@ use crate::client_common::tools::ToolSpec; use crate::features::Feature; use crate::features::Features; use crate::tools::handlers::PLAN_TOOL; +use crate::tools::handlers::agent_jobs::AgentJobsHandler; use crate::tools::handlers::apply_patch::create_apply_patch_freeform_tool; use crate::tools::handlers::apply_patch::create_apply_patch_json_tool; use crate::tools::handlers::collab::DEFAULT_WAIT_TIMEOUT_MS; @@ -481,6 +482,208 @@ fn create_spawn_agent_tool() -> ToolSpec { }) } +fn create_spawn_agents_on_csv_tool() -> ToolSpec { + let mut properties = BTreeMap::new(); + properties.insert( + "csv_path".to_string(), + JsonSchema::String { + description: Some("Path to the CSV file containing input rows.".to_string()), + }, + ); + properties.insert( + "instruction".to_string(), + JsonSchema::String { + description: Some("Instruction to apply to each CSV row.".to_string()), + }, + ); + properties.insert( + "id_column".to_string(), + JsonSchema::String { + description: Some("Optional column name to use as stable item id.".to_string()), + }, + ); + properties.insert( + "job_name".to_string(), + JsonSchema::String { + description: Some("Optional friendly name for the job.".to_string()), + }, + ); + properties.insert( + "output_csv_path".to_string(), + JsonSchema::String { + description: Some("Optional output CSV path for exported results.".to_string()), + }, + ); + properties.insert( + "output_schema".to_string(), + JsonSchema::Object { + properties: BTreeMap::new(), + required: None, + additional_properties: None, + }, + ); + properties.insert( + "max_concurrency".to_string(), + JsonSchema::Number { + description: Some( + "Maximum concurrent workers for this job. Defaults to a safe value capped by config." + .to_string(), + ), + }, + ); + ToolSpec::Function(ResponsesApiTool { + name: "spawn_agents_on_csv".to_string(), + description: "Create and run a batch agent job over a CSV file.".to_string(), + strict: false, + parameters: JsonSchema::Object { + properties, + required: Some(vec!["csv_path".to_string(), "instruction".to_string()]), + additional_properties: Some(false.into()), + }, + }) +} + +fn create_run_agent_job_tool() -> ToolSpec { + let mut properties = BTreeMap::new(); + properties.insert( + "job_id".to_string(), + JsonSchema::String { + description: Some("Identifier of the job to start.".to_string()), + }, + ); + properties.insert( + "max_concurrency".to_string(), + JsonSchema::Number { + description: Some( + "Maximum concurrent workers for this job. Defaults to a safe value capped by config." + .to_string(), + ), + }, + ); + ToolSpec::Function(ResponsesApiTool { + name: "run_agent_job".to_string(), + description: "Start or resume execution of an existing agent job.".to_string(), + strict: false, + parameters: JsonSchema::Object { + properties, + required: Some(vec!["job_id".to_string()]), + additional_properties: Some(false.into()), + }, + }) +} + +fn create_get_agent_job_status_tool() -> ToolSpec { + let mut properties = BTreeMap::new(); + properties.insert( + "job_id".to_string(), + JsonSchema::String { + description: Some("Identifier of the job to inspect.".to_string()), + }, + ); + ToolSpec::Function(ResponsesApiTool { + name: "get_agent_job_status".to_string(), + description: "Fetch job status and progress counters.".to_string(), + strict: false, + parameters: JsonSchema::Object { + properties, + required: Some(vec!["job_id".to_string()]), + additional_properties: Some(false.into()), + }, + }) +} + +fn create_wait_agent_job_tool() -> ToolSpec { + let mut properties = BTreeMap::new(); + properties.insert( + "job_id".to_string(), + JsonSchema::String { + description: Some("Identifier of the job to wait on.".to_string()), + }, + ); + properties.insert( + "timeout_ms".to_string(), + JsonSchema::Number { + description: Some( + "Maximum time in milliseconds to wait for job completion.".to_string(), + ), + }, + ); + ToolSpec::Function(ResponsesApiTool { + name: "wait_agent_job".to_string(), + description: "Wait for an agent job to complete or time out.".to_string(), + strict: false, + parameters: JsonSchema::Object { + properties, + required: Some(vec!["job_id".to_string()]), + additional_properties: Some(false.into()), + }, + }) +} + +fn create_export_agent_job_csv_tool() -> ToolSpec { + let mut properties = BTreeMap::new(); + properties.insert( + "job_id".to_string(), + JsonSchema::String { + description: Some("Identifier of the job to export.".to_string()), + }, + ); + properties.insert( + "path".to_string(), + JsonSchema::String { + description: Some("Optional output CSV path override.".to_string()), + }, + ); + ToolSpec::Function(ResponsesApiTool { + name: "export_agent_job_csv".to_string(), + description: "Export job results to a CSV file.".to_string(), + strict: false, + parameters: JsonSchema::Object { + properties, + required: Some(vec!["job_id".to_string()]), + additional_properties: Some(false.into()), + }, + }) +} + +fn create_report_agent_job_result_tool() -> ToolSpec { + let mut properties = BTreeMap::new(); + properties.insert( + "job_id".to_string(), + JsonSchema::String { + description: Some("Identifier of the job.".to_string()), + }, + ); + properties.insert( + "item_id".to_string(), + JsonSchema::String { + description: Some("Identifier of the job item.".to_string()), + }, + ); + properties.insert( + "result".to_string(), + JsonSchema::Object { + properties: BTreeMap::new(), + required: None, + additional_properties: None, + }, + ); + ToolSpec::Function(ResponsesApiTool { + name: "report_agent_job_result".to_string(), + description: "Report a worker result for an agent job item.".to_string(), + strict: false, + parameters: JsonSchema::Object { + properties, + required: Some(vec![ + "job_id".to_string(), + "item_id".to_string(), + "result".to_string(), + ]), + additional_properties: Some(false.into()), + }, + }) +} + fn create_send_input_tool() -> ToolSpec { let mut properties = BTreeMap::new(); properties.insert( @@ -1386,6 +1589,22 @@ pub(crate) fn build_specs( builder.register_handler("close_agent", collab_handler); } + if config.collab_tools { + let agent_jobs_handler = Arc::new(AgentJobsHandler); + builder.push_spec(create_spawn_agents_on_csv_tool()); + builder.push_spec(create_run_agent_job_tool()); + builder.push_spec(create_get_agent_job_status_tool()); + builder.push_spec(create_wait_agent_job_tool()); + builder.push_spec(create_export_agent_job_csv_tool()); + builder.push_spec(create_report_agent_job_result_tool()); + builder.register_handler("spawn_agents_on_csv", agent_jobs_handler.clone()); + builder.register_handler("run_agent_job", agent_jobs_handler.clone()); + builder.register_handler("get_agent_job_status", agent_jobs_handler.clone()); + builder.register_handler("wait_agent_job", agent_jobs_handler.clone()); + builder.register_handler("export_agent_job_csv", agent_jobs_handler.clone()); + builder.register_handler("report_agent_job_result", agent_jobs_handler); + } + if let Some(mcp_tools) = mcp_tools { let mut entries: Vec<(String, rmcp::model::Tool)> = mcp_tools.into_iter().collect(); entries.sort_by(|a, b| a.0.cmp(&b.0)); @@ -1638,7 +1857,18 @@ mod tests { let (tools, _) = build_specs(&tools_config, None, &[]).build(); assert_contains_tool_names( &tools, - &["spawn_agent", "send_input", "wait", "close_agent"], + &[ + "spawn_agent", + "send_input", + "wait", + "close_agent", + "spawn_agents_on_csv", + "run_agent_job", + "get_agent_job_status", + "wait_agent_job", + "export_agent_job_csv", + "report_agent_job_result", + ], ); } diff --git a/codex-rs/docs/agent_jobs.md b/codex-rs/docs/agent_jobs.md new file mode 100644 index 0000000000..81a8c09d69 --- /dev/null +++ b/codex-rs/docs/agent_jobs.md @@ -0,0 +1,70 @@ +# Agent Jobs + +This document describes the generic batch job engine used for large agentic workloads. +Agent jobs are designed to be: + +1. Resumable and durable via SQLite. +2. Bounded by configured concurrency and thread limits. +3. Observable via explicit status/progress tools. +4. Exportable to CSV at stage boundaries. + +## Tools + +All tools are function-style and gated by the `collab` feature. + +### `spawn_agents_on_csv` + +Create a new job from a CSV input and immediately start it. + +Required args: +- `csv_path`: path to the CSV file (first row is headers). +- `instruction`: instruction to apply to each row. + +Optional args: +- `id_column`: header column name to use as a stable item id. +- `job_name`: human-friendly label. +- `output_csv_path`: destination for CSV export (defaults to `.agent-job-.csv`). +- `output_schema`: JSON schema for result payloads (best-effort guidance). +- `max_concurrency`: cap on parallel workers. + +### `run_agent_job` + +Start or resume an existing job by id. + +### `get_agent_job_status` + +Return job status and progress counters. + +### `wait_agent_job` + +Wait for a job to complete, or return after a timeout. + +### `export_agent_job_csv` + +Export the current job results to CSV using the stored headers and results. + +### `report_agent_job_result` + +Workers must call this exactly once to report a JSON object for their assigned item. + +## Execution Model + +1. Jobs are stored in SQLite with per-item state (pending/running/completed/failed). +2. The job runner spawns subagents up to `max_concurrency`. +3. Each worker processes one item and calls `report_agent_job_result`. +4. The runner marks items completed after the worker finishes. +5. CSV export is generated by a single writer from the SQLite store. + +## CSV Output + +Exports include original input columns plus: +- `job_id` +- `item_id` +- `row_index` +- `source_id` +- `status` +- `attempt_count` +- `last_error` +- `result_json` +- `reported_at` +- `completed_at` diff --git a/codex-rs/state/migrations/0009_agent_jobs.sql b/codex-rs/state/migrations/0009_agent_jobs.sql new file mode 100644 index 0000000000..1ae1d72d00 --- /dev/null +++ b/codex-rs/state/migrations/0009_agent_jobs.sql @@ -0,0 +1,37 @@ +CREATE TABLE agent_jobs ( + id TEXT PRIMARY KEY, + name TEXT NOT NULL, + status TEXT NOT NULL, + instruction TEXT NOT NULL, + output_schema_json TEXT, + input_headers_json TEXT NOT NULL, + input_csv_path TEXT NOT NULL, + output_csv_path TEXT NOT NULL, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL, + started_at INTEGER, + completed_at INTEGER, + last_error TEXT +); + +CREATE TABLE agent_job_items ( + job_id TEXT NOT NULL, + item_id TEXT NOT NULL, + row_index INTEGER NOT NULL, + source_id TEXT, + row_json TEXT NOT NULL, + status TEXT NOT NULL, + assigned_thread_id TEXT, + attempt_count INTEGER NOT NULL DEFAULT 0, + result_json TEXT, + last_error TEXT, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL, + completed_at INTEGER, + reported_at INTEGER, + PRIMARY KEY (job_id, item_id), + FOREIGN KEY(job_id) REFERENCES agent_jobs(id) ON DELETE CASCADE +); + +CREATE INDEX idx_agent_jobs_status ON agent_jobs(status, updated_at DESC); +CREATE INDEX idx_agent_job_items_status ON agent_job_items(job_id, status, row_index ASC); diff --git a/codex-rs/state/src/lib.rs b/codex-rs/state/src/lib.rs index 1625554e29..099ed173c7 100644 --- a/codex-rs/state/src/lib.rs +++ b/codex-rs/state/src/lib.rs @@ -21,6 +21,13 @@ pub use runtime::StateRuntime; /// /// Most consumers should prefer [`StateRuntime`]. pub use extract::apply_rollout_item; +pub use model::AgentJob; +pub use model::AgentJobCreateParams; +pub use model::AgentJobItem; +pub use model::AgentJobItemCreateParams; +pub use model::AgentJobItemStatus; +pub use model::AgentJobProgress; +pub use model::AgentJobStatus; pub use model::Anchor; pub use model::BackfillState; pub use model::BackfillStats; diff --git a/codex-rs/state/src/model/agent_job.rs b/codex-rs/state/src/model/agent_job.rs new file mode 100644 index 0000000000..2976933c74 --- /dev/null +++ b/codex-rs/state/src/model/agent_job.rs @@ -0,0 +1,285 @@ +use anyhow::Result; +use chrono::DateTime; +use chrono::Utc; +use serde_json::Value; +use sqlx::Row; +use sqlx::sqlite::SqliteRow; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum AgentJobStatus { + Pending, + Running, + Completed, + Failed, + Cancelled, +} + +impl AgentJobStatus { + pub const fn as_str(self) -> &'static str { + match self { + AgentJobStatus::Pending => "pending", + AgentJobStatus::Running => "running", + AgentJobStatus::Completed => "completed", + AgentJobStatus::Failed => "failed", + AgentJobStatus::Cancelled => "cancelled", + } + } + + pub fn parse(value: &str) -> Result { + match value { + "pending" => Ok(Self::Pending), + "running" => Ok(Self::Running), + "completed" => Ok(Self::Completed), + "failed" => Ok(Self::Failed), + "cancelled" => Ok(Self::Cancelled), + _ => Err(anyhow::anyhow!("invalid agent job status: {value}")), + } + } + + pub fn is_final(self) -> bool { + matches!( + self, + AgentJobStatus::Completed | AgentJobStatus::Failed | AgentJobStatus::Cancelled + ) + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum AgentJobItemStatus { + Pending, + Running, + Completed, + Failed, +} + +impl AgentJobItemStatus { + pub const fn as_str(self) -> &'static str { + match self { + AgentJobItemStatus::Pending => "pending", + AgentJobItemStatus::Running => "running", + AgentJobItemStatus::Completed => "completed", + AgentJobItemStatus::Failed => "failed", + } + } + + pub fn parse(value: &str) -> Result { + match value { + "pending" => Ok(Self::Pending), + "running" => Ok(Self::Running), + "completed" => Ok(Self::Completed), + "failed" => Ok(Self::Failed), + _ => Err(anyhow::anyhow!("invalid agent job item status: {value}")), + } + } +} + +#[derive(Debug, Clone, PartialEq)] +pub struct AgentJob { + pub id: String, + pub name: String, + pub status: AgentJobStatus, + pub instruction: String, + pub output_schema_json: Option, + pub input_headers: Vec, + pub input_csv_path: String, + pub output_csv_path: String, + pub created_at: DateTime, + pub updated_at: DateTime, + pub started_at: Option>, + pub completed_at: Option>, + pub last_error: Option, +} + +#[derive(Debug, Clone, PartialEq)] +pub struct AgentJobItem { + pub job_id: String, + pub item_id: String, + pub row_index: i64, + pub source_id: Option, + pub row_json: Value, + pub status: AgentJobItemStatus, + pub assigned_thread_id: Option, + pub attempt_count: i64, + pub result_json: Option, + pub last_error: Option, + pub created_at: DateTime, + pub updated_at: DateTime, + pub completed_at: Option>, + pub reported_at: Option>, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AgentJobProgress { + pub total_items: usize, + pub pending_items: usize, + pub running_items: usize, + pub completed_items: usize, + pub failed_items: usize, +} + +#[derive(Debug, Clone)] +pub struct AgentJobCreateParams { + pub id: String, + pub name: String, + pub instruction: String, + pub output_schema_json: Option, + pub input_headers: Vec, + pub input_csv_path: String, + pub output_csv_path: String, +} + +#[derive(Debug, Clone)] +pub struct AgentJobItemCreateParams { + pub item_id: String, + pub row_index: i64, + pub source_id: Option, + pub row_json: Value, +} + +#[derive(Debug)] +pub(crate) struct AgentJobRow { + pub(crate) id: String, + pub(crate) name: String, + pub(crate) status: String, + pub(crate) instruction: String, + pub(crate) output_schema_json: Option, + pub(crate) input_headers_json: String, + pub(crate) input_csv_path: String, + pub(crate) output_csv_path: String, + pub(crate) created_at: i64, + pub(crate) updated_at: i64, + pub(crate) started_at: Option, + pub(crate) completed_at: Option, + pub(crate) last_error: Option, +} + +impl AgentJobRow { + pub(crate) fn try_from_row(row: &SqliteRow) -> Result { + Ok(Self { + id: row.try_get("id")?, + name: row.try_get("name")?, + status: row.try_get("status")?, + instruction: row.try_get("instruction")?, + output_schema_json: row.try_get("output_schema_json")?, + input_headers_json: row.try_get("input_headers_json")?, + input_csv_path: row.try_get("input_csv_path")?, + output_csv_path: row.try_get("output_csv_path")?, + created_at: row.try_get("created_at")?, + updated_at: row.try_get("updated_at")?, + started_at: row.try_get("started_at")?, + completed_at: row.try_get("completed_at")?, + last_error: row.try_get("last_error")?, + }) + } +} + +impl TryFrom for AgentJob { + type Error = anyhow::Error; + + fn try_from(value: AgentJobRow) -> Result { + let output_schema_json = value + .output_schema_json + .as_deref() + .map(serde_json::from_str) + .transpose()?; + let input_headers = serde_json::from_str(value.input_headers_json.as_str())?; + Ok(Self { + id: value.id, + name: value.name, + status: AgentJobStatus::parse(value.status.as_str())?, + instruction: value.instruction, + output_schema_json, + input_headers, + input_csv_path: value.input_csv_path, + output_csv_path: value.output_csv_path, + created_at: epoch_seconds_to_datetime(value.created_at)?, + updated_at: epoch_seconds_to_datetime(value.updated_at)?, + started_at: value + .started_at + .map(epoch_seconds_to_datetime) + .transpose()?, + completed_at: value + .completed_at + .map(epoch_seconds_to_datetime) + .transpose()?, + last_error: value.last_error, + }) + } +} + +#[derive(Debug)] +pub(crate) struct AgentJobItemRow { + pub(crate) job_id: String, + pub(crate) item_id: String, + pub(crate) row_index: i64, + pub(crate) source_id: Option, + pub(crate) row_json: String, + pub(crate) status: String, + pub(crate) assigned_thread_id: Option, + pub(crate) attempt_count: i64, + pub(crate) result_json: Option, + pub(crate) last_error: Option, + pub(crate) created_at: i64, + pub(crate) updated_at: i64, + pub(crate) completed_at: Option, + pub(crate) reported_at: Option, +} + +impl AgentJobItemRow { + pub(crate) fn try_from_row(row: &SqliteRow) -> Result { + Ok(Self { + job_id: row.try_get("job_id")?, + item_id: row.try_get("item_id")?, + row_index: row.try_get("row_index")?, + source_id: row.try_get("source_id")?, + row_json: row.try_get("row_json")?, + status: row.try_get("status")?, + assigned_thread_id: row.try_get("assigned_thread_id")?, + attempt_count: row.try_get("attempt_count")?, + result_json: row.try_get("result_json")?, + last_error: row.try_get("last_error")?, + created_at: row.try_get("created_at")?, + updated_at: row.try_get("updated_at")?, + completed_at: row.try_get("completed_at")?, + reported_at: row.try_get("reported_at")?, + }) + } +} + +impl TryFrom for AgentJobItem { + type Error = anyhow::Error; + + fn try_from(value: AgentJobItemRow) -> Result { + Ok(Self { + job_id: value.job_id, + item_id: value.item_id, + row_index: value.row_index, + source_id: value.source_id, + row_json: serde_json::from_str(value.row_json.as_str())?, + status: AgentJobItemStatus::parse(value.status.as_str())?, + assigned_thread_id: value.assigned_thread_id, + attempt_count: value.attempt_count, + result_json: value + .result_json + .as_deref() + .map(serde_json::from_str) + .transpose()?, + last_error: value.last_error, + created_at: epoch_seconds_to_datetime(value.created_at)?, + updated_at: epoch_seconds_to_datetime(value.updated_at)?, + completed_at: value + .completed_at + .map(epoch_seconds_to_datetime) + .transpose()?, + reported_at: value + .reported_at + .map(epoch_seconds_to_datetime) + .transpose()?, + }) + } +} + +fn epoch_seconds_to_datetime(secs: i64) -> Result> { + DateTime::::from_timestamp(secs, 0) + .ok_or_else(|| anyhow::anyhow!("invalid unix timestamp: {secs}")) +} diff --git a/codex-rs/state/src/model/mod.rs b/codex-rs/state/src/model/mod.rs index 6bec8875dc..82450ccd0d 100644 --- a/codex-rs/state/src/model/mod.rs +++ b/codex-rs/state/src/model/mod.rs @@ -1,8 +1,16 @@ +mod agent_job; mod backfill_state; mod log; mod thread_memory; mod thread_metadata; +pub use agent_job::AgentJob; +pub use agent_job::AgentJobCreateParams; +pub use agent_job::AgentJobItem; +pub use agent_job::AgentJobItemCreateParams; +pub use agent_job::AgentJobItemStatus; +pub use agent_job::AgentJobProgress; +pub use agent_job::AgentJobStatus; pub use backfill_state::BackfillState; pub use backfill_state::BackfillStatus; pub use log::LogEntry; @@ -17,6 +25,8 @@ pub use thread_metadata::ThreadMetadata; pub use thread_metadata::ThreadMetadataBuilder; pub use thread_metadata::ThreadsPage; +pub(crate) use agent_job::AgentJobItemRow; +pub(crate) use agent_job::AgentJobRow; pub(crate) use thread_memory::ThreadMemoryRow; pub(crate) use thread_metadata::ThreadRow; pub(crate) use thread_metadata::anchor_from_item; diff --git a/codex-rs/state/src/runtime.rs b/codex-rs/state/src/runtime.rs index ae28a02197..529be667d2 100644 --- a/codex-rs/state/src/runtime.rs +++ b/codex-rs/state/src/runtime.rs @@ -1,3 +1,10 @@ +use crate::AgentJob; +use crate::AgentJobCreateParams; +use crate::AgentJobItem; +use crate::AgentJobItemCreateParams; +use crate::AgentJobItemStatus; +use crate::AgentJobProgress; +use crate::AgentJobStatus; use crate::DB_ERROR_METRIC; use crate::LogEntry; use crate::LogQuery; @@ -9,6 +16,8 @@ use crate::ThreadMetadataBuilder; use crate::ThreadsPage; use crate::apply_rollout_item; use crate::migrations::MIGRATOR; +use crate::model::AgentJobItemRow; +use crate::model::AgentJobRow; use crate::model::ThreadMemoryRow; use crate::model::ThreadRow; use crate::model::anchor_from_item; @@ -712,6 +721,454 @@ ON CONFLICT(thread_id, position) DO NOTHING self.upsert_thread(&metadata).await } + pub async fn create_agent_job( + &self, + params: &AgentJobCreateParams, + items: &[AgentJobItemCreateParams], + ) -> anyhow::Result { + let now = Utc::now().timestamp(); + let input_headers_json = serde_json::to_string(¶ms.input_headers)?; + let output_schema_json = params + .output_schema_json + .as_ref() + .map(serde_json::to_string) + .transpose()?; + let mut tx = self.pool.begin().await?; + sqlx::query( + r#" +INSERT INTO agent_jobs ( + id, + name, + status, + instruction, + output_schema_json, + input_headers_json, + input_csv_path, + output_csv_path, + created_at, + updated_at, + started_at, + completed_at, + last_error +) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, NULL, NULL, NULL) + "#, + ) + .bind(params.id.as_str()) + .bind(params.name.as_str()) + .bind(AgentJobStatus::Pending.as_str()) + .bind(params.instruction.as_str()) + .bind(output_schema_json) + .bind(input_headers_json) + .bind(params.input_csv_path.as_str()) + .bind(params.output_csv_path.as_str()) + .bind(now) + .bind(now) + .execute(&mut *tx) + .await?; + + for item in items { + let row_json = serde_json::to_string(&item.row_json)?; + sqlx::query( + r#" +INSERT INTO agent_job_items ( + job_id, + item_id, + row_index, + source_id, + row_json, + status, + assigned_thread_id, + attempt_count, + result_json, + last_error, + created_at, + updated_at, + completed_at, + reported_at +) VALUES (?, ?, ?, ?, ?, ?, NULL, 0, NULL, NULL, ?, ?, NULL, NULL) + "#, + ) + .bind(params.id.as_str()) + .bind(item.item_id.as_str()) + .bind(item.row_index) + .bind(item.source_id.as_deref()) + .bind(row_json) + .bind(AgentJobItemStatus::Pending.as_str()) + .bind(now) + .bind(now) + .execute(&mut *tx) + .await?; + } + + tx.commit().await?; + + let job_id = params.id.as_str(); + self.get_agent_job(job_id) + .await? + .ok_or_else(|| anyhow::anyhow!("failed to load created agent job {job_id}")) + } + + pub async fn get_agent_job(&self, job_id: &str) -> anyhow::Result> { + let row = sqlx::query( + r#" +SELECT + id, + name, + status, + instruction, + output_schema_json, + input_headers_json, + input_csv_path, + output_csv_path, + created_at, + updated_at, + started_at, + completed_at, + last_error +FROM agent_jobs +WHERE id = ? + "#, + ) + .bind(job_id) + .fetch_optional(self.pool.as_ref()) + .await?; + row.map(|row| AgentJobRow::try_from_row(&row).and_then(AgentJob::try_from)) + .transpose() + } + + pub async fn list_agent_job_items( + &self, + job_id: &str, + status: Option, + limit: Option, + ) -> anyhow::Result> { + let mut builder = QueryBuilder::::new( + r#" +SELECT + job_id, + item_id, + row_index, + source_id, + row_json, + status, + assigned_thread_id, + attempt_count, + result_json, + last_error, + created_at, + updated_at, + completed_at, + reported_at +FROM agent_job_items +WHERE job_id = + "#, + ); + builder.push_bind(job_id); + if let Some(status) = status { + builder.push(" AND status = "); + builder.push_bind(status.as_str()); + } + builder.push(" ORDER BY row_index ASC"); + if let Some(limit) = limit { + builder.push(" LIMIT "); + builder.push_bind(limit as i64); + } + let rows = builder.build().fetch_all(self.pool.as_ref()).await?; + rows.into_iter() + .map(|row| AgentJobItemRow::try_from_row(&row).and_then(AgentJobItem::try_from)) + .collect() + } + + pub async fn get_agent_job_item( + &self, + job_id: &str, + item_id: &str, + ) -> anyhow::Result> { + let row = sqlx::query( + r#" +SELECT + job_id, + item_id, + row_index, + source_id, + row_json, + status, + assigned_thread_id, + attempt_count, + result_json, + last_error, + created_at, + updated_at, + completed_at, + reported_at +FROM agent_job_items +WHERE job_id = ? AND item_id = ? + "#, + ) + .bind(job_id) + .bind(item_id) + .fetch_optional(self.pool.as_ref()) + .await?; + row.map(|row| AgentJobItemRow::try_from_row(&row).and_then(AgentJobItem::try_from)) + .transpose() + } + + pub async fn mark_agent_job_running(&self, job_id: &str) -> anyhow::Result<()> { + let now = Utc::now().timestamp(); + sqlx::query( + r#" +UPDATE agent_jobs +SET + status = ?, + updated_at = ?, + started_at = COALESCE(started_at, ?), + completed_at = NULL, + last_error = NULL +WHERE id = ? + "#, + ) + .bind(AgentJobStatus::Running.as_str()) + .bind(now) + .bind(now) + .bind(job_id) + .execute(self.pool.as_ref()) + .await?; + Ok(()) + } + + pub async fn mark_agent_job_completed(&self, job_id: &str) -> anyhow::Result<()> { + let now = Utc::now().timestamp(); + sqlx::query( + r#" +UPDATE agent_jobs +SET status = ?, updated_at = ?, completed_at = ?, last_error = NULL +WHERE id = ? + "#, + ) + .bind(AgentJobStatus::Completed.as_str()) + .bind(now) + .bind(now) + .bind(job_id) + .execute(self.pool.as_ref()) + .await?; + Ok(()) + } + + pub async fn mark_agent_job_failed( + &self, + job_id: &str, + error_message: &str, + ) -> anyhow::Result<()> { + let now = Utc::now().timestamp(); + sqlx::query( + r#" +UPDATE agent_jobs +SET status = ?, updated_at = ?, completed_at = ?, last_error = ? +WHERE id = ? + "#, + ) + .bind(AgentJobStatus::Failed.as_str()) + .bind(now) + .bind(now) + .bind(error_message) + .bind(job_id) + .execute(self.pool.as_ref()) + .await?; + Ok(()) + } + + pub async fn mark_agent_job_item_running( + &self, + job_id: &str, + item_id: &str, + ) -> anyhow::Result { + let now = Utc::now().timestamp(); + let result = sqlx::query( + r#" +UPDATE agent_job_items +SET + status = ?, + assigned_thread_id = NULL, + attempt_count = attempt_count + 1, + updated_at = ?, + last_error = NULL +WHERE job_id = ? AND item_id = ? AND status = ? + "#, + ) + .bind(AgentJobItemStatus::Running.as_str()) + .bind(now) + .bind(job_id) + .bind(item_id) + .bind(AgentJobItemStatus::Pending.as_str()) + .execute(self.pool.as_ref()) + .await?; + Ok(result.rows_affected() > 0) + } + + pub async fn set_agent_job_item_thread( + &self, + job_id: &str, + item_id: &str, + thread_id: &str, + ) -> anyhow::Result { + let now = Utc::now().timestamp(); + let result = sqlx::query( + r#" +UPDATE agent_job_items +SET assigned_thread_id = ?, updated_at = ? +WHERE job_id = ? AND item_id = ? AND status = ? + "#, + ) + .bind(thread_id) + .bind(now) + .bind(job_id) + .bind(item_id) + .bind(AgentJobItemStatus::Running.as_str()) + .execute(self.pool.as_ref()) + .await?; + Ok(result.rows_affected() > 0) + } + + pub async fn report_agent_job_item_result( + &self, + job_id: &str, + item_id: &str, + reporting_thread_id: &str, + result_json: &Value, + ) -> anyhow::Result { + let now = Utc::now().timestamp(); + let serialized = serde_json::to_string(result_json)?; + let result = sqlx::query( + r#" +UPDATE agent_job_items +SET + result_json = ?, + reported_at = ?, + updated_at = ?, + assigned_thread_id = COALESCE(assigned_thread_id, ?), + last_error = NULL +WHERE + job_id = ? + AND item_id = ? + AND status = ? + AND (assigned_thread_id IS NULL OR assigned_thread_id = ?) + "#, + ) + .bind(serialized) + .bind(now) + .bind(now) + .bind(reporting_thread_id) + .bind(job_id) + .bind(item_id) + .bind(AgentJobItemStatus::Running.as_str()) + .bind(reporting_thread_id) + .execute(self.pool.as_ref()) + .await?; + Ok(result.rows_affected() > 0) + } + + pub async fn mark_agent_job_item_completed( + &self, + job_id: &str, + item_id: &str, + ) -> anyhow::Result { + let now = Utc::now().timestamp(); + let result = sqlx::query( + r#" +UPDATE agent_job_items +SET + status = ?, + completed_at = ?, + updated_at = ?, + assigned_thread_id = NULL +WHERE + job_id = ? + AND item_id = ? + AND status = ? + AND result_json IS NOT NULL + "#, + ) + .bind(AgentJobItemStatus::Completed.as_str()) + .bind(now) + .bind(now) + .bind(job_id) + .bind(item_id) + .bind(AgentJobItemStatus::Running.as_str()) + .execute(self.pool.as_ref()) + .await?; + Ok(result.rows_affected() > 0) + } + + pub async fn mark_agent_job_item_failed( + &self, + job_id: &str, + item_id: &str, + error_message: &str, + ) -> anyhow::Result { + let now = Utc::now().timestamp(); + let result = sqlx::query( + r#" +UPDATE agent_job_items +SET + status = ?, + completed_at = ?, + updated_at = ?, + last_error = ?, + assigned_thread_id = NULL +WHERE + job_id = ? + AND item_id = ? + AND status = ? + "#, + ) + .bind(AgentJobItemStatus::Failed.as_str()) + .bind(now) + .bind(now) + .bind(error_message) + .bind(job_id) + .bind(item_id) + .bind(AgentJobItemStatus::Running.as_str()) + .execute(self.pool.as_ref()) + .await?; + Ok(result.rows_affected() > 0) + } + + pub async fn get_agent_job_progress(&self, job_id: &str) -> anyhow::Result { + let row = sqlx::query( + r#" +SELECT + COUNT(*) AS total_items, + SUM(CASE WHEN status = ? THEN 1 ELSE 0 END) AS pending_items, + SUM(CASE WHEN status = ? THEN 1 ELSE 0 END) AS running_items, + SUM(CASE WHEN status = ? THEN 1 ELSE 0 END) AS completed_items, + SUM(CASE WHEN status = ? THEN 1 ELSE 0 END) AS failed_items +FROM agent_job_items +WHERE job_id = ? + "#, + ) + .bind(AgentJobItemStatus::Pending.as_str()) + .bind(AgentJobItemStatus::Running.as_str()) + .bind(AgentJobItemStatus::Completed.as_str()) + .bind(AgentJobItemStatus::Failed.as_str()) + .bind(job_id) + .fetch_one(self.pool.as_ref()) + .await?; + + let total_items: i64 = row.try_get("total_items")?; + let pending_items: Option = row.try_get("pending_items")?; + let running_items: Option = row.try_get("running_items")?; + let completed_items: Option = row.try_get("completed_items")?; + let failed_items: Option = row.try_get("failed_items")?; + Ok(AgentJobProgress { + total_items: usize::try_from(total_items).unwrap_or_default(), + pending_items: usize::try_from(pending_items.unwrap_or_default()).unwrap_or_default(), + running_items: usize::try_from(running_items.unwrap_or_default()).unwrap_or_default(), + completed_items: usize::try_from(completed_items.unwrap_or_default()) + .unwrap_or_default(), + failed_items: usize::try_from(failed_items.unwrap_or_default()).unwrap_or_default(), + }) + } + async fn ensure_backfill_state_row(&self) -> anyhow::Result<()> { sqlx::query( r#"