hyperhive/hive-c0re/src/hive_stats.rs

//! Hive-wide turn-stats aggregation for the dashboard's swarm stats
//! view. Reads every agent's per-agent
//! `hyperhive-turn-stats.sqlite` read-only and rolls the rows up into
//! swarm totals + a per-agent rollup + model mix + a *labelled*
//! cost estimate.
//!
//! Why re-read the rows here instead of reusing `hive-ag3nt`'s
//! `stats.rs`: that module lives in a different crate (the agent
//! harness) which hive-c0re can't import. The stable contract is the
//! turn-stats *schema*, so we run a focused query against it. If we
//! ever want a single source of truth, the row-read + aggregation can
//! be lifted into a shared crate — overkill for now.
//!
//! Privsep: the sqlite files are mode 0644 owned by the agent user;
//! `hive-core` reads them fine (same as `events_vacuum`). We open
//! read-only so an in-flight harness writer never blocks us.

use std::collections::HashMap;
use std::path::Path;
use std::time::{Duration, SystemTime, UNIX_EPOCH};

use rusqlite::{Connection, OpenFlags};
use serde::{Deserialize, Serialize};

use crate::coordinator::Coordinator;

/// Window accepted by `/api/stats-hive?window=`. Maps to a lookback
/// span; the hive view is a flat rollup (no per-bucket trend — the
/// per-agent `/stats` page owns the trend charts).
#[derive(Debug, Clone, Copy)]
pub enum Window {
    Hour,
    FourHour,
    Day,
    ThreeDay,
    Week,
    Month,
}

impl Window {
    #[must_use]
    pub fn parse(s: &str) -> Self {
        match s {
            "1h" => Self::Hour,
            "4h" => Self::FourHour,
            "3d" => Self::ThreeDay,
            "7d" => Self::Week,
            "30d" => Self::Month,
            _ => Self::Day,
        }
    }

    fn label(self) -> &'static str {
        match self {
            Self::Hour => "1h",
            Self::FourHour => "4h",
            Self::Day => "24h",
            Self::ThreeDay => "3d",
            Self::Week => "7d",
            Self::Month => "30d",
        }
    }

    fn span_secs(self) -> i64 {
        match self {
            Self::Hour => 3600,
            Self::FourHour => 4 * 3600,
            Self::Day => 24 * 3600,
            Self::ThreeDay => 3 * 24 * 3600,
            Self::Week => 7 * 24 * 3600,
            Self::Month => 30 * 24 * 3600,
        }
    }
}

/// Approximate USD price per **million** tokens for one model class.
/// The four fields map to the four token kinds the turn-stats schema
/// records. This is a deliberately rough estimate — Anthropic list
/// pricing drifts, so the dashboard labels the figure as an estimate.
///
/// Operators keep the table current without a code change via the
/// `services.hyperhive.modelPrices` nix option (passed to
/// `hive-c0re serve --model-prices <json>`, held on the `Coordinator`
/// as a [`PriceTable`]). Any model not covered by the operator table
/// falls back to [`builtin_prices`].
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
pub struct Prices {
    pub input: f64,
    pub output: f64,
    pub cache_read: f64,
    pub cache_write: f64,
}

/// Operator-tunable model→price map. Keys are matched
/// case-insensitively as a substring of the model id (e.g. `"sonnet"`
/// matches `"claude-sonnet-4-5"`); the longest matching key wins so a
/// specific entry beats a generic family name. An empty map means
/// every model uses [`builtin_prices`].
pub type PriceTable = HashMap<String, Prices>;

