ai-dispatch 9.3.2

// Batch retry logic for failed/skipped tasks.
// Exports: retry_failed
// Deps: crate::cmd::run, crate::store::Store, crate::types::Task
use crate::cmd::run::{self, 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, sync::Arc};
use tokio::time::Duration;
/// Caps the serialized retry waiter only; the retried task keeps running.
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::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);
    for bucket in bucket_retry_tasks(retry_tasks) {
        dispatch_retry_bucket(&store, bucket, group_id, agent_override).await?;
    }
    Ok(())
}

pub(super) fn should_retry_task(status: TaskStatus, include_waiting: bool) -> bool {
    matches!(status, TaskStatus::Failed | TaskStatus::Skipped)
        || (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 (dir, worktree) = retry_target(task);
    let agent_name = if let Some(override_name) = agent_override {
        override_name.to_string()
    } else {
        let original = task.agent_display_name().to_string();
        if let Some(kind) = crate::types::AgentKind::parse_str(&original) {
            if crate::rate_limit::is_rate_limited(&kind) {
                if let Some(fallback) = crate::agent::selection::coding_fallback_for(&kind) {
                    crate::aid_info!(
                        "[aid] {} is rate-limited, retrying with fallback: {}",
                        original,
                        fallback.as_str()
                    );
                    fallback.as_str().to_string()
                } else {
                    original
                }
            } else {
                original
            }
        } else {
            original
        }
    };
    let mut run_args = RunArgs::saved_for_task(store, task.id.as_str())?.unwrap_or_else(|| {
        RunArgs {
            repo: task.repo_path.clone(),
            output: task.output_path.clone(),
            model: task.model.clone(),
            verify: task.verify.clone(),
            read_only: task.read_only,
            budget: task.budget,
            ..Default::default()
        }
    });
    run_args.agent_name = agent_name;
    run_args.prompt = task.prompt.clone();
    run_args.dir = dir;
    run_args.worktree = worktree;
    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.existing_task_id = None;
    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<()> {
    if bucket.tasks.len() == 1 {
        let run_args = retry_task_to_run_args(store, &bucket.tasks[0], group_id, agent_override)?;
        let _ = run::run(store.clone(), run_args).await?;
        return Ok(());
    }
    let label = bucket.label();
    println!("[aid] Serializing {} retries in worktree {}", bucket.tasks.len(), label);
    for task in bucket.tasks {
        let run_args = retry_task_to_run_args(store, &task, group_id, agent_override)?;
        let task_id = run::run(store.clone(), run_args).await?;
        wait_for_retry_completion(store, task_id.as_str()).await?;
    }
    println!("[aid] Worktree {} bucket complete", label);
    Ok(())
}

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)
        }
    }
}

