clawmates/crates/cm-api/src/benchmark_runner.rs

//! Before/after benchmark runner — Slice 7.
//!
//! Executes a benchmark command inside the mission's team container
//! (or, when no team container exists, the shared runtime) and
//! records the resulting metrics as a `benchmark_snapshots` row
//! keyed on (phase_id, iteration).
//!
//! Two entry points:
//!   - `baseline(mission_id, phase_id)` — records iteration 0's
//!     `before_metrics` before any coding pass. Idempotent.
//!   - `after_iteration(mission_id, phase_id, iteration)` — records
//!     `after_metrics` for the given iteration and computes `delta`
//!     from iteration 0's baseline.
//!
//! Harnesses auto-detected via `phases.config.harness`:
//!   auto | criterion | cargo_bench | vitest_bench | pytest_bench | shell
//!
//! The `shell` fallback runs an arbitrary command and captures its
//! stdout as an opaque string; the canvas renders it verbatim so
//! bespoke harnesses aren't a blocker.

use serde_json::{json, Value};
use sqlx::PgPool;
use sqlx::Row;
use std::time::Duration;
use uuid::Uuid;

/// Ceiling for one benchmark command. Benchmarks are slow by nature — this is
/// a guard against a wedged run holding the phase open, not a performance
/// budget.
const BENCH_TIMEOUT: Duration = Duration::from_secs(1800);

/// Which slot in `benchmark_snapshots` the run should populate.
#[derive(Debug, Clone, Copy)]
pub enum Slot {
    Before,
    After,
}

/// Which invocation strategy to use — matches
/// `phases.config.harness` in templates/workflows/*.toml.
#[derive(Debug, Clone)]
pub enum Harness {
    /// Detect from repo files (Cargo.toml, package.json, pyproject).
    Auto,
    /// `cargo bench --bench <name> -- --output-format json`
    Criterion { bench_name: Option<String> },
    /// `cargo bench` — captures raw stdout.
    CargoBench,
    /// `pnpm exec vitest bench --run --reporter=json`
    VitestBench,
    /// `pytest --benchmark-only --benchmark-json=<tmp>`
    PytestBench,
    /// Raw shell command; stdout captured as a string blob.
    Shell { cmd: String },
}

impl Harness {
    fn command(&self) -> Vec<String> {
        match self {
            Harness::Criterion { bench_name } => {
                let bench = bench_name.as_deref().unwrap_or("");
                if bench.is_empty() {
                    vec![
                        "cargo".into(),
                        "bench".into(),
                        "--".into(),
                        "--output-format=bencher".into(),
                    ]
                } else {
                    vec![
                        "cargo".into(),
                        "bench".into(),
                        "--bench".into(),
                        bench.into(),
                        "--".into(),
                        "--output-format=bencher".into(),
                    ]
                }
            }
            Harness::CargoBench => vec!["cargo".into(), "bench".into()],
            Harness::VitestBench => vec![
                "pnpm".into(),
                "exec".into(),
                "vitest".into(),
                "bench".into(),
                "--run".into(),
                "--reporter=json".into(),
            ],
            Harness::PytestBench => vec![
                "pytest".into(),
                "--benchmark-only".into(),
                "--benchmark-json=/tmp/pytest-bench.json".into(),
            ],
            Harness::Shell { cmd } => vec!["sh".into(), "-c".into(), cmd.clone()],
            Harness::Auto => vec![
                "sh".into(),
                "-c".into(),
                "echo 'auto harness must be resolved before command()'".into(),
            ],
        }
    }

    fn driver_name(&self) -> &'static str {
        match self {
            Harness::Auto => "auto",
            Harness::Criterion { .. } => "criterion",
            Harness::CargoBench => "cargo_bench",
            Harness::VitestBench => "vitest_bench",
            Harness::PytestBench => "pytest_bench",
            Harness::Shell { .. } => "shell",
        }
    }
}

/// Parse phase config into a Harness. Unknown / missing → Auto.
pub fn harness_from_config(cfg: &Value) -> Harness {
    let harness = cfg
        .get("harness")
        .and_then(|v| v.as_str())
        .unwrap_or("auto");
    match harness {
        "criterion" => Harness::Criterion {
            bench_name: cfg
                .get("bench_name")
                .and_then(|v| v.as_str())
                .map(|s| s.to_string()),
        },
        "cargo_bench" => Harness::CargoBench,
        "vitest_bench" => Harness::VitestBench,
        "pytest_bench" => Harness::PytestBench,
        "shell" => Harness::Shell {
            cmd: cfg
                .get("cmd")
                .and_then(|v| v.as_str())
                .unwrap_or("echo 'no cmd configured'")
                .to_string(),
        },
        _ => Harness::Auto,
    }
}