/// Built-in fallback pricing, per **million** tokens. Used when the
/// operator's [`PriceTable`] has no key matching the model. Figures are
/// current Anthropic list prices for the Claude 4.x family (Opus 4.x,
/// Sonnet 4.x, Haiku 4.5); `cache_write` uses the 1-hour cache-TTL price,
/// which is the default through the Claude subscription the agents run
/// on. Keep in sync with the `services.hyperhive.modelPrices` nix default
/// (`nix/modules/hive-c0re.nix`).
fn builtin_prices(model: &str) -> Prices {
    let m = model.to_ascii_lowercase();
    if m.contains("opus") {
        Prices {
            input: 5.0,
            output: 25.0,
            cache_read: 0.5,
            cache_write: 10.0,
        }
    } else if m.contains("haiku") {
        Prices {
            input: 1.0,
            output: 5.0,
            cache_read: 0.1,
            cache_write: 2.0,
        }
    } else {
        // sonnet + unknown fallback
        Prices {
            input: 3.0,
            output: 15.0,
            cache_read: 0.3,
            cache_write: 6.0,
        }
    }
}

/// Resolve the price for `model`: prefer the operator-provided
/// `table` (longest matching key wins, so a specific
/// `claude-3-5-sonnet` entry beats a generic `sonnet`), else fall back
/// to [`builtin_prices`].
fn resolve_prices(model: &str, table: &PriceTable) -> Prices {
    let m = model.to_ascii_lowercase();
    let mut best: Option<(usize, Prices)> = None;
    for (k, v) in table {
        let key = k.to_ascii_lowercase();
        if key.is_empty() || !m.contains(&key) {
            continue;
        }
        let better = match best {
            Some((blen, _)) => key.len() > blen,
            None => true,
        };
        if better {
            best = Some((key.len(), *v));
        }
    }
    match best {
        Some((_, p)) => p,
        None => builtin_prices(&m),
    }
}

#[derive(Debug, Serialize)]
pub struct KeyCount {
    pub key: String,
    pub count: u64,
}

#[derive(Debug, Serialize)]
pub struct AgentRollup {
    pub name: String,
    pub turns: u64,
    pub input_tokens: u64,
    pub output_tokens: u64,
    pub cache_read_tokens: u64,
    pub cache_creation_tokens: u64,
    /// Labelled estimate — see [`Prices`].
    pub est_cost_usd: f64,
}

#[derive(Debug, Serialize)]
pub struct HiveStats {
    pub window: &'static str,
    pub from: i64,
    pub now: i64,
    /// Number of agents that had at least one turn in the window.
    pub active_agents: u64,
    pub total_turns: u64,
    pub total_input_tokens: u64,
    pub total_output_tokens: u64,
    pub total_cache_read_tokens: u64,
    pub total_cache_creation_tokens: u64,
    /// Labelled estimate — see [`Prices`].
    pub est_cost_usd: f64,
    /// Per-agent rollup, busiest (most turns) first.
    pub agents: Vec<AgentRollup>,
    /// Turns per model across the whole swarm, busiest first.
    pub model_mix: Vec<KeyCount>,
    /// Most-run normalised bash-command heads ("favorite tools") across
    /// the whole swarm, busiest first, capped to 10. Empty until the
    /// hive-bash-mcp capture has recorded `bash_commands` rows on at
    /// least one active agent.
    pub bash_mix: Vec<KeyCount>,
}

#[derive(Default)]
struct AgentAgg {
    turns: u64,
    input: u64,
    output: u64,
    cache_read: u64,
    cache_creation: u64,
    cost: f64,
    models: HashMap<String, u64>,
    /// Normalised bash-command head → invocation count, from the agent's
    /// `bash_commands` table (written by hive-bash-mcp). Empty when that
    /// capture hasn't run for this agent (table absent).
    bash: HashMap<String, u64>,
}

fn now_secs() -> i64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .map_or(0, |d| i64::try_from(d.as_secs()).unwrap_or(i64::MAX))
}

#[allow(
    clippy::cast_sign_loss,
    clippy::cast_possible_truncation,
    reason = "v.max(0) clamps to a non-negative value first, so the i64->u64 cast is exact: no sign loss, and no truncation since u64 covers all non-negative i64"
)]
fn u64_from_i64(v: i64) -> u64 {
    v.max(0) as u64
}

