use anyhow::Result;
use std::sync::Arc;
use tokio::time::{Duration, sleep};
use super::run::RunArgs;
use crate::paths;
use crate::store::Store;
use crate::types::{Task, TaskId, TaskStatus};
pub(crate) async fn prepare_retry(
store: Arc<Store>,
task_id: &TaskId,
args: &RunArgs,
) -> Result<Option<RunArgs>> {
if args.retry == 0 {
return Ok(None);
}
let Some(task) = store.get_task(task_id.as_str())? else { return Ok(None) };
if task.status != TaskStatus::Failed {
return Ok(None);
}
let stderr_tail = read_stderr_tail(task_id.as_str(), 5);
if let Some(parent_id) = args.parent_task_id.as_deref()
&& stderr_tail == read_stderr_tail(parent_id, 5)
{
println!("Retry stopped: identical stderr to previous attempt.");
return Ok(None);
}
let depth = retry_depth(&store, args.parent_task_id.as_deref())?;
let attempt = depth + 1;
let backoff_secs = backoff_for_attempt(attempt);
println!("Retry {attempt}/{}: re-dispatching after {backoff_secs}s...", depth + args.retry);
sleep(Duration::from_secs(backoff_secs)).await;
let prompt = root_prompt(&store, &task).unwrap_or_else(|| args.prompt.clone());
let mut retry_args = args.clone();
retry_args.prompt =
format!("[Previous attempt failed]\nError: {stderr_tail}\n\n[Original task]\n{prompt}");
retry_args.retry = args.retry.saturating_sub(1);
retry_args.background = false;
retry_args.parent_task_id = Some(task_id.as_str().to_string());
if let Some(ref wt) = task.worktree_path {
retry_args.worktree = None;
if std::path::Path::new(wt).is_dir() {
aid_info!("[aid] Retry reusing worktree: {wt}");
retry_args.dir = Some(wt.clone());
}
}
Ok(Some(retry_args))
}
pub(crate) fn read_stderr_tail(task_id: &str, lines: usize) -> String {
let Ok(stderr) = std::fs::read_to_string(paths::stderr_path(task_id)) else {
return "stderr unavailable".to_string();
};
let tail: Vec<_> = stderr.lines().rev().take(lines).collect();
if tail.is_empty() { "stderr unavailable".to_string() } else { tail.into_iter().rev().collect::<Vec<_>>().join("\n") }
}
fn retry_depth(store: &Store, parent_task_id: Option<&str>) -> Result<u32> {
let mut depth = 0;
let mut current = parent_task_id.map(str::to_string);
while let Some(task_id) = current {
let Some(task) = store.get_task(&task_id)? else { break };
depth += 1;
current = task.parent_task_id;
}
Ok(depth)
}
fn backoff_for_attempt(attempt: u32) -> u64 {
match attempt { 0 | 1 => 5, 2 => 15, _ => 45 }
}
pub(crate) fn root_prompt(store: &Store, task: &Task) -> Option<String> {
let mut prompt = task.prompt.clone();
let mut current = task.parent_task_id.clone();
while let Some(task_id) = current {
let Some(parent) = store.get_task(&task_id).ok().flatten() else { break };
prompt = parent.prompt;
current = parent.parent_task_id;
}
Some(prompt)
}
#[cfg(test)]
mod tests {
use super::*;
use chrono::Local;
use crate::types::{AgentKind, VerifyStatus};
fn task(id: &str) -> Task {
Task {
id: TaskId(id.to_string()), agent: AgentKind::Codex, prompt: "prompt".to_string(),
custom_agent_name: None,
resolved_prompt: None, category: None, status: TaskStatus::Failed, parent_task_id: None, workgroup_id: None,
caller_kind: None, caller_session_id: None, agent_session_id: None, repo_path: None, worktree_path: None,
worktree_branch: None, start_sha: None, log_path: None, output_path: None, tokens: None, prompt_tokens: None,
duration_ms: None, model: None, cost_usd: None, exit_code: None, created_at: Local::now(),
completed_at: None, verify: None, verify_status: VerifyStatus::Skipped, pending_reason: None, read_only: false, budget: false,
audit_verdict: None, audit_report_path: None, delivery_assessment: None,
}
}
#[test]
fn backoff_for_attempt_increases() {
assert!(backoff_for_attempt(1) < backoff_for_attempt(2));
assert!(backoff_for_attempt(2) < backoff_for_attempt(3));
}
#[test]
fn retry_depth_is_zero_without_parent() {
let store = Store::open_memory().unwrap();
store.insert_task(&task("t-root")).unwrap();
assert_eq!(retry_depth(&store, None).unwrap(), 0);
}
#[test]
fn test_retry_depth_with_chain() {
let store = Store::open_memory().unwrap();
let mut t_root = task("t-root");
t_root.parent_task_id = None;
store.insert_task(&t_root).unwrap();
let mut t_r1 = task("t-r1");
t_r1.parent_task_id = Some("t-root".to_string());
store.insert_task(&t_r1).unwrap();
let mut t_r2 = task("t-r2");
t_r2.parent_task_id = Some("t-r1".to_string());
store.insert_task(&t_r2).unwrap();
assert_eq!(retry_depth(&store, Some("t-root")).unwrap(), 1);
assert_eq!(retry_depth(&store, Some("t-r1")).unwrap(), 2);
}
#[test]
fn test_root_prompt_walks_chain() {
let store = Store::open_memory().unwrap();
let mut t_root = task("root");
t_root.prompt = "original".to_string();
t_root.parent_task_id = None;
store.insert_task(&t_root).unwrap();
let mut t_r1 = task("r1");
t_r1.prompt = "retry1".to_string();
t_r1.parent_task_id = Some("root".to_string());
store.insert_task(&t_r1).unwrap();
let mut t_r2 = task("r2");
t_r2.prompt = "retry2".to_string();
t_r2.parent_task_id = Some("r1".to_string());
store.insert_task(&t_r2).unwrap();
let r2_task = store.get_task("r2").unwrap().unwrap();
assert_eq!(root_prompt(&store, &r2_task), Some("original".to_string()));
}
#[test]
fn test_backoff_capped() {
assert_eq!(backoff_for_attempt(10), backoff_for_attempt(3));
}
}