fn retry_target(task: &Task) -> (Option<String>, Option<String>) {
    match task.worktree_path.as_ref() {
        Some(path) if std::path::Path::new(path).exists() => (Some(path.clone()), None),
        Some(_) => (None, task.worktree_branch.clone()),
        None => (task.repo_path.clone(), task.worktree_branch.clone()),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::paths;
    use crate::rate_limit::{clear_rate_limit, mark_rate_limited};
    use crate::types::{AgentKind, TaskId, VerifyStatus};
    use chrono::Local;
    fn make_task(id: &str, agent: AgentKind) -> Task {
        Task {
            id: TaskId(id.to_string()),
            agent,
            custom_agent_name: None,
            prompt: "test prompt".to_string(),
            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, final_head_sha: None, final_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,
        }
    }

    fn task_ids(bucket: &RetryBucket) -> Vec<String> { bucket.tasks.iter().map(|task| task.id.to_string()).collect() }
    fn aid_home_guard(name: &str) -> paths::AidHomeGuard {
        let temp_dir = std::env::temp_dir().join(name);
        let guard = paths::AidHomeGuard::set(&temp_dir);
        std::fs::create_dir_all(paths::aid_dir()).ok();
        guard
    }

    #[test]
    fn retry_uses_original_when_not_rate_limited() {
        let _guard = aid_home_guard("aid-retry-fallback-test-normal");
        let task = make_task("t-001", AgentKind::Codex);
        let args = retry_task_to_run_args(&Store::open_memory().unwrap(), &task, "wg-test", None).unwrap();
        assert_eq!(args.agent_name, "codex");
        clear_rate_limit(&AgentKind::Codex);
    }

    #[test]
    fn retry_uses_fallback_when_rate_limited() {
        let _guard = aid_home_guard("aid-retry-fallback-test-limited");
        // CI hosts have no agent binaries on PATH, so pin the detected set
        // to exercise the fallback logic deterministically.
        let _agents = crate::agent::DetectAgentsGuard::set(vec![
            AgentKind::Gemini,
            AgentKind::Qwen,
            AgentKind::Codex,
            AgentKind::Copilot,
        ]);
        mark_rate_limited(&AgentKind::Codex, "rate limit exceeded");
        let task = make_task("t-002", AgentKind::Codex);
        let args = retry_task_to_run_args(&Store::open_memory().unwrap(), &task, "wg-test", None).unwrap();
        assert_ne!(args.agent_name, "codex", "Should use fallback when rate-limited");
        clear_rate_limit(&AgentKind::Codex);
    }

    #[test]
    fn retry_uses_override_regardless_of_rate_limit() {
        let _guard = aid_home_guard("aid-retry-fallback-test-override");
        mark_rate_limited(&AgentKind::Codex, "rate limit exceeded");
        let task = make_task("t-003", AgentKind::Codex);
        let args = retry_task_to_run_args(&Store::open_memory().unwrap(), &task, "wg-test", Some("gemini")).unwrap();
        assert_eq!(args.agent_name, "gemini", "Override should bypass rate limit check");
        clear_rate_limit(&AgentKind::Codex);
    }

    #[test]
    fn retry_unchanged_for_unknown_agent() {
        let _guard = aid_home_guard("aid-retry-fallback-test-unknown");
        let task = make_task("t-004", AgentKind::Custom);
        let args = retry_task_to_run_args(&Store::open_memory().unwrap(), &task, "wg-test", None).unwrap();
        assert_eq!(args.agent_name, "custom");
    }

    #[test]
    fn retry_bucket_groups_shared_worktree_tasks() {
        let mut first = make_task("t-101", AgentKind::Codex);
        let mut second = make_task("t-102", AgentKind::Codex);
        let mut third = make_task("t-103", AgentKind::Codex);
        first.worktree_branch = Some("feat/gitbutler".to_string());
        second.worktree_branch = Some("feat/gitbutler".to_string());
        third.worktree_branch = Some("feat/gitbutler".to_string());
        let buckets = bucket_retry_tasks(vec![first, second, third]);
        assert_eq!(buckets.len(), 1);
        assert_eq!(task_ids(&buckets[0]), vec!["t-101".to_string(), "t-102".to_string(), "t-103".to_string()]);
    }

    #[test]
    fn retry_bucket_separates_distinct_worktrees() {
        let mut first = make_task("t-201", AgentKind::Codex);
        let mut second = make_task("t-202", AgentKind::Codex);
        first.worktree_branch = Some("feat/a".to_string());
        second.worktree_branch = Some("feat/b".to_string());
        let buckets = bucket_retry_tasks(vec![first, second]);
        assert_eq!(buckets.len(), 2);
        assert_eq!(task_ids(&buckets[0]), vec!["t-201".to_string()]);
        assert_eq!(task_ids(&buckets[1]), vec!["t-202".to_string()]);
    }

    #[test]
    fn retry_bucket_treats_none_worktree_as_unique() {
        let first = make_task("t-301", AgentKind::Codex);
        let second = make_task("t-302", AgentKind::Codex);
        let third = make_task("t-303", AgentKind::Codex);
        let buckets = bucket_retry_tasks(vec![first, second, third]);
        assert_eq!(buckets.len(), 3);
        assert_eq!(task_ids(&buckets[0]), vec!["t-301".to_string()]);
        assert_eq!(task_ids(&buckets[1]), vec!["t-302".to_string()]);
        assert_eq!(task_ids(&buckets[2]), vec!["t-303".to_string()]);
    }
}