hyperhive/hive-bash-mcp/src/stats.rs

//! Capture of normalised shell-command "heads" per bash task into the
//! per-agent turn-stats sqlite (`bash_commands` table), backing the
//! "favorite tools" doughnut on the /stats page.
//!
//! Best-effort: a missing db, a lock, or a parse miss is logged and
//! swallowed — recording a stat never affects task execution. The
//! read/aggregation side lives host-side in `hive-c0re` and reads the
//! same `bash_commands(ts INTEGER, head TEXT)` rows.
//!
//! Normalisation is deliberately rough (quoting-naive) — it just needs
//! to name the tool an operator would recognise. The dominant agent
//! pattern is `cd <repo> && <tool> …`, so we look past leading
//! `cd …`-style builtins and prefix-runners like `sudo`/`env` to find
//! the first real command, then take its basename.

use std::time::{SystemTime, UNIX_EPOCH};

use rusqlite::{Connection, params};

use crate::paths;

/// Builtins whose presence as a segment head means the *real* command
/// is in a later segment (e.g. the `cargo` in `cd repo && cargo build`),
/// so we skip the whole segment.
const SEGMENT_BUILTINS: &[&str] = &[
    "cd", "pushd", "popd", "export", "set", "unset", "source", ".", ":", "true", "false",
];

/// Prefix-runners that wrap the command following them in the SAME
/// segment (e.g. the `cargo` in `env FOO=bar cargo build`). We skip the
/// runner plus any following flags / env-assignments and unwrap to the
/// wrapped command.
const PREFIX_RUNNERS: &[&str] = &[
    "env", "exec", "sudo", "doas", "nice", "ionice", "stdbuf", "time", "command", "nohup",
];

/// Does `tok` look like a `VAR=value` shell assignment?
fn is_assignment(tok: &str) -> bool {
    match tok.split_once('=') {
        Some((name, _)) if !name.is_empty() => name.chars().enumerate().all(|(i, c)| {
            c == '_'
                || if i == 0 {
                    c.is_ascii_alphabetic()
                } else {
                    c.is_ascii_alphanumeric()
                }
        }),
        _ => false,
    }
}

/// Strip a path prefix: `/nix/store/…/bin/git` -> `git`.
fn basename(tok: &str) -> &str {
    tok.rsplit('/').next().unwrap_or(tok)
}

/// Resolve the head command of a single segment (no sequencing
/// operators inside). `None` for an empty or builtin-only segment.
fn segment_head(seg: &str) -> Option<String> {
    let mut tokens = seg.split_whitespace().peekable();
    // Skip leading env-assignments (`FOO=bar cmd`).
    while tokens.peek().is_some_and(|t| is_assignment(t)) {
        tokens.next();
    }
    let mut head = tokens.next()?;
    // Unwrap prefix-runners: skip the runner, its option flags, and any
    // further env-assignments, then take the wrapped command. Bounded
    // loop guards against a pathological `sudo sudo sudo …`.
    let mut guard = 0;
    while PREFIX_RUNNERS.contains(&basename(head)) && guard < 8 {
        guard += 1;
        while tokens.peek().is_some_and(|t| t.starts_with('-')) {
            tokens.next();
        }
        while tokens.peek().is_some_and(|t| is_assignment(t)) {
            tokens.next();
        }
        match tokens.next() {
            Some(next) => head = next,
            None => return None, // bare `sudo` with nothing after it
        }
    }
    let base = basename(head);
    if base.is_empty() || SEGMENT_BUILTINS.contains(&base) {
        return None;
    }
    Some(base.to_owned())
}

/// Extract a normalised tool name from a shell command line: the
/// basename of the first real command, looking past `cd repo &&`
/// prefixes, env-assignments, and prefix-runners like `sudo`/`env`.
/// `None` when no real command can be identified.
#[must_use]
pub fn command_head(cmd: &str) -> Option<String> {
    // Split into segments on shell sequencing operators. Naive: doesn't
    // honour quotes/escapes — acceptable for a best-effort tally. Order
    // matters: `||` before `|` so the former isn't shredded by the latter.
    let mut segments: Vec<&str> = vec![cmd];
    for sep in ["&&", "||", ";", "|", "\n"] {
        segments = segments.into_iter().flat_map(|s| s.split(sep)).collect();
    }
    segments.into_iter().find_map(segment_head)
}

/// Record one bash task's normalised command head. Best-effort: any
/// error (no db yet, locked, unparseable command) is logged at debug
/// and swallowed.
pub fn record_command(cmd: &str) {
    let Some(head) = command_head(cmd) else {
        return;
    };
    if let Err(e) = record_head(&head) {
        tracing::debug!(error = ?e, head = %head, "bash stats: record failed");
    }
}

fn record_head(head: &str) -> rusqlite::Result<()> {
    let conn = Connection::open(paths::turn_stats_db())?;
    // The harness writes turn rows to this same db (rollback-journal,
    // not WAL) and sibling bash tasks may write concurrently — wait out
    // a transient lock rather than dropping the stat on SQLITE_BUSY.
    conn.busy_timeout(std::time::Duration::from_secs(2))?;
    conn.execute(
        "CREATE TABLE IF NOT EXISTS bash_commands (ts INTEGER NOT NULL, head TEXT NOT NULL)",
        [],
    )?;
    let now = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default()
        .as_secs()
        .cast_signed();
    conn.execute(
        "INSERT INTO bash_commands (ts, head) VALUES (?1, ?2)",
        params![now, head],
    )?;
    Ok(())
}

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

    #[test]
    fn skips_leading_cd_to_the_real_tool() {
        assert_eq!(
            command_head("cd /repo && cargo build").as_deref(),
            Some("cargo")
        );
        assert_eq!(
            command_head("cd /a/b && export X=1; hive-forge view 1").as_deref(),
            Some("hive-forge")
        );
    }

    #[test]
    fn strips_store_path_to_basename() {
        assert_eq!(
            command_head("/nix/store/abc123/bin/git status").as_deref(),
            Some("git")
        );
        assert_eq!(
            command_head("VAR=1 ./script.sh").as_deref(),
            Some("script.sh")
        );
    }

    #[test]
    fn skips_leading_env_assignments() {
        assert_eq!(
            command_head("FOO=bar BAZ=qux make").as_deref(),
            Some("make")
        );
    }

    #[test]
    fn unwraps_prefix_runners() {
        assert_eq!(
            command_head("sudo systemctl restart nginx").as_deref(),
            Some("systemctl")
        );
        assert_eq!(
            command_head("env RUST_LOG=debug cargo test").as_deref(),
            Some("cargo")
        );
        assert_eq!(command_head("time nix build").as_deref(), Some("nix"));
    }

    #[test]
    fn first_command_of_a_pipeline() {
        assert_eq!(
            command_head("git log --oneline | head -5").as_deref(),
            Some("git")
        );
    }

    #[test]
    fn none_for_empty_or_builtin_only() {
        assert_eq!(command_head(""), None);
        assert_eq!(command_head("   "), None);
        assert_eq!(command_head("cd /only/dir"), None);
        assert_eq!(command_head("sudo"), None);
    }

    #[test]
    fn assignment_detector() {
        assert!(is_assignment("FOO=bar"));
        assert!(is_assignment("RUST_LOG=debug"));
        assert!(!is_assignment("--flag"));
        assert!(!is_assignment("=leading"));
        assert!(!is_assignment("git"));
    }
}