agent-manager/src/lab.rs

//! am lab (issue #62): run the same task on several agents and compare
//! duration, exit code, output and estimated cost. Tasks are plain YAML
//! files living in `<state_dir>/lab/` — versionable and replayable:
//! running the same file again reproduces the same benchmark.
//!
//! A task is the list of arguments passed to each agent's run command
//! (plus an optional working directory and timeout).

use crate::app::App;
use anyhow::{bail, Context, Result};
use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;
use std::path::{Path, PathBuf};
use std::process::{Command, Stdio};
use std::time::{Duration, Instant};
use wait_timeout::ChildExt;

/// Default timeout for one task run (seconds).
pub const DEFAULT_TASK_TIMEOUT_SECS: u64 = 600;
/// Name of the tasks directory inside the state directory.
pub const DIR_NAME: &str = "lab";

/// A benchmark task: a versionable YAML file run on every agent.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LabTask {
    /// Short name of the task (required).
    pub name: String,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub description: Option<String>,
    /// Arguments passed to every agent's run command.
    #[serde(default)]
    pub args: Vec<String>,
    /// Working directory of the run (default: the current directory).
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub cwd: Option<String>,
    /// Timeout in seconds (default 600).
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub timeout: Option<u64>,
}

/// Outcome of one agent on the task.
#[derive(Debug, Clone, Default, Serialize)]
pub struct LabRun {
    pub agent: String,
    /// ok | failed | timeout | error
    pub status: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub exit_code: Option<i32>,
    pub duration_ms: u64,
    /// Captured stdout of the run.
    pub output: String,
    #[serde(skip_serializing_if = "String::is_empty")]
    pub stderr: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tokens_in: Option<u64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tokens_out: Option<u64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub cost_usd: Option<f64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub error: Option<String>,
}

/// The comparative report (also the --json contract: stable fields,
/// runs sorted by agent name).
#[derive(Debug, Clone, Serialize)]
pub struct LabReport {
    pub task: String,
    pub runs: Vec<LabRun>,
    pub summary: LabSummary,
}

#[derive(Debug, Clone, Default, Serialize)]
pub struct LabSummary {
    pub runs: usize,
    pub ok: usize,
    pub failed: usize,
    pub timeout: usize,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub fastest: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub cheapest: Option<String>,
}

/// The tasks directory of the state directory.
pub fn dir(app: &App) -> PathBuf {
    app.events_dir().join(DIR_NAME)
}

/// Resolve a task path: a bare name goes into the lab directory (with
/// .yaml/.yml appended when the exact name does not exist), an explicit
/// path is honored as-is.
pub fn resolve_task(app: &App, dest: &Path) -> PathBuf {
    let has_sep = dest.components().count() > 1
        || dest.components().any(|c| {
            matches!(
                c,
                std::path::Component::ParentDir | std::path::Component::CurDir
            )
        });
    if has_sep || dest.is_absolute() {
        return dest.to_path_buf();
    }
    let base = dir(app).join(dest);
    if base.exists() {
        return base;
    }
    for ext in ["yaml", "yml"] {
        let candidate = dir(app).join(format!("{}.{}", dest.display(), ext));
        if candidate.exists() {
            return candidate;
        }
    }
    base
}

/// Task files of the lab directory (yaml/yml), sorted by name.
pub fn list_tasks(app: &App) -> Vec<PathBuf> {
    let mut out: Vec<PathBuf> = std::fs::read_dir(dir(app))
        .map(|rd| {
            rd.flatten()
                .map(|e| e.path())
                .filter(|p| p.is_file())
                .filter(|p| {
                    matches!(
                        p.extension().and_then(|e| e.to_str()),
                        Some("yaml" | "yml")
                    )
                })
                .collect()
        })
        .unwrap_or_default();
    out.sort();
    out
}

/// Load and validate a task file.
pub fn load_task(path: &Path) -> Result<LabTask> {
    let text = std::fs::read_to_string(path)
        .with_context(|| format!("cannot read {}", path.display()))?;
    let task: LabTask = serde_yaml::from_str(&text)
        .with_context(|| format!("{}: invalid task YAML", path.display()))?;
    if task.name.trim().is_empty() {
        bail!("{}: 'name' is required", path.display());
    }
    Ok(task)
}