/// Execute the harness inside the team's container and return
/// `(metrics_json, driver_name)`. Errors are formatted for the
/// canvas — a per-run failure shows up as text, not a stack trace.
pub async fn run(
    pool: &PgPool,
    mission_id: Uuid,
    phase_id: Uuid,
    _slot: Slot,
) -> Result<(Value, String), String> {
    let phase_cfg = phase_config(pool, phase_id).await?;
    let harness = harness_from_config(&phase_cfg);
    let harness = match harness {
        Harness::Auto => auto_detect(pool, mission_id)
            .await
            .unwrap_or(Harness::CargoBench),
        other => other,
    };

    let (container, workdir) = exec_target(pool, mission_id).await?;
    let cmd = harness.command();
    let raw = docker_exec(&container, &workdir, &cmd)
        .await
        .map_err(|e| format!("exec {cmd:?}: {e}"))?;
    let metrics = parse_output(&raw, &harness);
    Ok((metrics, harness.driver_name().to_string()))
}

/// Record iteration 0's baseline (before_metrics). Idempotent — subsequent
/// baseline calls overwrite the before slot only, never the after slot.
pub async fn baseline(pool: &PgPool, mission_id: Uuid, phase_id: Uuid) -> Result<(), String> {
    let (metrics, driver) = run(pool, mission_id, phase_id, Slot::Before).await?;
    cm_db::repo::missions::upsert_benchmark_snapshot(
        pool,
        cm_db::repo::missions::UpsertBenchmarkSnapshot {
            mission_id,
            phase_id,
            iteration: 0,
            before_metrics: Some(&metrics),
            after_metrics: None,
            delta: None,
            driver: Some(&driver),
        },
    )
    .await
    .map_err(|e| format!("upsert baseline: {e}"))?;
    Ok(())
}

/// Record post-iteration metrics and compute the delta vs. baseline.
pub async fn after_iteration(
    pool: &PgPool,
    mission_id: Uuid,
    phase_id: Uuid,
    iteration: i32,
) -> Result<(), String> {
    if iteration <= 0 {
        return Err("iteration must be ≥ 1 (0 is the baseline slot)".into());
    }
    let (metrics, driver) = run(pool, mission_id, phase_id, Slot::After).await?;
    let baseline = load_before(pool, phase_id).await?;
    let delta = baseline
        .as_ref()
        .map(|b| compute_delta(b, &metrics))
        .unwrap_or_else(|| json!({ "note": "no baseline recorded — run baseline() first" }));
    cm_db::repo::missions::upsert_benchmark_snapshot(
        pool,
        cm_db::repo::missions::UpsertBenchmarkSnapshot {
            mission_id,
            phase_id,
            iteration,
            before_metrics: None,
            after_metrics: Some(&metrics),
            delta: Some(&delta),
            driver: Some(&driver),
        },
    )
    .await
    .map_err(|e| format!("upsert after: {e}"))?;
    Ok(())
}

// ── Internals ───────────────────────────────────────────────────

async fn phase_config(pool: &PgPool, phase_id: Uuid) -> Result<Value, String> {
    let row = sqlx::query("SELECT config FROM mission_phases WHERE id = $1")
        .bind(phase_id)
        .fetch_optional(pool)
        .await
        .map_err(|e| format!("load phase: {e}"))?;
    Ok(row
        .map(|r| r.get::<Value, _>("config"))
        .unwrap_or_else(|| json!({})))
}

/// Post-task-#23: shared runtime container + per-mission working dir.
/// See security_scan::exec_target for the same convention.
async fn exec_target(
    pool: &PgPool,
    mission_id: Uuid,
) -> Result<(String, std::path::PathBuf), String> {
    let repo_id: Option<Uuid> = sqlx::query_scalar("SELECT repo_id FROM missions WHERE id = $1")
        .bind(mission_id)
        .fetch_optional(pool)
        .await
        .map_err(|e| format!("resolve mission repo: {e}"))?
        .flatten();
    if repo_id.is_none() {
        return Err(
            "mission has no repo bound — benchmark requires a repository under mission.repo_id"
                .into(),
        );
    }
    let container = std::env::var("CLAWMATES_RUNTIME_CONTAINER")
        .unwrap_or_else(|_| "clawmates-runtime".to_string());
    let root = std::env::var("CLAWMATES_MISSIONS_ROOT")
        .unwrap_or_else(|_| "/var/lib/clawmates-missions".to_string());
    let workdir = std::path::PathBuf::from(root)
        .join(mission_id.to_string())
        .join("repo");
    Ok((container, workdir))
}

async fn load_before(pool: &PgPool, phase_id: Uuid) -> Result<Option<Value>, String> {
    let row = sqlx::query(
        "SELECT before_metrics
         FROM benchmark_snapshots
         WHERE phase_id = $1 AND iteration = 0",
    )
    .bind(phase_id)
    .fetch_optional(pool)
    .await
    .map_err(|e| format!("load baseline: {e}"))?;
    Ok(row.and_then(|r| {
        r.try_get::<Option<Value>, _>("before_metrics")
            .ok()
            .flatten()
    }))
}