/// Aggregate one agent's turn-stats over `[from, now]`. Errors bubble
/// up so the caller can skip a single bad/locked db without failing
/// the whole endpoint.
fn read_agent(path: &Path, from: i64, prices: &PriceTable) -> rusqlite::Result<AgentAgg> {
    let conn = Connection::open_with_flags(path, OpenFlags::SQLITE_OPEN_READ_ONLY)?;
    // turn_stats is rollback-journal (not WAL): a read landing while the
    // harness is mid-INSERT would get SQLITE_BUSY and drop that active
    // agent from the rollup. Wait out the brief write instead.
    conn.busy_timeout(Duration::from_millis(500))?;
    let mut stmt = conn.prepare(
        "SELECT model, input_tokens, output_tokens,
                cache_read_input_tokens, cache_creation_input_tokens
         FROM turn_stats
         WHERE started_at >= ?1",
    )?;
    let mut agg = AgentAgg::default();
    let rows = stmt.query_map([from], |row| {
        Ok((
            row.get::<_, String>(0)?,
            u64_from_i64(row.get::<_, i64>(1)?),
            u64_from_i64(row.get::<_, i64>(2)?),
            u64_from_i64(row.get::<_, i64>(3)?),
            u64_from_i64(row.get::<_, i64>(4)?),
        ))
    })?;
    for r in rows {
        let (model, input, output, cache_read, cache_creation) = r?;
        agg.turns += 1;
        agg.input = agg.input.saturating_add(input);
        agg.output = agg.output.saturating_add(output);
        agg.cache_read = agg.cache_read.saturating_add(cache_read);
        agg.cache_creation = agg.cache_creation.saturating_add(cache_creation);
        let p = resolve_prices(&model, prices);
        #[allow(
            clippy::cast_precision_loss,
            reason = "token counts stay well under f64's 2^53 exact-integer range, so this cost computation loses no precision"
        )]
        {
            agg.cost += (input as f64 * p.input
                + output as f64 * p.output
                + cache_read as f64 * p.cache_read
                + cache_creation as f64 * p.cache_write)
                / 1_000_000.0;
        }
        *agg.models.entry(model).or_insert(0) += 1;
    }
    agg.bash = read_bash_heads(&conn, from);
    Ok(agg)
}

/// Tally normalised bash-command heads from the agent's `bash_commands`
/// table (`ts INTEGER, head TEXT`, written by hive-bash-mcp) over
/// `[from, now]`. Best-effort + isolated from `read_agent`'s error path:
/// a missing table (capture hasn't run for this agent) or any read error
/// yields an empty map rather than propagating, so the favorite-tools
/// rollup simply omits that agent and never fails the whole endpoint.
/// `ts` is unix seconds, matching the `from` cutoff.
fn read_bash_heads(conn: &Connection, from: i64) -> HashMap<String, u64> {
    let mut out: HashMap<String, u64> = HashMap::new();
    let Ok(mut stmt) =
        conn.prepare("SELECT head, COUNT(*) FROM bash_commands WHERE ts >= ?1 GROUP BY head")
    else {
        return out;
    };
    let Ok(rows) = stmt.query_map([from], |row| {
        Ok((
            row.get::<_, String>(0)?,
            u64_from_i64(row.get::<_, i64>(1)?),
        ))
    }) else {
        return out;
    };
    for (head, count) in rows.flatten() {
        *out.entry(head).or_insert(0) += count;
    }
    out
}