/// Split the --agents list ("a,b, c") into trimmed names.
pub fn split_agents(list: &str) -> Vec<String> {
    list.split(',')
        .map(|s| s.trim().to_string())
        .filter(|s| !s.is_empty())
        .collect()
}

/// A pre-resolved run context: owned data, safe to move across threads.
struct RunCtx {
    agent: String,
    program: String,
    args: Vec<String>,
    env: BTreeMap<String, String>,
    cwd: Option<PathBuf>,
    timeout: Duration,
    cost_model: Option<String>,
    sheet: BTreeMap<String, crate::config::CostModel>,
}

/// Run the task on every agent and build the comparative report.
/// Sequential by default; `parallel` spawns one thread per agent.
pub fn run_benchmark(
    app: &App,
    agents: &[String],
    task: &LabTask,
    parallel: bool,
    timeout_override: Option<u64>,
) -> Result<LabReport> {
    let timeout = Duration::from_secs(
        timeout_override.unwrap_or(task.timeout.unwrap_or(DEFAULT_TASK_TIMEOUT_SECS)),
    );
    let task_cwd = task.cwd.as_deref().map(crate::config::expand_path);
    if let Some(cwd) = &task_cwd {
        if !cwd.is_dir() {
            bail!("task cwd does not exist: {}", cwd.display());
        }
    }
    let sheet = app
        .config
        .settings
        .cost_models
        .clone()
        .unwrap_or_else(crate::config::default_cost_models);

    let mut ctxs = Vec::new();
    for name in agents {
        let agent = crate::commands::require_agent(app, name)?;
        let exec = crate::commands::resolve_exec(app, agent, &task.args, &BTreeMap::new())?;
        ctxs.push(RunCtx {
            agent: agent.name.clone(),
            program: exec.program,
            args: exec.args,
            env: exec.env,
            cwd: task_cwd.clone(),
            timeout,
            cost_model: agent.cost_model.clone(),
            sheet: sheet.clone(),
        });
    }

    let mut runs: Vec<LabRun> = if parallel {
        let handles: Vec<_> = ctxs
            .into_iter()
            .map(|ctx| std::thread::spawn(move || run_one(&ctx)))
            .collect();
        handles
            .into_iter()
            .map(|h| {
                h.join().unwrap_or_else(|_| LabRun {
                    agent: "?".to_string(),
                    status: "error".to_string(),
                    error: Some("runner thread panicked".to_string()),
                    ..Default::default()
                })
            })
            .collect()
    } else {
        ctxs.iter().map(run_one).collect()
    };
    // Stable order for --json: always sorted by agent name.
    runs.sort_by(|a, b| a.agent.cmp(&b.agent));

    // Journal one Lab event per run (axe 1: everything is traceable).
    for run in &runs {
        let mut ev = crate::events::Event::now(crate::events::EventKind::Lab)
            .with_agent(run.agent.clone())
            .with_reason(run.status.clone())
            .with_args(task.args.clone())
            .with_duration(run.duration_ms / 1000);
        if let Some(code) = run.exit_code {
            ev = ev.with_exit_code(code);
        }
        if let Some(cost) = run.cost_usd {
            ev = ev
                .with_usage(run.tokens_in.unwrap_or(0), run.tokens_out.unwrap_or(0))
                .with_cost(cost);
        }
        app.emit(&ev);
    }

    Ok(LabReport {
        task: task.name.clone(),
        summary: summarize(&runs),
        runs,
    })
}

