mirror of
https://github.com/pchuan98/codex.git
synced 2026-07-01 00:31:56 +08:00
Agent jobs (spawn_agents_on_csv) + progress UI (#10935)
## Summary - Add agent job support: spawn a batch of sub-agents from CSV, auto-run, auto-export, and store results in SQLite. - Simplify workflow: remove run/resume/get-status/export tools; spawn is deterministic and completes in one call. - Improve exec UX: stable, single-line progress bar with ETA; suppress sub-agent chatter in exec. ## Why Enables map-reduce style workflows over arbitrarily large repos using the existing Codex orchestrator. This addresses review feedback about overly complex job controls and non-deterministic monitoring. ## Demo (progress bar) ``` ./codex-rs/target/debug/codex exec \ --enable collab \ --enable sqlite \ --full-auto \ --progress-cursor \ -c agents.max_threads=16 \ -C /Users/daveaitel/code/codex \ - <<'PROMPT' Create /tmp/agent_job_progress_demo.csv with columns: path,area and 30 rows: path = item-01..item-30, area = test. Then call spawn_agents_on_csv with: - csv_path: /tmp/agent_job_progress_demo.csv - instruction: "Run `python - <<'PY'` to sleep a random 0.3–1.2s, then output JSON with keys: path, score (int). Set score = 1." - output_csv_path: /tmp/agent_job_progress_demo_out.csv PROMPT ``` ## Review feedback addressed - Auto-start jobs on spawn; removed run/resume/status/export tools. - Auto-export on success. - More descriptive tool spec + clearer prompts. - Avoid deadlocks on spawn failure; pending/running handled safely. - Progress bar no longer scrolls; stable single-line redraw. ## Tests - `cd codex-rs && cargo test -p codex-exec` - `cd codex-rs && cargo build -p codex-cli`
This commit is contained in:
@@ -22,6 +22,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;
|
||||
@@ -38,6 +45,9 @@ pub use model::ThreadsPage;
|
||||
pub use runtime::state_db_filename;
|
||||
pub use runtime::state_db_path;
|
||||
|
||||
/// Environment variable for overriding the SQLite state database home directory.
|
||||
pub const SQLITE_HOME_ENV: &str = "CODEX_SQLITE_HOME";
|
||||
|
||||
pub const STATE_DB_FILENAME: &str = "state";
|
||||
pub const STATE_DB_VERSION: u32 = 5;
|
||||
|
||||
|
||||
@@ -0,0 +1,256 @@
|
||||
use anyhow::Result;
|
||||
use chrono::DateTime;
|
||||
use chrono::Utc;
|
||||
use serde_json::Value;
|
||||
|
||||
#[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<Self> {
|
||||
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<Self> {
|
||||
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 auto_export: bool,
|
||||
pub max_runtime_seconds: Option<u64>,
|
||||
// TODO(jif-oai): Convert to JSON Schema and enforce structured outputs.
|
||||
pub output_schema_json: Option<Value>,
|
||||
pub input_headers: Vec<String>,
|
||||
pub input_csv_path: String,
|
||||
pub output_csv_path: String,
|
||||
pub created_at: DateTime<Utc>,
|
||||
pub updated_at: DateTime<Utc>,
|
||||
pub started_at: Option<DateTime<Utc>>,
|
||||
pub completed_at: Option<DateTime<Utc>>,
|
||||
pub last_error: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct AgentJobItem {
|
||||
pub job_id: String,
|
||||
pub item_id: String,
|
||||
pub row_index: i64,
|
||||
pub source_id: Option<String>,
|
||||
pub row_json: Value,
|
||||
pub status: AgentJobItemStatus,
|
||||
pub assigned_thread_id: Option<String>,
|
||||
pub attempt_count: i64,
|
||||
pub result_json: Option<Value>,
|
||||
pub last_error: Option<String>,
|
||||
pub created_at: DateTime<Utc>,
|
||||
pub updated_at: DateTime<Utc>,
|
||||
pub completed_at: Option<DateTime<Utc>>,
|
||||
pub reported_at: Option<DateTime<Utc>>,
|
||||
}
|
||||
|
||||
#[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 auto_export: bool,
|
||||
pub max_runtime_seconds: Option<u64>,
|
||||
pub output_schema_json: Option<Value>,
|
||||
pub input_headers: Vec<String>,
|
||||
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<String>,
|
||||
pub row_json: Value,
|
||||
}
|
||||
|
||||
#[derive(Debug, sqlx::FromRow)]
|
||||
pub(crate) struct AgentJobRow {
|
||||
pub(crate) id: String,
|
||||
pub(crate) name: String,
|
||||
pub(crate) status: String,
|
||||
pub(crate) instruction: String,
|
||||
pub(crate) auto_export: i64,
|
||||
pub(crate) max_runtime_seconds: Option<i64>,
|
||||
pub(crate) output_schema_json: Option<String>,
|
||||
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<i64>,
|
||||
pub(crate) completed_at: Option<i64>,
|
||||
pub(crate) last_error: Option<String>,
|
||||
}
|
||||
|
||||
impl TryFrom<AgentJobRow> for AgentJob {
|
||||
type Error = anyhow::Error;
|
||||
|
||||
fn try_from(value: AgentJobRow) -> Result<Self, Self::Error> {
|
||||
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())?;
|
||||
let max_runtime_seconds = value
|
||||
.max_runtime_seconds
|
||||
.map(u64::try_from)
|
||||
.transpose()
|
||||
.map_err(|_| anyhow::anyhow!("invalid max_runtime_seconds value"))?;
|
||||
Ok(Self {
|
||||
id: value.id,
|
||||
name: value.name,
|
||||
status: AgentJobStatus::parse(value.status.as_str())?,
|
||||
instruction: value.instruction,
|
||||
auto_export: value.auto_export != 0,
|
||||
max_runtime_seconds,
|
||||
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, sqlx::FromRow)]
|
||||
pub(crate) struct AgentJobItemRow {
|
||||
pub(crate) job_id: String,
|
||||
pub(crate) item_id: String,
|
||||
pub(crate) row_index: i64,
|
||||
pub(crate) source_id: Option<String>,
|
||||
pub(crate) row_json: String,
|
||||
pub(crate) status: String,
|
||||
pub(crate) assigned_thread_id: Option<String>,
|
||||
pub(crate) attempt_count: i64,
|
||||
pub(crate) result_json: Option<String>,
|
||||
pub(crate) last_error: Option<String>,
|
||||
pub(crate) created_at: i64,
|
||||
pub(crate) updated_at: i64,
|
||||
pub(crate) completed_at: Option<i64>,
|
||||
pub(crate) reported_at: Option<i64>,
|
||||
}
|
||||
|
||||
impl TryFrom<AgentJobItemRow> for AgentJobItem {
|
||||
type Error = anyhow::Error;
|
||||
|
||||
fn try_from(value: AgentJobItemRow) -> Result<Self, Self::Error> {
|
||||
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<Utc>> {
|
||||
DateTime::<Utc>::from_timestamp(secs, 0)
|
||||
.ok_or_else(|| anyhow::anyhow!("invalid unix timestamp: {secs}"))
|
||||
}
|
||||
@@ -1,8 +1,16 @@
|
||||
mod agent_job;
|
||||
mod backfill_state;
|
||||
mod log;
|
||||
mod memories;
|
||||
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;
|
||||
@@ -21,6 +29,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 memories::Stage1OutputRow;
|
||||
pub(crate) use thread_metadata::ThreadRow;
|
||||
pub(crate) use thread_metadata::anchor_from_item;
|
||||
|
||||
@@ -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;
|
||||
@@ -11,6 +18,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::ThreadRow;
|
||||
use crate::model::anchor_from_item;
|
||||
use crate::model::datetime_to_epoch_seconds;
|
||||
@@ -901,6 +910,564 @@ ON CONFLICT(thread_id, position) DO NOTHING
|
||||
Ok(result.rows_affected())
|
||||
}
|
||||
|
||||
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_error
|
||||
FROM agent_jobs
|
||||
WHERE 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_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_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 = 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_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(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_cancelled(
|
||||
&self,
|
||||
job_id: &str,
|
||||
reason: &str,
|
||||
) -> anyhow::Result<bool> {
|
||||
let now = Utc::now().timestamp();
|
||||
let result = sqlx::query(
|
||||
r#"
|
||||
UPDATE agent_jobs
|
||||
SET status = ?, updated_at = ?, completed_at = ?, last_error = ?
|
||||
WHERE id = ? AND status IN (?, ?)
|
||||
"#,
|
||||
)
|
||||
.bind(AgentJobStatus::Cancelled.as_str())
|
||||
.bind(now)
|
||||
.bind(now)
|
||||
.bind(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 status
|
||||
FROM agent_jobs
|
||||
WHERE 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_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 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_items
|
||||
SET
|
||||
status = ?,
|
||||
assigned_thread_id = ?,
|
||||
attempt_count = attempt_count + 1,
|
||||
updated_at = ?,
|
||||
last_error = NULL
|
||||
WHERE 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_items
|
||||
SET
|
||||
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_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<bool> {
|
||||
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 = ?,
|
||||
last_error = NULL
|
||||
WHERE
|
||||
job_id = ?
|
||||
AND item_id = ?
|
||||
AND status = ?
|
||||
AND assigned_thread_id = ?
|
||||
"#,
|
||||
)
|
||||
.bind(serialized)
|
||||
.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_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<bool> {
|
||||
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<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_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<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(),
|
||||
})
|
||||
}
|
||||
|
||||
async fn ensure_backfill_state_row(&self) -> anyhow::Result<()> {
|
||||
sqlx::query(
|
||||
r#"
|
||||
|
||||
Reference in New Issue
Block a user