mirror of
https://github.com/openai/codex.git
synced 2026-11-03 19:53:25 +00:00
Agent jobs: template instructions and resume safety
This commit is contained in:
@@ -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}");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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()),
|
||||
},
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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();
|
||||
|
||||
Reference in New Issue
Block a user