ai-dispatch 11.2.0

// Handler for `aid wait` — block until tasks finish.
// Prints status transitions and per-task completion summaries.
// Deps: crate::store::Store, crate::types, crate::cost

use anyhow::Result;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::time::{sleep, Duration, Instant};

use crate::cost;
use crate::store::Store;
use crate::types::{verify_required, Task, TaskFilter, TaskOutcome, TaskStatus, VerifyStatus};
use crate::task::verify::VERIFY_TIMEOUT;
#[path = "wait_settlement.rs"]
mod settlement;
use settlement::worker_is_settling;

#[derive(Debug, PartialEq, Eq)]
pub(crate) enum WaitOutcome {
    Completed,
    Failed(Vec<String>),
    TimedOut(Vec<String>),
}

pub async fn run(
    store: &Arc<Store>,
    task_ids: &[String],
    group: Option<&str>,
    exit_on_await: bool,
    timeout_secs: Option<u64>,
) -> Result<()> {
    let tracked_group = if task_ids.is_empty() { group } else { None };
    let task_ids = if task_ids.is_empty() {
        current_running_ids(store, tracked_group)?
    } else {
        task_ids.to_vec()
    };

    if task_ids.is_empty() {
        println!("No running tasks.");
        return Ok(());
    }

    println!("Waiting for {} task(s): {}", task_ids.len(), task_ids.join(", "));
    match wait_for_task_ids(
        store,
        &task_ids,
        tracked_group,
        exit_on_await,
        timeout_secs.map(Duration::from_secs),
    )
    .await?
    {
        WaitOutcome::Completed => {
            if let Some(group_id) = tracked_group
                && let Ok(repo_path) = crate::support::repo_root::resolve_git_root_string(".")
                && let Some(hint) =
                    crate::cmd::batch_gitbutler::merge_back_hint(std::path::Path::new(&repo_path), group_id)
            {
                println!("{hint}");
            }
            Ok(())
        }
        WaitOutcome::Failed(task_ids) => {
            aid_error!("[aid] {} task(s) did not succeed: {}", task_ids.len(), task_ids.join(", "));
            std::process::exit(1);
        }
        WaitOutcome::TimedOut(running) => {
            let secs = timeout_secs.unwrap_or_default();
            aid_error!(
                "[aid] Timeout after {}s. Still running: {}",
                secs,
                running.join(", ")
            );
            std::process::exit(124);
        }
    }
}

pub(crate) async fn wait_for_task_ids(
    store: &Arc<Store>,
    task_ids: &[String],
    group: Option<&str>,
    exit_on_await: bool,
    timeout: Option<Duration>,
) -> Result<WaitOutcome> {
    if let Some(timeout) = timeout {
        match tokio::time::timeout(timeout, wait_for_task_ids_inner(store, task_ids, group, exit_on_await)).await {
            Ok(result) => result,
            Err(_) => Ok(WaitOutcome::TimedOut(still_running_task_ids(store, task_ids, group)?)),
        }
    } else {
        wait_for_task_ids_inner(store, task_ids, group, exit_on_await).await
    }
}

