use super::*;use crate::model::AgentJobItemRow;impl StateRuntime { pub async fn create_agent_job( &self, params: &AgentJobCreateParams, items: &[AgentJobItemCreateParams], ) -> anyhow::Result<AgentJob> { 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 max_runtime_seconds = params .max_runtime_seconds .map(i64::try_from) .transpose() .map_err(|_| anyhow::anyhow!("invalid max_runtime_seconds value"))?; let mut tx = self.pool.begin().await?; sqlx::query( r#"INSERT INTO agent_jobs ( id, name, status, instruction, auto_export, max_runtime_seconds, 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(i64::from(params.auto_export)) .bind(max_runtime_seconds) .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<Option<AgentJob>> { let row = sqlx::query_as::<_, AgentJobRow>( r#"SELECT id, name, status, instruction, auto_export, max_runtime_seconds, output_schema_json, input_headers_json, input_csv_path, output_csv_path, created_at, updated_at, started_at, completed_at, last_errorFROM agent_jobsWHERE id = ? "#, ) .bind(job_id) .fetch_optional(self.pool.as_ref()) .await?; row.map(AgentJob::try_from).transpose() } pub async fn list_agent_job_items( &self, job_id: &str, status: Option<AgentJobItemStatus>, limit: Option<usize>, ) -> anyhow::Result<Vec<AgentJobItem>> { let mut builder = QueryBuilder::<Sqlite>::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_atFROM agent_job_itemsWHERE 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: Vec<AgentJobItemRow> = builder .build_query_as::<AgentJobItemRow>() .fetch_all(self.pool.as_ref()) .await?; rows.into_iter().map(AgentJobItem::try_from).collect() } pub async fn get_agent_job_item( &self, job_id: &str, item_id: &str, ) -> anyhow::Result<Option<AgentJobItem>> { let row: Option<AgentJobItemRow> = sqlx::query_as::<_, AgentJobItemRow>( 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_atFROM agent_job_itemsWHERE job_id = ? AND item_id = ? "#, ) .bind(job_id) .bind(item_id) .fetch_optional(self.pool.as_ref()) .await?; row.map(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_jobsSET status = ?, updated_at = ?, started_at = COALESCE(started_at, ?), completed_at = NULL, last_error = NULLWHERE 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_jobsSET status = ?, updated_at = ?, completed_at = ?, last_error = NULLWHERE 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_jobsSET 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_cancelled( &self, job_id: &str, reason: &str, ) -> anyhow::Result<bool> { let now = Utc::now().timestamp(); let result = sqlx::query( r#"UPDATE agent_jobsSET status = ?, updated_at = ?, completed_at = ?, last_error = ?WHERE id = ? AND status IN (?, ?) "#, ) .bind(AgentJobStatus::Cancelled.as_str()) .bind(now) .bind(now) .bind(reason) .bind(job_id) .bind(AgentJobStatus::Pending.as_str()) .bind(AgentJobStatus::Running.as_str()) .execute(self.pool.as_ref()) .await?; Ok(result.rows_affected() > 0) } pub async fn is_agent_job_cancelled(&self, job_id: &str) -> anyhow::Result<bool> { let row = sqlx::query( r#"SELECT statusFROM agent_jobsWHERE id = ? "#, ) .bind(job_id) .fetch_optional(self.pool.as_ref()) .await?; let Some(row) = row else { return Ok(false); }; let status: String = row.try_get("status")?; Ok(AgentJobStatus::parse(status.as_str())? == AgentJobStatus::Cancelled) } pub async fn mark_agent_job_item_running( &self, job_id: &str, item_id: &str, ) -> anyhow::Result<bool> { let now = Utc::now().timestamp(); let result = sqlx::query( r#"UPDATE agent_job_itemsSET status = ?, assigned_thread_id = NULL, attempt_count = attempt_count + 1, updated_at = ?, last_error = NULLWHERE 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 mark_agent_job_item_running_with_thread( &self, job_id: &str, item_id: &str, thread_id: &str, ) -> anyhow::Result<bool> { let now = Utc::now().timestamp(); let result = sqlx::query( r#"UPDATE agent_job_itemsSET status = ?, assigned_thread_id = ?, attempt_count = attempt_count + 1, updated_at = ?, last_error = NULLWHERE job_id = ? AND item_id = ? AND status = ? "#, ) .bind(AgentJobItemStatus::Running.as_str()) .bind(thread_id) .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 mark_agent_job_item_pending( &self, job_id: &str, item_id: &str, error_message: Option<&str>, ) -> anyhow::Result<bool> { let now = Utc::now().timestamp(); let result = sqlx::query( r#"UPDATE agent_job_itemsSET status = ?, assigned_thread_id = NULL, updated_at = ?, last_error = ?WHERE job_id = ? AND item_id = ? AND status = ? "#, ) .bind(AgentJobItemStatus::Pending.as_str()) .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 set_agent_job_item_thread( &self, job_id: &str, item_id: &str, thread_id: &str, ) -> anyhow::Result<bool> { let now = Utc::now().timestamp(); let result = sqlx::query( r#"UPDATE agent_job_itemsSET 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<bool> { let now = Utc::now().timestamp(); let serialized = serde_json::to_string(result_json)?; let result = sqlx::query( r#"UPDATE agent_job_itemsSET status = ?, result_json = ?, reported_at = ?, completed_at = ?, updated_at = ?, last_error = NULL, assigned_thread_id = NULLWHERE job_id = ? AND item_id = ? AND status = ? AND assigned_thread_id = ? "#, ) .bind(AgentJobItemStatus::Completed.as_str()) .bind(serialized) .bind(now) .bind(now) .bind(now) .bind(job_id) .bind(item_id) .bind(AgentJobItemStatus::Running.as_str()) .bind(reporting_thread_id) .execute(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<bool> { let now = Utc::now().timestamp(); let result = sqlx::query( r#"UPDATE agent_job_itemsSET status = ?, completed_at = ?, updated_at = ?, assigned_thread_id = NULLWHERE 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<bool> { let now = Utc::now().timestamp(); let result = sqlx::query( r#"UPDATE agent_job_itemsSET status = ?, completed_at = ?, updated_at = ?, last_error = ?, assigned_thread_id = NULLWHERE 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<AgentJobProgress> { 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_itemsFROM agent_job_itemsWHERE 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<i64> = row.try_get("pending_items")?; let running_items: Option<i64> = row.try_get("running_items")?; let completed_items: Option<i64> = row.try_get("completed_items")?; let failed_items: Option<i64> = 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(), }) }}#[cfg(test)]mod tests { use super::*; use crate::runtime::test_support::unique_temp_dir; use pretty_assertions::assert_eq; use serde_json::json; async fn create_running_single_item_job( runtime: &StateRuntime, ) -> anyhow::Result<(String, String, String)> { let job_id = "job-1".to_string(); let item_id = "item-1".to_string(); let thread_id = "thread-1".to_string(); runtime .create_agent_job( &AgentJobCreateParams { id: job_id.clone(), name: "test-job".to_string(), instruction: "Return a result".to_string(), auto_export: true, max_runtime_seconds: None, output_schema_json: None, input_headers: vec!["path".to_string()], input_csv_path: "/tmp/in.csv".to_string(), output_csv_path: "/tmp/out.csv".to_string(), }, &[AgentJobItemCreateParams { item_id: item_id.clone(), row_index: 0, source_id: None, row_json: json!({"path":"file-1"}), }], ) .await?; runtime.mark_agent_job_running(job_id.as_str()).await?; let marked_running = runtime .mark_agent_job_item_running_with_thread( job_id.as_str(), item_id.as_str(), thread_id.as_str(), ) .await?; assert!(marked_running); Ok((job_id, item_id, thread_id)) } #[tokio::test] async fn report_agent_job_item_result_completes_item_atomically() -> anyhow::Result<()> { let codex_home = unique_temp_dir(); let runtime = StateRuntime::init(codex_home, "test-provider".to_string()).await?; let (job_id, item_id, thread_id) = create_running_single_item_job(runtime.as_ref()).await?; let accepted = runtime .report_agent_job_item_result( job_id.as_str(), item_id.as_str(), thread_id.as_str(), &json!({"ok": true}), ) .await?; assert!(accepted); let item = runtime .get_agent_job_item(job_id.as_str(), item_id.as_str()) .await? .expect("job item should exist"); assert_eq!(item.status, AgentJobItemStatus::Completed); assert_eq!(item.result_json, Some(json!({"ok": true}))); assert_eq!(item.assigned_thread_id, None); assert_eq!(item.last_error, None); assert!(item.reported_at.is_some()); assert!(item.completed_at.is_some()); let progress = runtime.get_agent_job_progress(job_id.as_str()).await?; assert_eq!( progress, AgentJobProgress { total_items: 1, pending_items: 0, running_items: 0, completed_items: 1, failed_items: 0, } ); Ok(()) } #[tokio::test] async fn report_agent_job_item_result_rejects_late_reports() -> anyhow::Result<()> { let codex_home = unique_temp_dir(); let runtime = StateRuntime::init(codex_home, "test-provider".to_string()).await?; let (job_id, item_id, thread_id) = create_running_single_item_job(runtime.as_ref()).await?; let marked_failed = runtime .mark_agent_job_item_failed(job_id.as_str(), item_id.as_str(), "missing report") .await?; assert!(marked_failed); let accepted = runtime .report_agent_job_item_result( job_id.as_str(), item_id.as_str(), thread_id.as_str(), &json!({"late": true}), ) .await?; assert!(!accepted); let item = runtime .get_agent_job_item(job_id.as_str(), item_id.as_str()) .await? .expect("job item should exist"); assert_eq!(item.status, AgentJobItemStatus::Failed); assert_eq!(item.result_json, None); assert_eq!(item.last_error, Some("missing report".to_string())); Ok(()) }}