Agent jobs: template instructions and resume safety

This commit is contained in:
Dave Aitel
2026-02-06 10:51:55 -05:00
parent 263cacc2bb
commit 6630587da3
4 changed files with 269 additions and 25 deletions
+81 -10
View File
@@ -77,7 +77,8 @@ struct WaitAgentJobArgs {
#[derive(Debug, Deserialize)]
struct ExportAgentJobCsvArgs {
job_id: String,
path: Option<String>,
#[serde(alias = "path")]
output_csv_path: Option<String>,
}
#[derive(Debug, Deserialize)]
@@ -117,7 +118,7 @@ struct WaitAgentJobResult {
#[derive(Debug, Serialize)]
struct ExportAgentJobCsvResult {
job_id: String,
path: String,
output_csv_path: String,
row_count: usize,
}
@@ -332,6 +333,15 @@ mod run_agent_job {
.ok_or_else(|| {
FunctionCallError::RespondToModel(format!("agent job {job_id} not found"))
})?;
if !job_runner_active(job_id.as_str()).await {
db.reset_agent_job_running_items(job_id.as_str())
.await
.map_err(|err| {
FunctionCallError::RespondToModel(format!(
"failed to reset running items for agent job {job_id}: {err}"
))
})?;
}
let options = build_runner_options(
&session,
&turn,
@@ -465,7 +475,7 @@ mod export_agent_job_csv {
"failed to load items for agent job {job_id}: {err}"
))
})?;
let output_path = args.path.map_or_else(
let output_path = args.output_csv_path.map_or_else(
|| PathBuf::from(job.output_csv_path.clone()),
|path| turn.resolve_path(Some(path)),
);
@@ -488,7 +498,7 @@ mod export_agent_job_csv {
})?;
let content = serde_json::to_string(&ExportAgentJobCsvResult {
job_id,
path: output_display,
output_csv_path: output_display,
row_count: items.len(),
})
.map_err(|err| {
@@ -652,6 +662,14 @@ async fn cleanup_finished_runners() {
runners.retain(|_, handle| !handle.is_finished());
}
async fn job_runner_active(job_id: &str) -> bool {
cleanup_finished_runners().await;
let runners = ACTIVE_JOB_RUNNERS.lock().await;
runners
.get(job_id)
.is_some_and(|handle| !handle.is_finished())
}
async fn run_agent_job_loop(
session: Arc<Session>,
db: Arc<codex_state::StateRuntime>,
@@ -707,6 +725,12 @@ async fn run_agent_job_loop(
{
Ok(thread_id) => thread_id,
Err(CodexErr::AgentLimitReached { .. }) => {
db.mark_agent_job_item_pending(
job_id.as_str(),
item.item_id.as_str(),
None,
)
.await?;
break;
}
Err(err) => {
@@ -923,7 +947,7 @@ fn build_worker_prompt(
) -> anyhow::Result<String> {
let job_id = job.id.as_str();
let item_id = item.item_id.as_str();
let instruction = job.instruction.as_str();
let instruction = render_instruction_template(job.instruction.as_str(), &item.row_json);
let output_schema = job
.output_schema_json
.as_ref()
@@ -949,6 +973,31 @@ After the tool call succeeds, stop.",
))
}
fn render_instruction_template(instruction: &str, row_json: &Value) -> String {
const OPEN_BRACE_SENTINEL: &str = "__CODEX_OPEN_BRACE__";
const CLOSE_BRACE_SENTINEL: &str = "__CODEX_CLOSE_BRACE__";
let mut rendered = instruction
.replace("{{", OPEN_BRACE_SENTINEL)
.replace("}}", CLOSE_BRACE_SENTINEL);
let Some(row) = row_json.as_object() else {
return rendered
.replace(OPEN_BRACE_SENTINEL, "{")
.replace(CLOSE_BRACE_SENTINEL, "}");
};
for (key, value) in row {
let placeholder = format!("{{{key}}}");
let replacement = value
.as_str()
.map(str::to_string)
.unwrap_or_else(|| value.to_string());
rendered = rendered.replace(placeholder.as_str(), replacement.as_str());
}
rendered
.replace(OPEN_BRACE_SENTINEL, "{")
.replace(CLOSE_BRACE_SENTINEL, "}")
}
async fn render_job_status(
db: Arc<codex_state::StateRuntime>,
job_id: &str,
@@ -965,11 +1014,7 @@ async fn render_job_status(
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());
let runner_active = job_runner_active(job_id).await;
Ok(AgentJobToolResult {
job_id: job.id,
status: job.status.as_str().to_string(),
@@ -1170,4 +1215,30 @@ mod tests {
assert_eq!(csv_escape("a,b"), "\"a,b\"");
assert_eq!(csv_escape("a\"b"), "\"a\"\"b\"");
}
#[test]
fn render_instruction_template_expands_placeholders_and_escapes_braces() {
let row = json!({
"path": "src/lib.rs",
"area": "test",
"file path": "docs/readme.md",
});
let rendered = render_instruction_template(
"Review {path} in {area}. Also see {file path}. Use {{literal}}.",
&row,
);
assert_eq!(
rendered,
"Review src/lib.rs in test. Also see docs/readme.md. Use {literal}."
);
}
#[test]
fn render_instruction_template_leaves_unknown_placeholders() {
let row = json!({
"path": "src/lib.rs",
});
let rendered = render_instruction_template("Check {path} then {missing}", &row);
assert_eq!(rendered, "Check src/lib.rs then {missing}");
}
}
+11 -7
View File
@@ -493,7 +493,10 @@ fn create_spawn_agents_on_csv_tool() -> ToolSpec {
properties.insert(
"instruction".to_string(),
JsonSchema::String {
description: Some("Instruction to apply to each CSV row.".to_string()),
description: Some(
"Instruction template to apply to each CSV row. Use {column_name} placeholders to inject values from the row."
.to_string(),
),
},
);
properties.insert(
@@ -526,7 +529,7 @@ fn create_spawn_agents_on_csv_tool() -> ToolSpec {
"max_concurrency".to_string(),
JsonSchema::Number {
description: Some(
"Maximum concurrent workers for this job. Defaults to a safe value capped by config."
"Maximum concurrent workers for this job. Defaults to 64 and is capped by config."
.to_string(),
),
},
@@ -542,7 +545,8 @@ fn create_spawn_agents_on_csv_tool() -> ToolSpec {
);
ToolSpec::Function(ResponsesApiTool {
name: "spawn_agents_on_csv".to_string(),
description: "Create and run a batch agent job over a CSV file.".to_string(),
description: "Create and start a job that spawns subagents to process each CSV row."
.to_string(),
strict: false,
parameters: JsonSchema::Object {
properties,
@@ -557,14 +561,14 @@ fn create_run_agent_job_tool() -> ToolSpec {
properties.insert(
"job_id".to_string(),
JsonSchema::String {
description: Some("Identifier of the job to start.".to_string()),
description: Some("Identifier of the job to resume.".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."
"Maximum concurrent workers for this job. Defaults to 64 and is capped by config."
.to_string(),
),
},
@@ -580,7 +584,7 @@ fn create_run_agent_job_tool() -> ToolSpec {
);
ToolSpec::Function(ResponsesApiTool {
name: "run_agent_job".to_string(),
description: "Start or resume execution of an existing agent job.".to_string(),
description: "Resume execution of an existing agent job.".to_string(),
strict: false,
parameters: JsonSchema::Object {
properties,
@@ -647,7 +651,7 @@ fn create_export_agent_job_csv_tool() -> ToolSpec {
},
);
properties.insert(
"path".to_string(),
"output_csv_path".to_string(),
JsonSchema::String {
description: Some("Optional output CSV path override.".to_string()),
},
+14 -8
View File
@@ -18,7 +18,9 @@ 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.
- `instruction`: instruction template to apply to each row. Use `{column_name}` placeholders to
inject values from the CSV row (column names are case-sensitive and may include spaces).
Use `{{` and `}}` for literal braces.
Optional args:
- `id_column`: header column name to use as a stable item id.
@@ -30,11 +32,13 @@ Optional args:
### `run_agent_job`
Start or resume an existing job by id.
Resume an existing job by id. Jobs auto-start when created. When resuming, any items
left in `running` state are reset to `pending` unless they already reported a result.
### `get_agent_job_status`
Return job status and progress counters.
Return job status and progress counters. Most flows should prefer `wait_agent_job`
to deterministically block until completion.
### `wait_agent_job`
@@ -42,7 +46,8 @@ 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.
Export the current job results to CSV using the stored headers and results. You can
optionally override the destination path.
### `report_agent_job_result`
@@ -52,10 +57,11 @@ Workers must call this exactly once to report a JSON object for their assigned i
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. If `auto_export` is true, the runner writes a CSV snapshot on successful completion.
6. CSV export can also be generated manually by a single writer from the SQLite store.
3. The job instruction is rendered per row by substituting `{column_name}` placeholders.
4. Each worker processes one item and calls `report_agent_job_result`.
5. The runner marks items completed after the worker finishes.
6. If `auto_export` is true, the runner writes a CSV snapshot on successful completion.
7. CSV export can also be generated manually by a single writer from the SQLite store.
## CSV Output
+163
View File
@@ -1008,6 +1008,35 @@ WHERE job_id = ? AND item_id = ? AND status = ?
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,
@@ -1136,6 +1165,54 @@ WHERE
Ok(result.rows_affected() > 0)
}
pub async fn reset_agent_job_running_items(&self, job_id: &str) -> anyhow::Result<()> {
let now = Utc::now().timestamp();
sqlx::query(
r#"
UPDATE agent_job_items
SET
status = ?,
completed_at = ?,
updated_at = ?,
assigned_thread_id = NULL,
last_error = NULL
WHERE
job_id = ?
AND status = ?
AND result_json IS NOT NULL
"#,
)
.bind(AgentJobItemStatus::Completed.as_str())
.bind(now)
.bind(now)
.bind(job_id)
.bind(AgentJobItemStatus::Running.as_str())
.execute(self.pool.as_ref())
.await?;
sqlx::query(
r#"
UPDATE agent_job_items
SET
status = ?,
updated_at = ?,
assigned_thread_id = NULL,
last_error = ?
WHERE
job_id = ?
AND status = ?
"#,
)
.bind(AgentJobItemStatus::Pending.as_str())
.bind(now)
.bind("reset to pending after resume")
.bind(job_id)
.bind(AgentJobItemStatus::Running.as_str())
.execute(self.pool.as_ref())
.await?;
Ok(())
}
pub async fn get_agent_job_progress(&self, job_id: &str) -> anyhow::Result<AgentJobProgress> {
let row = sqlx::query(
r#"
@@ -1422,12 +1499,16 @@ mod tests {
use super::StateRuntime;
use super::ThreadMetadata;
use super::state_db_filename;
use crate::AgentJobCreateParams;
use crate::AgentJobItemCreateParams;
use crate::AgentJobItemStatus;
use chrono::DateTime;
use chrono::Utc;
use codex_protocol::ThreadId;
use codex_protocol::protocol::AskForApproval;
use codex_protocol::protocol::SandboxPolicy;
use pretty_assertions::assert_eq;
use serde_json::json;
use sqlx::Row;
use std::path::Path;
use std::path::PathBuf;
@@ -1520,6 +1601,88 @@ mod tests {
let _ = tokio::fs::remove_dir_all(codex_home).await;
}
#[tokio::test]
async fn reset_running_items_for_resume_updates_statuses() {
let codex_home = unique_temp_dir();
tokio::fs::create_dir_all(&codex_home)
.await
.expect("create codex_home");
let runtime = StateRuntime::init(codex_home.clone(), "test-provider".to_string(), None)
.await
.expect("initialize runtime");
let job_id = Uuid::new_v4().to_string();
let job = runtime
.create_agent_job(
&AgentJobCreateParams {
id: job_id.clone(),
name: "resume-test".to_string(),
instruction: "do work".to_string(),
auto_export: true,
output_schema_json: None,
input_headers: vec!["path".to_string()],
input_csv_path: "/tmp/input.csv".to_string(),
output_csv_path: "/tmp/output.csv".to_string(),
},
&[
AgentJobItemCreateParams {
item_id: "row-1".to_string(),
row_index: 0,
source_id: None,
row_json: json!({ "path": "a.txt" }),
},
AgentJobItemCreateParams {
item_id: "row-2".to_string(),
row_index: 1,
source_id: None,
row_json: json!({ "path": "b.txt" }),
},
],
)
.await
.expect("create job");
runtime
.mark_agent_job_item_running(job.id.as_str(), "row-1")
.await
.expect("mark running");
runtime
.report_agent_job_item_result(job.id.as_str(), "row-1", "thread-1", &json!({"ok":true}))
.await
.expect("report result");
runtime
.mark_agent_job_item_running(job.id.as_str(), "row-2")
.await
.expect("mark running");
runtime
.set_agent_job_item_thread(job.id.as_str(), "row-2", "thread-2")
.await
.expect("set thread");
runtime
.reset_agent_job_running_items(job.id.as_str())
.await
.expect("reset running items");
let completed = runtime
.list_agent_job_items(job.id.as_str(), Some(AgentJobItemStatus::Completed), None)
.await
.expect("list completed");
let pending = runtime
.list_agent_job_items(job.id.as_str(), Some(AgentJobItemStatus::Pending), None)
.await
.expect("list pending");
assert_eq!(completed.len(), 1);
assert_eq!(completed[0].item_id, "row-1");
assert_eq!(completed[0].assigned_thread_id, None);
assert_eq!(pending.len(), 1);
assert_eq!(pending[0].item_id, "row-2");
assert_eq!(pending[0].assigned_thread_id, None);
}
#[tokio::test]
async fn backfill_state_persists_progress_and_completion() {
let codex_home = unique_temp_dir();