diff --git a/crates/agentflare-jobs/src/executor.rs b/crates/agentflare-jobs/src/executor.rs index c27b6016..d02757be 100644 --- a/crates/agentflare-jobs/src/executor.rs +++ b/crates/agentflare-jobs/src/executor.rs @@ -17,5 +17,35 @@ pub trait InProcessExecutor: Send + Sync { job_id: &str, args: &[String], log: &mut dyn std::io::Write, - ) -> Result<(), String>; + ) -> Result<(), JobFailure>; +} + +/// What went wrong, plus an optional hint for how long to wait before this +/// job (if it still has retries left) becomes eligible again. Set by a +/// caller that classifies the failure as transient/rate-limit shaped (see +/// the main binary's `auth_runner`/`auth_db` — this crate stays agnostic to +/// *why* a delay is warranted, it just carries the number through to +/// `Queue::fail`). Left `None` for anything else, which keeps `Queue::fail`'s +/// existing instant-requeue behavior. +pub struct JobFailure { + pub message: String, + pub retry_after_secs: Option, +} + +impl From for JobFailure { + fn from(message: String) -> Self { + JobFailure { + message, + retry_after_secs: None, + } + } +} + +impl From<&str> for JobFailure { + fn from(message: &str) -> Self { + JobFailure { + message: message.to_string(), + retry_after_secs: None, + } + } } diff --git a/crates/agentflare-jobs/src/lib.rs b/crates/agentflare-jobs/src/lib.rs index 6c1083ad..6b66d8e5 100644 --- a/crates/agentflare-jobs/src/lib.rs +++ b/crates/agentflare-jobs/src/lib.rs @@ -4,7 +4,7 @@ pub mod supervisor; pub mod types; pub mod worker; -pub use executor::InProcessExecutor; +pub use executor::{InProcessExecutor, JobFailure}; pub use queue::Queue; pub use supervisor::Supervisor; pub use types::{AgentJob, JobInfo, JobOutput, JobState}; diff --git a/crates/agentflare-jobs/src/queue.rs b/crates/agentflare-jobs/src/queue.rs index 289a6fa4..58f0b559 100644 --- a/crates/agentflare-jobs/src/queue.rs +++ b/crates/agentflare-jobs/src/queue.rs @@ -59,6 +59,7 @@ pub fn migrations() -> rusqlite_migration::Migrations<'static> { "ALTER TABLE agent_jobs ADD COLUMN stdout_bytes INTEGER NOT NULL DEFAULT 0; ALTER TABLE agent_jobs ADD COLUMN stderr_bytes INTEGER NOT NULL DEFAULT 0;", ), + rusqlite_migration::M::up("ALTER TABLE agent_jobs ADD COLUMN not_before INTEGER;"), ]) } @@ -145,21 +146,21 @@ impl Queue { pub fn dequeue(&self) -> Result, Error> { let mut conn = self.conn.lock(); let tx = conn.transaction_with_behavior(rusqlite::TransactionBehavior::Immediate)?; + let now = db_kit::ids::now(); let row: Option<(String, String)> = { let mut stmt = tx.prepare( "SELECT id, payload FROM agent_jobs - WHERE state = 'queued' + WHERE state = 'queued' AND (not_before IS NULL OR not_before <= ?1) ORDER BY created_at ASC LIMIT 1", )?; - stmt.query_row([], |r| Ok((r.get(0)?, r.get(1)?))) + stmt.query_row(params![now], |r| Ok((r.get(0)?, r.get(1)?))) .optional()? }; let (id, payload_json) = match row { Some(r) => r, None => return Ok(None), }; - let now = db_kit::ids::now(); tx.execute( "UPDATE agent_jobs SET state = 'running', started_at = ?1 WHERE id = ?2", params![now, id], @@ -199,7 +200,7 @@ impl Queue { Ok(()) } - pub fn fail(&self, id: &str, error: &str) -> Result<(), Error> { + pub fn fail(&self, id: &str, error: &str, retry_after_secs: Option) -> Result<(), Error> { let now = db_kit::ids::now(); let conn = self.conn.lock(); let (retries, max_retries): (u32, u32) = conn.query_row( @@ -209,11 +210,12 @@ impl Queue { )?; let retried = retries < max_retries; if retried { + let not_before = retry_after_secs.map(|s| now + s as i64); conn.execute( "UPDATE agent_jobs - SET state = 'queued', retries = retries + 1, error = ?1, started_at = NULL - WHERE id = ?2", - params![error, id], + SET state = 'queued', retries = retries + 1, error = ?1, started_at = NULL, not_before = ?2 + WHERE id = ?3", + params![error, not_before, id], )?; } else { conn.execute( @@ -225,7 +227,11 @@ impl Queue { } drop(conn); // A retry goes back to 'queued' — wake workers so it's picked up - // promptly instead of waiting out the fallback poll interval. + // promptly instead of waiting out the fallback poll interval. Waking + // them even when `not_before` is in the future is harmless (their + // next `dequeue` simply finds nothing and re-parks) and keeps this + // function simple; the fallback poll in `worker_loop` is the real + // backstop for the delayed case regardless. if retried { self.wake_workers(); } @@ -391,3 +397,70 @@ impl OptionalExt for Result { } } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::types::AgentJob; + + fn test_queue() -> (Queue, tempfile::TempDir) { + let dir = tempfile::tempdir().unwrap(); + let queue = Queue::open_memory(dir.path().join("logs")).unwrap(); + (queue, dir) + } + + #[test] + fn fail_without_retry_delay_requeues_immediately() { + let (queue, _dir) = test_queue(); + let job = AgentJob::new("true").max_retries(1); + let info = queue.enqueue(&job).unwrap(); + queue.dequeue().unwrap(); + queue.fail(&info.id, "boom", None).unwrap(); + let (id, _) = queue + .dequeue() + .unwrap() + .expect("retried job should be immediately eligible"); + assert_eq!(id, info.id); + } + + #[test] + fn fail_with_retry_delay_hides_the_job_until_not_before_passes() { + let (queue, _dir) = test_queue(); + let job = AgentJob::new("true").max_retries(1); + let info = queue.enqueue(&job).unwrap(); + queue.dequeue().unwrap(); + queue.fail(&info.id, "rate limited", Some(3600)).unwrap(); + + assert!( + queue.dequeue().unwrap().is_none(), + "a job with a future not_before must not be dequeued yet" + ); + + // Simulate the delay having elapsed, without a real sleep. + { + let conn = queue.conn.lock(); + conn.execute( + "UPDATE agent_jobs SET not_before = 0 WHERE id = ?1", + params![info.id], + ) + .unwrap(); + } + let (id, _) = queue + .dequeue() + .unwrap() + .expect("job should be eligible once not_before has passed"); + assert_eq!(id, info.id); + } + + #[test] + fn fail_past_max_retries_marks_failed_regardless_of_retry_delay() { + let (queue, _dir) = test_queue(); + let job = AgentJob::new("true").max_retries(0); + let info = queue.enqueue(&job).unwrap(); + queue.dequeue().unwrap(); + queue.fail(&info.id, "boom", Some(60)).unwrap(); + let jobs = queue.list(Some(JobState::Failed)).unwrap(); + assert_eq!(jobs.len(), 1); + assert_eq!(jobs[0].id, info.id); + } +} diff --git a/crates/agentflare-jobs/src/worker.rs b/crates/agentflare-jobs/src/worker.rs index 720ff583..8f087341 100644 --- a/crates/agentflare-jobs/src/worker.rs +++ b/crates/agentflare-jobs/src/worker.rs @@ -1,4 +1,4 @@ -use crate::executor::InProcessExecutor; +use crate::executor::{InProcessExecutor, JobFailure}; use crate::queue::Queue; use crate::supervisor::Supervisor; use crate::types::JobOutput; @@ -99,7 +99,7 @@ fn worker_loop(queue: &Queue, running: &AtomicBool, executor: Option<&Arc { - if let Err(qe) = queue.fail(&id, &e.to_string()) { + if let Err(qe) = queue.fail(&id, &e.to_string(), None) { eprintln!("agentflare-jobs: failed to record failure for {id}: {qe}"); } } @@ -147,6 +147,7 @@ fn run_in_process( if let Err(e) = queue.fail( id, "job is marked in_process but no InProcessExecutor is registered on this WorkerPool", + None, ) { eprintln!("agentflare-jobs: failed to record failure for {id}: {e}"); } @@ -158,14 +159,14 @@ fn run_in_process( let mut log_file = match std::fs::File::create(&stdout_path) { Ok(f) => f, Err(e) => { - if let Err(qe) = queue.fail(id, &format!("failed to open job log file: {e}")) { + if let Err(qe) = queue.fail(id, &format!("failed to open job log file: {e}"), None) { eprintln!("agentflare-jobs: failed to record failure for {id}: {qe}"); } return; } }; - let (tx, rx) = std::sync::mpsc::channel::>(); + let (tx, rx) = std::sync::mpsc::channel::>(); let executor = executor.clone(); let job_id = id.to_string(); let args = job.args.clone(); @@ -192,8 +193,8 @@ fn run_in_process( eprintln!("agentflare-jobs: failed to complete job {id}: {e}"); } } - Ok(Err(msg)) => { - if let Err(e) = queue.fail(id, &msg) { + Ok(Err(failure)) => { + if let Err(e) = queue.fail(id, &failure.message, failure.retry_after_secs) { eprintln!("agentflare-jobs: failed to record failure for {id}: {e}"); } } @@ -205,7 +206,7 @@ fn run_in_process( what this timeout measures)", job.timeout_secs ); - if let Err(e) = queue.fail(id, &msg) { + if let Err(e) = queue.fail(id, &msg, None) { eprintln!("agentflare-jobs: failed to record failure for {id}: {e}"); } } diff --git a/crates/agentflare-jobs/tests/in_process_test.rs b/crates/agentflare-jobs/tests/in_process_test.rs index 72a36e53..d0ad47a3 100644 --- a/crates/agentflare-jobs/tests/in_process_test.rs +++ b/crates/agentflare-jobs/tests/in_process_test.rs @@ -5,7 +5,9 @@ //! separate from `worker_test.rs`'s existing subprocess-path coverage, which //! stays unchanged and passing to prove that path is untouched. -use agentflare_jobs::{AgentJob, InProcessExecutor, JobInfo, JobState, Queue, WorkerPool}; +use agentflare_jobs::{ + AgentJob, InProcessExecutor, JobFailure, JobInfo, JobState, Queue, WorkerPool, +}; use std::sync::Arc; fn test_queue() -> Queue { @@ -34,7 +36,7 @@ impl InProcessExecutor for EchoExecutor { job_id: &str, args: &[String], log: &mut dyn std::io::Write, - ) -> Result<(), String> { + ) -> Result<(), JobFailure> { let _ = writeln!(log, "running job {job_id} with args {args:?}"); Ok(()) } @@ -81,8 +83,8 @@ impl InProcessExecutor for FailingExecutor { _job_id: &str, _args: &[String], _log: &mut dyn std::io::Write, - ) -> Result<(), String> { - Err("deliberate failure".to_string()) + ) -> Result<(), JobFailure> { + Err("deliberate failure".into()) } } @@ -135,7 +137,7 @@ impl InProcessExecutor for SlowExecutor { _job_id: &str, _args: &[String], _log: &mut dyn std::io::Write, - ) -> Result<(), String> { + ) -> Result<(), JobFailure> { std::thread::sleep(std::time::Duration::from_secs(5)); Ok(()) } diff --git a/crates/agentflare-jobs/tests/queue_test.rs b/crates/agentflare-jobs/tests/queue_test.rs index 0c9e78f4..8a03c0f1 100644 --- a/crates/agentflare-jobs/tests/queue_test.rs +++ b/crates/agentflare-jobs/tests/queue_test.rs @@ -122,19 +122,19 @@ fn fail_retries_then_permanent() { let id = info.id; // First failure → retry (queued again) - q.fail(&id, "err1").unwrap(); + q.fail(&id, "err1", None).unwrap(); let info = q.get(&id).unwrap(); assert_eq!(info.state, JobState::Queued); assert_eq!(info.retries, 1); // Second failure → retry - q.fail(&id, "err2").unwrap(); + q.fail(&id, "err2", None).unwrap(); let info = q.get(&id).unwrap(); assert_eq!(info.state, JobState::Queued); assert_eq!(info.retries, 2); // Third failure → permanent - q.fail(&id, "err3").unwrap(); + q.fail(&id, "err3", None).unwrap(); let info = q.get(&id).unwrap(); assert_eq!(info.state, JobState::Failed); assert_eq!(info.error.as_deref(), Some("err3")); @@ -177,7 +177,7 @@ fn cleanup_removes_old_jobs() { let q = test_queue(); q.enqueue(&AgentJob::new("a")).unwrap(); let (id, _) = q.dequeue().unwrap().unwrap(); - q.fail(&id, "done").unwrap(); + q.fail(&id, "done", None).unwrap(); // Use a large negative cutoff to simulate "now = 0" // The jobs have created_at = now (positive), so anything older than diff --git a/src/agent_launch.rs b/src/agent_launch.rs index 5d3904d8..9e25306a 100644 --- a/src/agent_launch.rs +++ b/src/agent_launch.rs @@ -375,7 +375,11 @@ pub fn run_headless( diagnostic_suffix(&c) )) } - Ok(_) => HeadlessOutcome::Failed(format!("{} exited non-zero", spec.display_name)), + Ok(c) => HeadlessOutcome::Failed(format!( + "{} exited non-zero{}", + spec.display_name, + diagnostic_suffix(&c) + )), Err(e) => HeadlessOutcome::Failed(format!("failed to run {}: {e}", spec.display_name)), } } @@ -862,4 +866,25 @@ mod tests { assert!(suffix.contains("last stderr before kill")); assert!(suffix.contains("panic: something broke")); } + + #[test] + fn failed_non_zero_exit_includes_captured_output() { + let captured = Captured { + stdout: String::new(), + stderr: "HTTP 429 Too Many Requests".to_string(), + success: false, + timed_out: false, + idle_killed: false, + }; + let msg = match Ok::<_, std::io::Error>(captured) { + Ok(c) if c.success => unreachable!(), + Ok(c) if c.timed_out => unreachable!(), + Ok(c) => format!("test exited non-zero{}", diagnostic_suffix(&c)), + Err(_) => unreachable!(), + }; + assert!( + msg.contains("HTTP 429 Too Many Requests"), + "expected captured stderr in the failure message, got: {msg}" + ); + } } diff --git a/src/auth_db.rs b/src/auth_db.rs index 951e7193..092b25f1 100644 --- a/src/auth_db.rs +++ b/src/auth_db.rs @@ -260,6 +260,19 @@ pub fn list_cooldowns(conn: &Connection, agent: Option<&str>) -> Vec bool { + !list_cooldowns(conn, Some(agent)).is_empty() +} + pub fn clear_cooldown(conn: &Connection, agent: &str, profile: &str) { conn.execute( "DELETE FROM cooldowns WHERE agent = ?1 AND profile = ?2", @@ -447,6 +460,16 @@ mod tests { }); } + #[test] + fn is_cooling_down_reflects_active_cooldowns() { + with_temp_home(|| { + let conn = open_or_rebuild(); + assert!(!is_cooling_down(&conn, "claude-code")); + set_cooldown(&conn, "claude-code", "alice", 30, "rate limit"); + assert!(is_cooling_down(&conn, "claude-code")); + }); + } + #[test] fn clear_cooldown_removes() { with_temp_home(|| { diff --git a/src/auth_runner.rs b/src/auth_runner.rs index 1500b1e1..9b6f7c99 100644 --- a/src/auth_runner.rs +++ b/src/auth_runner.rs @@ -65,12 +65,22 @@ enum ExitKind { Failure(i32), } +/// Pattern-matches `text` against the same rate-limit phrase list the +/// interactive retry loop above uses (`RATE_LIMIT_PATTERNS`), independent of +/// any exit code — used by the autonomous work-dispatch path +/// (`src/cli/work.rs`'s `classify_and_cooldown`) to classify a headless +/// run's captured failure text the same way this file's own retry loop +/// already classifies a subprocess's stderr. +pub(crate) fn is_rate_limited(text: &str) -> bool { + let lower = text.to_lowercase(); + RATE_LIMIT_PATTERNS.iter().any(|p| lower.contains(p)) +} + fn categorize_exit(code: i32, stderr: &str) -> ExitKind { if code == 0 { return ExitKind::Success; } - let lower = stderr.to_lowercase(); - if RATE_LIMIT_PATTERNS.iter().any(|p| lower.contains(p)) { + if is_rate_limited(stderr) { return ExitKind::RateLimited; } ExitKind::Failure(code) @@ -187,4 +197,11 @@ mod tests { let result = categorize_exit(1, "something went wrong"); assert!(matches!(result, ExitKind::Failure(1))); } + + #[test] + fn is_rate_limited_matches_common_phrases() { + assert!(is_rate_limited("HTTP 429 Too Many Requests")); + assert!(is_rate_limited("quota exceeded for today")); + assert!(!is_rate_limited("something went wrong")); + } } diff --git a/src/cli/work.rs b/src/cli/work.rs index 11e1065f..28e5bf03 100644 --- a/src/cli/work.rs +++ b/src/cli/work.rs @@ -339,10 +339,63 @@ fn notify(recipient: &str, body: &str, item_id: &str) { impl WorkArgs { pub fn run(self) { - std::process::exit(execute_work(self, &mut std::io::stdout())); + std::process::exit(execute_work(self, &mut std::io::stdout()).exit_code); } } +/// `execute_work`'s result: the process exit code (0 = success), plus — set +/// only when the failure was classified as rate-limit shaped — a hint for +/// how long the job queue should wait before retrying this item. +/// `WorkItemExecutor` converts this into `agentflare_jobs::JobFailure`. +pub(crate) struct WorkOutcome { + pub exit_code: i32, + pub retry_after_secs: Option, +} + +impl From for WorkOutcome { + fn from(exit_code: i32) -> Self { + WorkOutcome { + exit_code, + retry_after_secs: None, + } + } +} + +/// Cooldown-table key used when there's no active vault rotation profile for +/// `agent` — keeps `auth_db`'s `(agent, profile)`-keyed cooldown table as the +/// single source of truth for both the interactive (`auth_runner`) and +/// autonomous (this file) dispatch paths, even for the common single- +/// credential setup that never configured vault profiles. +const DEFAULT_COOLDOWN_PROFILE: &str = "__default__"; +/// Matches the cooldown length `auth_runner::run` already uses for the +/// interactive path's rate-limit rotation. +const RATE_LIMIT_COOLDOWN_MINUTES: u32 = 30; + +/// Classifies a headless run's failure message the same way the interactive +/// `agentflare run` path does (`auth_runner::is_rate_limited`) and, if it +/// looks rate-limit shaped, records a cooldown so `auth_db::is_cooling_down` +/// (checked by the discovery tick before dispatching the next item for this +/// agent, and by `auth rotate`) sees it too. Returns the seconds until that +/// cooldown clears, for the caller to pass through as the job queue's +/// retry-after delay. +fn classify_and_cooldown(agent: &str, failure_message: &str) -> Option { + if !crate::auth_runner::is_rate_limited(failure_message) { + return None; + } + let conn = crate::auth_db::open_or_rebuild(); + let profile = crate::auth_db::get_rotation_last(&conn, agent) + .map(|(profile, _)| profile) + .unwrap_or_else(|| DEFAULT_COOLDOWN_PROFILE.to_string()); + crate::auth_db::set_cooldown( + &conn, + agent, + &profile, + RATE_LIMIT_COOLDOWN_MINUTES, + "rate limit", + ); + Some(RATE_LIMIT_COOLDOWN_MINUTES as u64 * 60) +} + /// Claims `args.target`, runs the resolved agent on it, and reports the /// outcome back onto the item — the whole body of `agentflare work`. /// Progress lines that used to go straight to stdout now go through `log` @@ -351,7 +404,7 @@ impl WorkArgs { /// progress captured into that job's own log file — the exact same file the /// dashboard already tails for subprocess-dispatched jobs — rather than only /// working when there's a real subprocess's stdout to capture. -pub(crate) fn execute_work(args: WorkArgs, log: &mut dyn std::io::Write) -> i32 { +pub(crate) fn execute_work(args: WorkArgs, log: &mut dyn std::io::Write) -> WorkOutcome { let mcp = AgentflareMcp::default(); let timeout = Duration::from_secs(args.timeout); let idle_timeout = Duration::from_secs(args.idle_timeout); @@ -365,7 +418,7 @@ pub(crate) fn execute_work(args: WorkArgs, log: &mut dyn std::io::Write) -> i32 crate::ui::error(&format!( "unknown agent: {explicit} — use `agentflare agents list`" )); - return 1; + return 1.into(); }; // The claim below identifies its own owner via `claims::owner_id()`, // which falls back to agent-detector's parent-process/env sniffing @@ -406,7 +459,7 @@ pub(crate) fn execute_work(args: WorkArgs, log: &mut dyn std::io::Write) -> i32 Ok(json) => json, Err(e) => { crate::ui::error(&format!("claim failed: {}", e.message)); - return 1; + return 1.into(); } }; let claim: serde_json::Value = @@ -416,7 +469,7 @@ pub(crate) fn execute_work(args: WorkArgs, log: &mut dyn std::io::Write) -> i32 let owner = claim["owner"].as_str().unwrap_or("?"); let age = claim["age_secs"].as_i64().unwrap_or(0); crate::ui::error(&format!("item held by {owner} ({age}s) — cannot claim")); - return 1; + return 1.into(); } let item_id = claim["item_id"] .as_str() @@ -433,7 +486,7 @@ pub(crate) fn execute_work(args: WorkArgs, log: &mut dyn std::io::Write) -> i32 let msg = "claim succeeded but no worktree was created (bad git state?)"; release_and_comment(&mcp, item_id, msg, args.notify.as_deref()); crate::ui::error(msg); - return 1; + return 1.into(); }; let _ = writeln!(log, "worktree: {}", wpath.display()); @@ -456,7 +509,7 @@ pub(crate) fn execute_work(args: WorkArgs, log: &mut dyn std::io::Write) -> i32 let msg = "failed to read item details after claim"; release_and_comment(&mcp, item_id, msg, args.notify.as_deref()); crate::ui::error(msg); - return 1; + return 1.into(); } }; @@ -489,14 +542,14 @@ pub(crate) fn execute_work(args: WorkArgs, log: &mut dyn std::io::Write) -> i32 Err(msg) => { release_and_comment(&mcp, item_id, &msg, args.notify.as_deref()); crate::ui::error(&msg); - return 1; + return 1.into(); } }; if headless_args(agent_enum).is_none() { let msg = format!("agent {} has no headless print mode", agent_enum.as_str()); release_and_comment(&mcp, item_id, &msg, args.notify.as_deref()); crate::ui::error(&msg); - return 1; + return 1.into(); } let _ = writeln!(log, "agent: {} ({route_reason})", agent_enum.as_str()); @@ -511,7 +564,7 @@ pub(crate) fn execute_work(args: WorkArgs, log: &mut dyn std::io::Write) -> i32 let msg = format!("failed to chdir into {}", wpath.display()); release_and_comment(&mcp, item_id, &msg, args.notify.as_deref()); crate::ui::error(&msg); - return 1; + return 1.into(); } let outcome = agent_launch::run_headless( @@ -554,7 +607,7 @@ pub(crate) fn execute_work(args: WorkArgs, log: &mut dyn std::io::Write) -> i32 Ok(j) => j, Err(e) => { crate::ui::error(&format!("item_done failed: {}", e.message)); - return 1; + return 1.into(); } }; let done_val: serde_json::Value = @@ -581,14 +634,18 @@ pub(crate) fn execute_work(args: WorkArgs, log: &mut dyn std::io::Write) -> i32 if let Some(url) = &pr_url { let _ = writeln!(log, "pr: {url}"); } - 0 + 0.into() } other => { let msg = failure_message(&other); release_and_comment(&mcp, item_id, &msg, args.notify.as_deref()); crate::ui::error(&msg); let _ = writeln!(log, "failed: {msg}"); - 1 + let retry_after_secs = classify_and_cooldown(agent_enum.as_str(), &msg); + WorkOutcome { + exit_code: 1, + retry_after_secs, + } } } } @@ -606,11 +663,12 @@ impl agentflare_jobs::InProcessExecutor for WorkItemExecutor { job_id: &str, args: &[String], log: &mut dyn std::io::Write, - ) -> Result<(), String> { + ) -> Result<(), agentflare_jobs::JobFailure> { let (Some(item_id), Some(agent)) = (args.first(), args.get(1)) else { return Err(format!( "malformed in-process work job: expected [item_id, agent], got {args:?}" - )); + ) + .into()); }; let work_args = WorkArgs { target: item_id.clone(), @@ -625,13 +683,17 @@ impl agentflare_jobs::InProcessExecutor for WorkItemExecutor { // discriminator, playing the role a subprocess's unique pid plays // for `claims::owner_id()` in the CLI path (see its doc comment). let owner = format!("{agent}:{job_id}"); - let exit_code = crate::claims::with_owner_override(owner, || execute_work(work_args, log)); - if exit_code == 0 { + let outcome = crate::claims::with_owner_override(owner, || execute_work(work_args, log)); + if outcome.exit_code == 0 { Ok(()) } else { - Err(format!( - "agentflare work exited with code {exit_code} — see the job log for details" - )) + Err(agentflare_jobs::JobFailure { + message: format!( + "agentflare work exited with code {} — see the job log for details", + outcome.exit_code + ), + retry_after_secs: outcome.retry_after_secs, + }) } } } @@ -902,6 +964,36 @@ use = "opencode" assert_eq!(failure_message(&outcome), "claude not found"); } + #[test] + fn failure_message_includes_diagnostic_suffix_for_plain_failures() { + let outcome = HeadlessOutcome::Failed(format!( + "claude-code exited non-zero — last stderr before kill:\n{}", + "HTTP 429 Too Many Requests" + )); + let msg = failure_message(&outcome); + assert!(msg.contains("HTTP 429 Too Many Requests")); + } + + #[test] + fn classify_and_cooldown_ignores_non_rate_limit_failures() { + crate::paths::test_support::with_temp_home(|| { + let retry = classify_and_cooldown("claude-code", "something went wrong"); + assert!(retry.is_none()); + let conn = crate::auth_db::open_or_rebuild(); + assert!(!crate::auth_db::is_cooling_down(&conn, "claude-code")); + }); + } + + #[test] + fn classify_and_cooldown_sets_a_cooldown_on_rate_limit_shaped_failures() { + crate::paths::test_support::with_temp_home(|| { + let retry = classify_and_cooldown("claude-code", "HTTP 429 Too Many Requests"); + assert_eq!(retry, Some(RATE_LIMIT_COOLDOWN_MINUTES as u64 * 60)); + let conn = crate::auth_db::open_or_rebuild(); + assert!(crate::auth_db::is_cooling_down(&conn, "claude-code")); + }); + } + #[test] fn build_extra_args_includes_bypass_and_streaming_output_for_claude() { // Plain `--output-format json` writes NOTHING to stdout/stderr until diff --git a/src/dashboard/server.rs b/src/dashboard/server.rs index 7b5656a5..76640269 100644 --- a/src/dashboard/server.rs +++ b/src/dashboard/server.rs @@ -182,7 +182,8 @@ fn spawn_supervisor_discovery( let queue = queue.clone(); let mcp = mcp.clone(); let result = tokio::task::spawn_blocking(move || { - crate::supervisor::run_discovery_tick(&mcp, &queue) + let auth_conn = crate::auth_db::open_or_rebuild(); + crate::supervisor::run_discovery_tick(&mcp, &queue, &auth_conn) }) .await; match result { @@ -582,6 +583,27 @@ fn is_local_bind(host: &str) -> bool { matches!(host, "127.0.0.1" | "localhost" | "::1") } +/// Parses `AGENTFLARE_WORK_MAX_CONCURRENCY` (default `2`, matching the +/// hardcoded value this replaces). Split into a pure parse step so the +/// override logic is testable without mutating process-global env state — +/// env vars are shared across the whole test binary, unlike this narrow +/// seam. Zero and unparseable values fall back to the default rather than +/// silently starting a `WorkerPool` with no workers, which would wedge the +/// queue forever with no error. +fn parse_work_max_concurrency(raw: Option<&str>) -> usize { + raw.and_then(|s| s.parse::().ok()) + .filter(|n| *n > 0) + .unwrap_or(2) +} + +fn work_max_concurrency() -> usize { + parse_work_max_concurrency( + std::env::var("AGENTFLARE_WORK_MAX_CONCURRENCY") + .ok() + .as_deref(), + ) +} + pub async fn run(host: &str, port: u16, open: bool, yes_expose: bool) { if !is_local_bind(host) && !yes_expose { eprintln!( @@ -607,7 +629,7 @@ pub async fn run(host: &str, port: u16, open: bool, yes_expose: bool) { // which relies on SIGTERM/SIGKILL), so neither does this. let mut worker_pool = agentflare_jobs::WorkerPool::new(queue.clone()) .with_executor(std::sync::Arc::new(crate::cli::work::WorkItemExecutor)); - worker_pool.start(2); + worker_pool.start(work_max_concurrency()); spawn_job_cleanup(queue.clone()); spawn_supervisor_discovery( queue.clone(), @@ -636,6 +658,22 @@ pub async fn run(host: &str, port: u16, open: bool, yes_expose: bool) { mod tests { use super::*; + #[test] + fn parse_work_max_concurrency_defaults_to_two() { + assert_eq!(parse_work_max_concurrency(None), 2); + } + + #[test] + fn parse_work_max_concurrency_honors_a_valid_override() { + assert_eq!(parse_work_max_concurrency(Some("5")), 5); + } + + #[test] + fn parse_work_max_concurrency_falls_back_on_garbage_or_zero() { + assert_eq!(parse_work_max_concurrency(Some("not-a-number")), 2); + assert_eq!(parse_work_max_concurrency(Some("0")), 2); + } + fn test_queue() -> Queue { // `.keep()` so the dir outlives this function — otherwise the // returned `Queue`'s `log_dir` would point at an already-deleted diff --git a/src/supervisor.rs b/src/supervisor.rs index e9a2e859..8d478887 100644 --- a/src/supervisor.rs +++ b/src/supervisor.rs @@ -58,6 +58,7 @@ pub(crate) struct DiscoveryTickResult { pub(crate) fn run_discovery_tick( mcp: &AgentflareMcp, queue: &agentflare_jobs::Queue, + auth_conn: &rusqlite::Connection, ) -> DiscoveryTickResult { let mut result = DiscoveryTickResult { dispatched: 0, @@ -101,6 +102,13 @@ pub(crate) fn run_discovery_tick( result.skipped += 1; continue; }; + if crate::auth_db::is_cooling_down(auth_conn, agent.as_str()) { + // Leave the ready-for-work label in place, same as the + // Wait branch below: the cooldown may clear before the + // next tick, and the item must still be visible to that + // tick's discovery query. + continue; + } if dispatch_item(mcp, queue, &item, agent, &label_id_by_name, &ready_id) { result.dispatched += 1; } @@ -275,6 +283,12 @@ mod tests { agentflare_jobs::Queue::open_memory(dir.join("logs")).unwrap() } + fn test_auth_conn() -> rusqlite::Connection { + let conn = rusqlite::Connection::open_in_memory().unwrap(); + crate::auth_db::migrate(&conn).unwrap(); + conn + } + fn seed_ready_item(mcp: &AgentflareMcp, assignee: Option<&str>) -> String { mcp.with_backend_db(|conn| { let project = mcp.resolve_project(conn).unwrap(); @@ -344,7 +358,8 @@ mod tests { let queue = test_queue(); let item_id = seed_ready_item(&mcp, Some("claude-code")); - let result = run_discovery_tick(&mcp, &queue); + let auth_conn = test_auth_conn(); + let result = run_discovery_tick(&mcp, &queue, &auth_conn); assert_eq!(result.dispatched, 1); assert_eq!(result.skipped, 0); @@ -365,6 +380,31 @@ mod tests { assert!(labels_contain_name(&mcp, &labels, "dispatched")); } + #[test] + fn agent_in_cooldown_is_skipped_not_dispatched() { + let mcp = test_mcp(); + let queue = test_queue(); + let auth_conn = test_auth_conn(); + let item_id = seed_ready_item(&mcp, Some("claude-code")); + crate::auth_db::set_cooldown(&auth_conn, "claude-code", "__default__", 30, "rate limit"); + + let result = run_discovery_tick(&mcp, &queue, &auth_conn); + + assert_eq!(result.dispatched, 0); + assert!( + queue.list(None).unwrap().is_empty(), + "a cooling-down agent must not be dispatched" + ); + + let labels = mcp + .with_backend_db(|conn| agentflare_backend::item::list_labels(conn, &item_id).unwrap()) + .unwrap(); + assert!( + labels_contain_name(&mcp, &labels, "ready-for-work"), + "the item must stay ready-for-work so a later tick can pick it up once the cooldown clears" + ); + } + fn seed_ready_item_under_gated_goal(mcp: &AgentflareMcp) -> String { mcp.with_backend_db(|conn| { let project = mcp.resolve_project(conn).unwrap(); @@ -458,7 +498,8 @@ mod tests { let queue = test_queue(); let item_id = seed_ready_item_under_gated_goal(&mcp); - let result = run_discovery_tick(&mcp, &queue); + let auth_conn = test_auth_conn(); + let result = run_discovery_tick(&mcp, &queue, &auth_conn); assert_eq!(result.dispatched, 0); assert!( @@ -584,7 +625,8 @@ mod tests { let queue = test_queue(); let (_item_id, _goal_id) = seed_ready_item_under_active_goal_with_repairs(&mcp, 0); - let result = run_discovery_tick(&mcp, &queue); + let auth_conn = test_auth_conn(); + let result = run_discovery_tick(&mcp, &queue, &auth_conn); assert_eq!(result.dispatched, 1, "self-repair still dispatches the job"); assert_eq!(queue.list(None).unwrap().len(), 1); @@ -599,7 +641,8 @@ mod tests { crate::quota::decide::SELF_REPAIR_CAP, ); - let result = run_discovery_tick(&mcp, &queue); + let auth_conn = test_auth_conn(); + let result = run_discovery_tick(&mcp, &queue, &auth_conn); assert_eq!( result.dispatched, 0, @@ -620,7 +663,8 @@ mod tests { // at all) — this is the plan's explicit no-regression guarantee. let item_id = seed_ready_item(&mcp, Some("claude-code")); - let result = run_discovery_tick(&mcp, &queue); + let auth_conn = test_auth_conn(); + let result = run_discovery_tick(&mcp, &queue, &auth_conn); assert_eq!(result.dispatched, 1); assert_eq!(result.skipped, 0); @@ -634,7 +678,8 @@ mod tests { let queue = test_queue(); let item_id = seed_ready_item(&mcp, Some("opencode")); - let result = run_discovery_tick(&mcp, &queue); + let auth_conn = test_auth_conn(); + let result = run_discovery_tick(&mcp, &queue, &auth_conn); assert_eq!(result.dispatched, 0); assert_eq!(result.skipped, 1);