use crate::cmd::run::{self, apply_retry_target, switch_agent, RunArgs};
use crate::cmd::wait::{wait_for_task_ids, WaitOutcome};
use crate::store::Store;
use crate::types::{Task, TaskStatus};
use anyhow::Result;
use std::{collections::HashMap, future::Future, sync::Arc};
use tokio::time::Duration;
const SERIAL_RETRY_TIMEOUT_SECS: u64 = 30 * 60;
type WorktreeIdentity = (Option<String>, Option<String>);
#[derive(Debug)]
struct RetryBucket {
worktree_path: Option<String>,
worktree_branch: Option<String>,
tasks: Vec<Task>,
}
impl RetryBucket {
fn new(task: Task) -> Self {
Self {
worktree_path: task.worktree_path.clone(),
worktree_branch: task.worktree_branch.clone(),
tasks: vec![task],
}
}
fn label(&self) -> String {
self.worktree_branch
.clone()
.or_else(|| self.worktree_path.clone())
.unwrap_or_else(|| "unknown".to_string())
}
}
pub async fn retry_failed(
store: Arc<Store>,
group_id: &str,
agent_override: Option<&str>,
include_waiting: bool,
) -> Result<()> {
crate::support::sanitize::validate_workgroup_id(group_id)?;
let tasks = store.list_tasks_by_group(group_id)?;
let total = tasks.len();
let retry_tasks: Vec<_> = tasks
.into_iter()
.filter(|task| should_retry_task(task.status, include_waiting))
.collect();
if retry_tasks.is_empty() {
println!("No retryable tasks in {group_id}");
return Ok(());
}
println!("[batch] Retrying {}/{} task(s) in {group_id}", retry_tasks.len(), total);
let mut failed_retries = Vec::new();
for bucket in bucket_retry_tasks(retry_tasks) {
failed_retries.extend(
dispatch_retry_bucket(&store, bucket, group_id, agent_override).await?,
);
}
if !failed_retries.is_empty() {
anyhow::bail!("Retried task(s) did not succeed: {}", failed_retries.join(", "))
}
Ok(())
}
pub(super) fn should_retry_task(status: TaskStatus, include_waiting: bool) -> bool {
matches!(status, TaskStatus::Failed | TaskStatus::Skipped | TaskStatus::Stopped)
|| (include_waiting && status == TaskStatus::Waiting)
}
pub(crate) fn retry_task_to_run_args(store: &Store, task: &Task, group_id: &str, agent_override: Option<&str>) -> Result<RunArgs> {
let mut run_args = RunArgs::for_retry(store, task)?;
run_args.prompt = task.prompt.clone();
if let Some(name) = agent_override {
switch_agent(&mut run_args, name.to_string());
} else {
let (kind, custom) = crate::routing::rate_limit::resolve_agent(&run_args.agent_name);
if crate::routing::rate_limit::is_rate_limited(&kind, custom)
&& run_args.cascade.is_empty()
&& let Some(candidate) = run::advice_route::automatic_candidate(Some(store), &run_args)
{
run::advice_route::apply_candidate(&mut run_args, &candidate);
}
}
apply_retry_target(task, &mut run_args)?;
run_args.group = Some(group_id.to_string());
run_args.background = true;
run_args.announce = true;
run_args.parent_task_id = Some(task.id.to_string());
run_args.origin = crate::types::DispatchOrigin::ExplicitRetry;
Ok(run_args)
}
fn bucket_retry_tasks(tasks: Vec<Task>) -> Vec<RetryBucket> {
let mut buckets = Vec::new();
let mut bucket_indexes = HashMap::<WorktreeIdentity, usize>::new();
for task in tasks {
let Some(identity) = retry_bucket_identity(&task) else {
buckets.push(RetryBucket::new(task));
continue;
};
if let Some(bucket_idx) = bucket_indexes.get(&identity).copied() {
buckets[bucket_idx].tasks.push(task);
continue;
}
let bucket_idx = buckets.len();
bucket_indexes.insert(identity, bucket_idx);
buckets.push(RetryBucket::new(task));
}
buckets
}
fn retry_bucket_identity(task: &Task) -> Option<WorktreeIdentity> {
if task.worktree_path.is_none() && task.worktree_branch.is_none() {
None
} else {
Some((task.worktree_path.clone(), task.worktree_branch.clone()))
}
}
async fn dispatch_retry_bucket(
store: &Arc<Store>,
bucket: RetryBucket,
group_id: &str,
agent_override: Option<&str>,
) -> Result<Vec<String>> {
let wait_for_completion = bucket.tasks.len() > 1;
let label = bucket.label();
if wait_for_completion {
println!("[aid] Serializing {} retries in worktree {}", bucket.tasks.len(), label);
}
let store = store.clone();
let group_id = group_id.to_string();
let agent_override = agent_override.map(str::to_string);
let failures = retry_tasks_serially(bucket.tasks, move |task| {
let store = store.clone();
let group_id = group_id.clone();
let agent_override = agent_override.clone();
async move {
let run_args = retry_task_to_run_args(
store.as_ref(),
&task,
&group_id,
agent_override.as_deref(),
)?;
let task_id = run::run(store.clone(), run_args).await?;
if wait_for_completion {
wait_for_retry_completion(&store, task_id.as_str()).await?;
}
Ok(())
}
})
.await;
if wait_for_completion {
println!("[aid] Worktree {} bucket complete", label);
}
Ok(failures)
}
pub(super) async fn retry_tasks_serially<F, Fut>(tasks: Vec<Task>, mut retry_one: F) -> Vec<String>
where
F: FnMut(Task) -> Fut,
Fut: Future<Output = Result<()>>,
{
let mut failures = Vec::new();
for task in tasks {
let task_id = task.id.to_string();
if let Err(error) = retry_one(task).await {
aid_error!("[aid] Retry failed for {task_id}: {error}");
failures.push(task_id);
}
}
failures
}
async fn wait_for_retry_completion(store: &Arc<Store>, task_id: &str) -> Result<()> {
let timeout = Duration::from_secs(SERIAL_RETRY_TIMEOUT_SECS);
let task_ids = [task_id.to_string()];
match wait_for_task_ids(store, &task_ids, None, false, Some(timeout)).await? {
WaitOutcome::Completed => Ok(()),
WaitOutcome::TimedOut(_) => {
anyhow::bail!(
"Timed out waiting for retried task {} after {}s",
task_id,
SERIAL_RETRY_TIMEOUT_SECS
)
}
WaitOutcome::Failed(task_ids) => {
anyhow::bail!("Retried task(s) did not succeed: {}", task_ids.join(", "))
}
}
}
#[cfg(test)]
#[path = "retry_tests.rs"]
mod tests;