/// Run one agent on the task with a timeout; captures stdout/stderr.
fn run_one(ctx: &RunCtx) -> LabRun {
    let started = Instant::now();
    let mut cmd = Command::new(&ctx.program);
    cmd.args(&ctx.args)
        .envs(&ctx.env)
        .stdin(Stdio::null())
        .stdout(Stdio::piped())
        .stderr(Stdio::piped());
    if let Some(cwd) = &ctx.cwd {
        cmd.current_dir(cwd);
    }
    set_process_group(&mut cmd);

    let mut child = match cmd.spawn() {
        Ok(c) => c,
        Err(e) => {
            return LabRun {
                agent: ctx.agent.clone(),
                status: "error".to_string(),
                error: Some(format!("cannot spawn: {e:#}")),
                duration_ms: started.elapsed().as_millis() as u64,
                ..Default::default()
            }
        }
    };
    let stdout = read_pipe(child.stdout.take());
    let stderr = read_pipe(child.stderr.take());

    let deadline = started + ctx.timeout;
    let outcome = loop {
        match child.wait_timeout(Duration::from_millis(200)) {
            Ok(Some(status)) => break Ok(status),
            Ok(None) => {
                if Instant::now() >= deadline {
                    crate::process::kill_tree(child.id());
                    let _ = child.wait();
                    break Err("timeout".to_string());
                }
            }
            Err(e) => {
                crate::process::kill_tree(child.id());
                let _ = child.wait();
                break Err(format!("cannot wait: {e:#}"));
            }
        }
    };
    let out = stdout.join().unwrap_or_default();
    let err = stderr.join().unwrap_or_default();

    let (status, exit_code, error) = match outcome {
        Ok(st) => {
            let code = st.code().unwrap_or(-1);
            (
                if code == 0 { "ok" } else { "failed" }.to_string(),
                Some(code),
                None,
            )
        }
        Err(e) => {
            let kind = if e == "timeout" { "timeout" } else { "error" };
            (kind.to_string(), None, Some(e))
        }
    };
    let mut run = LabRun {
        agent: ctx.agent.clone(),
        status,
        exit_code,
        duration_ms: started.elapsed().as_millis() as u64,
        output: out,
        stderr: err,
        error,
        ..Default::default()
    };

    // Estimated cost when the agent exposes its usage (axe 1, issue #49).
    let cwd_str = ctx.cwd.as_deref().map(|p| p.display().to_string());
    if let Some(usage) = crate::costs::collect_usage(&ctx.agent, cwd_str.as_deref()) {
        let cost = crate::costs::estimate_cost(ctx.cost_model.as_deref(), usage, &ctx.sheet);
        run.tokens_in = Some(usage.tokens_in);
        run.tokens_out = Some(usage.tokens_out);
        run.cost_usd = Some(cost);
    }
    run
}

/// Compute the summary rows of the report.
fn summarize(runs: &[LabRun]) -> LabSummary {
    let mut s = LabSummary {
        runs: runs.len(),
        ..Default::default()
    };
    let mut fastest: Option<(&str, u64)> = None;
    let mut cheapest: Option<(&str, f64)> = None;
    for r in runs {
        match r.status.as_str() {
            "ok" => s.ok += 1,
            "failed" => s.failed += 1,
            "timeout" => s.timeout += 1,
            _ => {}
        }
        if r.status == "ok" && fastest.map(|(_, d)| r.duration_ms < d).unwrap_or(true) {
            fastest = Some((&r.agent, r.duration_ms));
        }
        if let Some(cost) = r.cost_usd {
            if cheapest.map(|(_, c)| cost < c).unwrap_or(true) {
                cheapest = Some((&r.agent, cost));
            }
        }
    }
    s.fastest = fastest.map(|(a, _)| a.to_string());
    s.cheapest = cheapest.map(|(a, _)| a.to_string());
    s
}

/// Read a child pipe to the end on a dedicated thread (avoids pipe
/// deadlocks while the parent waits with a timeout).
fn read_pipe<R: std::io::Read + Send + 'static>(pipe: Option<R>) -> std::thread::JoinHandle<String> {
    std::thread::spawn(move || {
        use std::io::Read;
        let mut buf = String::new();
        if let Some(mut p) = pipe {
            let _ = p.read_to_string(&mut buf);
        }
        buf
    })
}