/// Build the swarm-wide rollup. Best-effort: a missing or unreadable
/// per-agent db is skipped (logged), never fatal.
#[must_use]
pub fn hive_snapshot(window: Window, prices: &PriceTable) -> HiveStats {
    let now = now_secs();
    let from = now - window.span_secs();

    let mut agents: Vec<AgentRollup> = Vec::new();
    let mut model_mix: HashMap<String, u64> = HashMap::new();
    let mut bash_mix: HashMap<String, u64> = HashMap::new();
    let mut total_turns = 0u64;
    let mut total_input = 0u64;
    let mut total_output = 0u64;
    let mut total_cache_read = 0u64;
    let mut total_cache_creation = 0u64;
    let mut total_cost = 0.0f64;
    let mut active_agents = 0u64;

    for name in Coordinator::kept_state_names() {
        let path = Coordinator::agent_harness_dir(&name).join("hyperhive-turn-stats.sqlite");
        if !path.exists() {
            continue;
        }
        let agg = match read_agent(&path, from, prices) {
            Ok(a) => a,
            Err(e) => {
                tracing::warn!(agent = %name, error = ?e, "hive-stats: read failed; skipping");
                continue;
            }
        };
        if agg.turns == 0 {
            continue;
        }
        active_agents += 1;
        total_turns += agg.turns;
        total_input = total_input.saturating_add(agg.input);
        total_output = total_output.saturating_add(agg.output);
        total_cache_read = total_cache_read.saturating_add(agg.cache_read);
        total_cache_creation = total_cache_creation.saturating_add(agg.cache_creation);
        total_cost += agg.cost;
        for (m, c) in &agg.models {
            *model_mix.entry(m.clone()).or_insert(0) += c;
        }
        for (h, c) in &agg.bash {
            *bash_mix.entry(h.clone()).or_insert(0) += c;
        }
        agents.push(AgentRollup {
            name,
            turns: agg.turns,
            input_tokens: agg.input,
            output_tokens: agg.output,
            cache_read_tokens: agg.cache_read,
            cache_creation_tokens: agg.cache_creation,
            est_cost_usd: agg.cost,
        });
    }

    // Busiest agents first.
    agents.sort_by(|a, b| b.turns.cmp(&a.turns).then_with(|| a.name.cmp(&b.name)));

    let mut model_mix: Vec<KeyCount> = model_mix
        .into_iter()
        .map(|(key, count)| KeyCount { key, count })
        .collect();
    model_mix.sort_by(|a, b| b.count.cmp(&a.count).then_with(|| a.key.cmp(&b.key)));

    // Busiest commands first, capped to a top-10 "favorite tools" list.
    let mut bash_mix: Vec<KeyCount> = bash_mix
        .into_iter()
        .map(|(key, count)| KeyCount { key, count })
        .collect();
    bash_mix.sort_by(|a, b| b.count.cmp(&a.count).then_with(|| a.key.cmp(&b.key)));
    bash_mix.truncate(10);

    HiveStats {
        window: window.label(),
        from,
        now,
        active_agents,
        total_turns,
        total_input_tokens: total_input,
        total_output_tokens: total_output,
        total_cache_read_tokens: total_cache_read,
        total_cache_creation_tokens: total_cache_creation,
        est_cost_usd: total_cost,
        agents,
        model_mix,
        bash_mix,
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    /// `read_bash_heads` tallies per-head counts over the window and is
    /// isolated from a missing table (returns empty, never errors).
    #[test]
    fn bash_heads_tally_and_window() {
        let conn = Connection::open_in_memory().unwrap();
        conn.execute_batch("CREATE TABLE bash_commands (ts INTEGER NOT NULL, head TEXT NOT NULL);")
            .unwrap();
        let now = now_secs();
        for (ts, head) in [
            (now - 100, "cargo"),
            (now - 200, "cargo"),
            (now - 300, "git"),
            (now - 10_000, "rg"), // outside a 1h-ish cutoff below
        ] {
            conn.execute(
                "INSERT INTO bash_commands (ts, head) VALUES (?1, ?2)",
                rusqlite::params![ts, head],
            )
            .unwrap();
        }
        let heads = read_bash_heads(&conn, now - 3600);
        assert_eq!(heads.get("cargo").copied(), Some(2));
        assert_eq!(heads.get("git").copied(), Some(1));
        assert_eq!(heads.get("rg").copied(), None); // window cutoff excludes it
    }

    /// A missing `bash_commands` table degrades to an empty tally rather
    /// than erroring — the pre-capture window on a fresh agent.
    #[test]
    fn bash_heads_missing_table_is_empty() {
        let conn = Connection::open_in_memory().unwrap();
        assert!(read_bash_heads(&conn, 0).is_empty());
    }
}