/// Peek at repo layout inside the container to pick a sensible
/// harness — Cargo.toml → CargoBench, package.json with vitest →
/// VitestBench, pyproject with pytest-benchmark → PytestBench.
async fn auto_detect(pool: &PgPool, mission_id: Uuid) -> Result<Harness, String> {
    let (container, workdir) = exec_target(pool, mission_id).await?;
    let listing = docker_exec(&container, &workdir, &["ls".into()])
        .await
        .unwrap_or_default();
    if listing.contains("Cargo.toml") {
        return Ok(Harness::CargoBench);
    }
    if listing.contains("package.json") {
        return Ok(Harness::VitestBench);
    }
    if listing.contains("pyproject.toml") {
        return Ok(Harness::PytestBench);
    }
    Ok(Harness::Shell {
        cmd: "echo 'no harness detected — configure phase.config.harness'".into(),
    })
}

/// Run a benchmark command in the runtime container.
///
/// Uses the Docker API, not the `docker` CLI — the server image ships no such
/// binary, so this previously failed to spawn and every benchmark returned a
/// spawn error as its "result".
async fn docker_exec(
    container: &str,
    workdir: &std::path::Path,
    cmd: &[String],
) -> Result<String, String> {
    let docker = crate::container_exec::connect()?;
    let workdir_s = workdir.display().to_string();
    let out = crate::container_exec::exec(&docker, container, Some(&workdir_s), cmd, BENCH_TIMEOUT)
        .await?;
    // Benchmark harnesses split their reporting across both streams (criterion
    // writes results to stdout, cargo writes compilation to stderr), so the
    // caller needs both to make sense of a run.
    Ok(out.combined())
}

fn parse_output(raw: &str, harness: &Harness) -> Value {
    // Try JSON first (criterion --output-format=bencher lines are
    // NOT strictly JSON; but pytest-benchmark + vitest-bench are).
    if let Ok(v) = serde_json::from_str::<Value>(raw) {
        return v;
    }
    // Bencher-format lines: `test foo::bar ... bench:  123 ns/iter (+/- 45)`
    let mut samples: Vec<Value> = Vec::new();
    for line in raw.lines() {
        if let Some((name, nanos, plusminus)) = parse_bencher_line(line) {
            samples.push(json!({
                "name": name,
                "ns_per_iter": nanos,
                "plusminus_ns": plusminus,
            }));
        }
    }
    if !samples.is_empty() {
        return json!({
            "format": "bencher",
            "samples": samples,
        });
    }
    // Fallback: opaque text blob so nothing is lost.
    json!({
        "format": "text",
        "driver": harness.driver_name(),
        "stdout": raw,
    })
}

fn parse_bencher_line(line: &str) -> Option<(&str, u64, u64)> {
    // `test <name> ... bench:  <ns> ns/iter (+/- <plus>)`
    let after_test = line.strip_prefix("test ")?;
    let (name, tail) = after_test.split_once(" ... bench:")?;
    let tail = tail.trim();
    let (ns_str, rest) = tail.split_once(" ns/iter")?;
    let nanos: u64 = ns_str.trim().replace(',', "").parse().ok()?;
    let pm = rest
        .split_once("(+/-")
        .and_then(|(_, r)| r.split_once(')'))
        .and_then(|(v, _)| v.trim().replace(',', "").parse::<u64>().ok())
        .unwrap_or(0);
    Some((name, nanos, pm))
}

/// Compute a delta between before + after metrics. For structured
/// samples we compute per-sample percent change; for opaque text we
/// return a marker noting the shapes couldn't be diffed.
fn compute_delta(before: &Value, after: &Value) -> Value {
    let before_samples = before.get("samples").and_then(|v| v.as_array());
    let after_samples = after.get("samples").and_then(|v| v.as_array());
    if let (Some(bs), Some(as_)) = (before_samples, after_samples) {
        let mut diffs: Vec<Value> = Vec::new();
        for a in as_ {
            let Some(name) = a.get("name").and_then(|v| v.as_str()) else {
                continue;
            };
            let b_ns = bs
                .iter()
                .find(|b| b.get("name").and_then(|v| v.as_str()) == Some(name))
                .and_then(|b| b.get("ns_per_iter").and_then(|v| v.as_u64()));
            let a_ns = a.get("ns_per_iter").and_then(|v| v.as_u64());
            if let (Some(b_ns), Some(a_ns)) = (b_ns, a_ns) {
                if b_ns > 0 {
                    let pct = (a_ns as f64 - b_ns as f64) / (b_ns as f64) * 100.0;
                    diffs.push(json!({
                        "name": name,
                        "before_ns": b_ns,
                        "after_ns": a_ns,
                        "delta_pct": pct,
                        "direction": if pct < 0.0 { "improved" } else { "regressed" },
                    }));
                }
            }
        }
        return json!({ "kind": "bencher_diff", "samples": diffs });
    }
    json!({ "kind": "opaque", "note": "before/after not structurally comparable" })
}