/// Put the child in its own process group so `kill_tree` can stop it with
/// its whole subtree on timeout.
fn set_process_group(cmd: &mut Command) {
    #[cfg(unix)]
    {
        use std::os::unix::process::CommandExt;
        cmd.process_group(0);
    }
    #[cfg(windows)]
    {
        use std::os::windows::process::CommandExt;
        const CREATE_NEW_PROCESS_GROUP: u32 = 0x0000_0200;
        cmd.creation_flags(CREATE_NEW_PROCESS_GROUP);
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    use super::*;
    use crate::cli::Cli;
    use clap::Parser;

    fn test_app(tag: &str) -> App {
        let dir = tempfile::tempdir().unwrap();
        let cfg = dir.path().join("config.yaml");
        std::fs::write(
            &cfg,
            concat!(
                "version: \"1.0\"\n",
                "settings:\n",
                "  auto_install_deps: false\n",
                "  confirm_before_run: false\n",
                "agents: []\n",
            ),
        )
        .unwrap();
        let cli = Cli::parse_from(["am", "--config", cfg.to_str().unwrap(), "list"]);
        let mut app = App::from_cli(cli).unwrap();
        app.paths.state_file = dir.path().join(format!("state-{tag}.json"));
        app.paths.log_dir = dir.path().join("logs");
        app
    }

    /// An app whose catalog defines three shell-based fake agents:
    /// bench-ok (echoes its first argument), bench-fail (exit 3),
    /// bench-slow (infinite loop, for the timeout test).
    fn bench_app(tag: &str) -> App {
        let dir = tempfile::tempdir().unwrap();
        let cfg = dir.path().join("config.yaml");
        std::fs::write(
            &cfg,
            concat!(
                "version: \"1.0\"\n",
                "settings:\n",
                "  auto_install_deps: false\n",
                "  confirm_before_run: false\n",
                "agents:\n",
                "  - name: bench-ok\n",
                "    installable: false\n",
                "    run: sh\n",
                "    args: [\"-c\", \"echo out: $0\"]\n",
                "  - name: bench-fail\n",
                "    installable: false\n",
                "    run: sh\n",
                "    args: [\"-c\", \"echo boom >&2; exit 3\"]\n",
                "  - name: bench-slow\n",
                "    installable: false\n",
                "    run: sh\n",
                "    args: [\"-c\", \"while true; do :; done\"]\n",
            ),
        )
        .unwrap();
        let cli = Cli::parse_from(["am", "--config", cfg.to_str().unwrap(), "list"]);
        let mut app = App::from_cli(cli).unwrap();
        app.paths.state_file = dir.path().join(format!("state-{tag}.json"));
        app.paths.log_dir = dir.path().join("logs");
        app
    }

    #[test]
    fn split_agents_trims_and_filters() {
        assert_eq!(split_agents("a,b"), vec!["a".to_string(), "b".to_string()]);
        assert_eq!(split_agents(" a , b, c "), vec!["a", "b", "c"]);
        assert_eq!(split_agents(" , "), Vec::<String>::new());
    }

    #[test]
    fn task_loads_with_defaults() {
        let dir = tempfile::tempdir().unwrap();
        let f = dir.path().join("hello.yaml");
        std::fs::write(
            &f,
            "name: hello\ndescription: test\nargs: [\"--prompt\", \"coucou\"]\n",
        )
        .unwrap();
        let t = load_task(&f).unwrap();
        assert_eq!(t.name, "hello");
        assert_eq!(t.args, vec!["--prompt".to_string(), "coucou".to_string()]);
        assert_eq!(t.timeout, None);
        assert!(load_task(&dir.path().join("missing.yaml")).is_err());
        std::fs::write(&f, "args: [x]\n").unwrap();
        assert!(load_task(&f).is_err(), "name is required");
    }

    #[test]
    fn bare_name_resolves_into_the_lab_dir() {
        let app = test_app("dest");
        let resolved = resolve_task(&app, Path::new("hello"));
        assert_eq!(resolved, app.events_dir().join(DIR_NAME).join("hello"));
        let abs = PathBuf::from("C:/tmp/hello.yaml");
        assert_eq!(resolve_task(&app, &abs), abs);
    }

    #[test]
    fn list_tasks_filters_yaml_files() {
        let app = test_app("list");
        let d = dir(&app);
        std::fs::create_dir_all(&d).unwrap();
        std::fs::write(d.join("a.yaml"), "name: a\n").unwrap();
        std::fs::write(d.join("b.yml"), "name: b\n").unwrap();
        std::fs::write(d.join("notes.md"), "nope").unwrap();
        let tasks = list_tasks(&app);
        assert_eq!(tasks.len(), 2);
        assert!(tasks[0].display().to_string().ends_with("a.yaml"));
        assert!(tasks[1].display().to_string().ends_with("b.yml"));
    }

    #[test]
    fn benchmark_runs_sequentially_and_reports() {
        let app = bench_app("seq");
        let task = LabTask {
            name: "hello".into(),
            description: None,
            args: vec!["hello".to_string()],
            cwd: None,
            timeout: Some(30),
        };
        let agents = vec!["bench-ok".to_string(), "bench-fail".to_string()];
        let report = run_benchmark(&app, &agents, &task, false, None).unwrap();
        assert_eq!(report.task, "hello");
        assert_eq!(report.runs.len(), 2);
        let ok = report.runs.iter().find(|r| r.agent == "bench-ok").unwrap();
        assert_eq!(ok.status, "ok");
        assert_eq!(ok.exit_code, Some(0));
        assert!(ok.output.contains("out: hello"), "got: {}", ok.output);
        let fail = report.runs.iter().find(|r| r.agent == "bench-fail").unwrap();
        assert_eq!(fail.status, "failed");
        assert_eq!(fail.exit_code, Some(3));
        assert!(fail.stderr.contains("boom"));
        // Summary counts.
        assert_eq!(report.summary.ok, 1);
        assert_eq!(report.summary.failed, 1);
        assert_eq!(report.summary.fastest.as_deref(), Some("bench-ok"));
    }

    #[test]
    fn benchmark_is_replayable_identically() {
        let app = bench_app("replay");
        let task = LabTask {
            name: "hello".into(),
            description: None,
            args: vec!["hello".to_string()],
            cwd: None,
            timeout: Some(30),
        };
        let agents = vec!["bench-ok".to_string()];
        let r1 = run_benchmark(&app, &agents, &task, false, None).unwrap();
        let r2 = run_benchmark(&app, &agents, &task, false, None).unwrap();
        assert_eq!(r1.runs.len(), r2.runs.len());
        assert_eq!(r1.runs[0].status, r2.runs[0].status);
        assert_eq!(r1.runs[0].exit_code, r2.runs[0].exit_code);
        assert_eq!(r1.runs[0].output, r2.runs[0].output);
        // The --json contract is stable: same fields on both serializations.
        let j1 = serde_json::to_string(&r1).unwrap();
        let j2 = serde_json::to_string(&r2).unwrap();
        assert!(j1.contains("\"task\":\"hello\""));
        assert!(j1.contains("\"status\":\"ok\""));
        assert!(j2.contains("\"status\":\"ok\""));
    }

    #[test]
    fn timeout_kills_the_run() {
        let app = bench_app("to");
        let task = LabTask {
            name: "slow".into(),
            description: None,
            args: vec![],
            cwd: None,
            timeout: Some(1),
        };
        let agents = vec!["bench-slow".to_string()];
        let report = run_benchmark(&app, &agents, &task, false, None).unwrap();
        assert_eq!(report.runs[0].status, "timeout");
        assert_eq!(report.summary.timeout, 1);
    }

    #[test]
    fn parallel_mode_runs_every_agent() {
        let app = bench_app("par");
        let task = LabTask {
            name: "hello".into(),
            description: None,
            args: vec!["hello".to_string()],
            cwd: None,
            timeout: Some(30),
        };
        let agents = vec!["bench-ok".to_string(), "bench-fail".to_string()];
        let report = run_benchmark(&app, &agents, &task, true, None).unwrap();
        assert_eq!(report.runs.len(), 2);
        // Runs are sorted by agent name: stable --json.
        assert_eq!(report.runs[0].agent, "bench-fail");
        assert_eq!(report.runs[1].agent, "bench-ok");
        assert_eq!(report.summary.ok, 1);
    }

    #[test]
    fn missing_cwd_is_an_error() {
        let app = bench_app("cwd");
        let task = LabTask {
            name: "bad".into(),
            description: None,
            args: vec![],
            cwd: Some("C:/does/not/exist-anywhere".to_string()),
            timeout: None,
        };
        assert!(run_benchmark(&app, &["bench-ok".to_string()], &task, false, None).is_err());
    }

    #[test]
    fn unknown_agent_is_an_error() {
        let app = bench_app("unk");
        let task = LabTask {
            name: "t".into(),
            description: None,
            args: vec![],
            cwd: None,
            timeout: None,
        };
        assert!(run_benchmark(&app, &["ghost".to_string()], &task, false, None).is_err());
    }
}