async fn wait_for_task_ids_inner(
    store: &Arc<Store>,
    task_ids: &[String],
    group: Option<&str>,
    exit_on_await: bool,
) -> Result<WaitOutcome> {
    let mut task_ids = task_ids.to_vec();
    let mut last_status: HashMap<String, String> = HashMap::new();
    let mut pending_since: HashMap<String, Instant> = HashMap::new();
    let mut completed = 0usize;
    let track_group_tasks = |task_ids: &mut Vec<String>, completed: usize| -> Result<bool> {
        let Some(group_id) = group else { return Ok(false); };
        let mut found = false;
        for task_id in current_running_ids(store, Some(group_id))? {
            if task_ids.contains(&task_id) {
                continue;
            }
            task_ids.push(task_id.clone());
            println!("[{}/{}] new task {} added to watch set", completed, task_ids.len(), task_id);
            found = true;
        }
        Ok(found)
    };

    loop {
        let _ = crate::runtime::background::check_zombie_tasks(store);
        track_group_tasks(&mut task_ids, completed)?;
        let mut remaining = 0usize;
        let total = task_ids.len();
        let mut failed = Vec::new();

        for task_id in &task_ids {
            // The worker removes its spec only after delivery/verification settle.
            // Observe that barrier BEFORE reading status, never reuse a pre-barrier Done.
            let settling = worker_is_settling(task_id)?;
            let Some(task) = store.get_task(task_id)? else {
                continue;
            };
            if settling && task.status.is_terminal() {
                remaining += 1;
                continue;
            }
            let status = task.status;
            let outcome = task.outcome();
            let status_text = status.label().to_string();
            let status_changed = last_status.insert(task_id.clone(), status_text.clone()) != Some(status_text);

            // Print status transitions with completion summary
            if status_changed {
                if status.is_terminal() {
                    completed += 1;
                    let duration = task.duration_ms
                        .map(|ms| {
                            let secs = ms / 1000;
                            if secs < 60 { format!("{secs}s") } else { format!("{}m {:02}s", secs / 60, secs % 60) }
                        })
                        .unwrap_or_else(|| "-".to_string());
                    let tokens = task.tokens
                        .map(|t| if t >= 1_000_000 { format!("{:.1}M", t as f64 / 1_000_000.0) } else if t >= 1_000 { format!("{:.1}k", t as f64 / 1_000.0) } else { t.to_string() })
                        .unwrap_or_else(|| "-".to_string());
                    let fail_reason = if status == TaskStatus::Failed {
                        store.latest_error(task_id)
                            .map(|r| format!(" — {r}"))
                            .unwrap_or_default()
                    } else {
                        String::new()
                    };
                    let status_label = outcome
                        .verification_tag()
                        .map(|tag| format!("{} [{tag}]", status.label()))
                        .unwrap_or_else(|| status.label().to_string());
                    println!(
                        "[{}/{}] {} {} ({}, {}tok, {}){fail_reason}",
                        completed, total, task_id, status_label, duration, tokens, cost::format_cost(task.cost_usd),
                    );
                } else {
                    println!("[{}/{}] {} {}", completed, total, task_id, status.label());
                }
            }

            if exit_on_await && status == TaskStatus::AwaitingInput && status_changed {
                let events = store.get_events(task_id)?;
                let prompt = events
                    .iter()
                    .rev()
                    .find_map(|e| {
                        e.metadata
                            .as_ref()
                            .and_then(|m| m.get("awaiting_prompt"))
                            .and_then(|v| v.as_str())
                    })
                    .unwrap_or("");
                println!("{} {}", task_id, prompt);
                println!("Use: aid respond {} \"your answer\"", task_id);
                return Ok(WaitOutcome::Completed);
            }

            if verification_is_in_flight(&task) {
                let started_at = pending_since.entry(task_id.clone()).or_insert_with(Instant::now);
                if started_at.elapsed() < VERIFY_TIMEOUT {
                    remaining += 1;
                } else {
                    failed.push(task_id.clone());
                }
            } else if matches!(outcome, TaskOutcome::InProgress) {
                pending_since.remove(task_id);
                remaining += 1;
            } else if !outcome.is_success() {
                pending_since.remove(task_id);
                failed.push(task_id.clone());
            } else {
                pending_since.remove(task_id);
            }
        }

        if remaining == 0 {
            if track_group_tasks(&mut task_ids, completed)? {
                continue;
            }
            println!("All {} task(s) completed.", total);
            failed.sort();
            failed.dedup();
            return Ok(if failed.is_empty() {
                WaitOutcome::Completed
            } else {
                WaitOutcome::Failed(failed)
            });
        }

        sleep(Duration::from_secs(2)).await;
    }
}

fn current_running_ids(store: &Arc<Store>, group: Option<&str>) -> Result<Vec<String>> {
    let mut tasks = store.list_tasks(TaskFilter::Active)?;
    settlement::include_worker_tasks(store, &mut tasks)?;
    if let Some(group_id) = group {
        tasks.retain(|t| t.workgroup_id.as_deref() == Some(group_id));
    }
    Ok(tasks.into_iter().map(|task| task.id.to_string()).collect())
}

fn still_running_task_ids(store: &Arc<Store>, task_ids: &[String], group: Option<&str>) -> Result<Vec<String>> {
    let mut tracked = task_ids.to_vec();
    if let Some(group_id) = group {
        for task_id in current_running_ids(store, Some(group_id))? {
            if !tracked.contains(&task_id) {
                tracked.push(task_id);
            }
        }
    }
    let mut running = Vec::new();
    for task_id in tracked {
        if worker_is_settling(&task_id)? {
            running.push(task_id);
            continue;
        }
        if let Some(task) = store.get_task(&task_id)? {
            if verification_is_in_flight(&task) {
                running.push(task_id);
                continue;
            }
            let outcome = task.outcome();
            if matches!(outcome, TaskOutcome::InProgress) {
                running.push(task_id);
            }
        }
    }
    Ok(running)
}
fn verification_is_in_flight(task: &Task) -> bool {
    matches!(task.status, TaskStatus::Done | TaskStatus::Merged)
        && task.verify_status == VerifyStatus::Pending
        && verify_required(task.verify.as_deref())
}

#[cfg(test)]
#[path = "wait_tests.rs"]
mod tests;

#[cfg(test)]
#[path = "wait_settlement_tests.rs"]
mod settlement_tests;