Big Brain v0.4.0

hooks.rs

Rust · 259 lines · 10.7 KB
//! Claude Code hooks. All must be fast and must never fail loudly: any error
//! means "say nothing" so a broken brain can't break an agent's session.

use crate::librarian::Librarian;
use crate::store::Brain;
use serde_json::{json, Value};
use sha2::{Digest, Sha256};
use std::collections::HashMap;
use std::fs;
use std::io::Read;
use std::path::{Path, PathBuf};

/// Edits before the first nudge, and between later nudges.
const FIRST_NUDGE: usize = 5;
const NEXT_NUDGE: usize = 25;
/// Stale entries listed per message; the rest follow on later messages.
const SHOW_MAX: usize = 10;
/// Messages after which a still-stale entry the session was told about is mentioned again.
const REMIND_AFTER: usize = 10;

fn stdin_json() -> Value {
    let mut s = String::new();
    let _ = std::io::stdin().read_to_string(&mut s);
    serde_json::from_str(&s).unwrap_or(Value::Null)
}

fn session_id(input: &Value) -> Option<String> {
    input["session_id"]
        .as_str()
        .filter(|s| !s.is_empty() && s.chars().all(|c| c.is_ascii_alphanumeric() || c == '-'))
        .map(str::to_string)
}

fn state_file(brain: &Brain, session: &str, ext: &str) -> PathBuf {
    let dir = brain.home.join("hook-state");
    let _ = fs::create_dir_all(&dir);
    dir.join(format!("{session}.{ext}"))
}

fn cwd_of(input: &Value) -> PathBuf {
    input["cwd"]
        .as_str()
        .map(PathBuf::from)
        .unwrap_or_else(|| std::env::current_dir().unwrap_or_default())
}

struct StaleEntry {
    id: String,
    /// Changes when the referenced files change again, so a new change is re-announced.
    sig: String,
    line: String,
}

/// Entries owned by `project` whose refs changed, plus the project's entry count.
fn project_stale(brain: &Brain, project: &str) -> (usize, Vec<StaleEntry>) {
    let mut total = 0;
    let mut out = vec![];
    let graph = crate::links::Graph::build(brain);
    for (id, meta) in &graph.metas {
        if graph.is_retired(id) {
            continue; // retired entries don't need keeping up to date
        }
        let id = id.clone();
        let e = crate::store::Entry { id: id.clone(), meta: meta.clone(), body: String::new(), raw: String::new() };
        let owner = if e.meta.owner.is_empty() { &e.meta.project } else { &e.meta.owner };
        if owner != project {
            continue;
        }
        total += 1;
        let stale = brain.stale_refs(&e.meta);
        if stale.is_empty() {
            continue;
        }
        let mut h = Sha256::new();
        for r in &e.meta.refs {
            h.update(brain.ref_hash_now(r).unwrap_or_default());
        }
        // File names come from entries (untrusted text), so automatic context only gets a count.
        let files = stale.len();
        // Only the verdict word goes into agent context; the note itself is model output and is read
        // deliberately with `brain get`.
        let note = Librarian::current_review(brain, &e)
            .map(|r| {
                // reviews/*.json can be hand-written: only the three real verdicts reach agent context.
                let v = if ["accurate", "needs_update", "unsure"].contains(&r.verdict.as_str()) { r.verdict.as_str() } else { "unsure" };
                format!(" (librarian review: {v}; see `brain get {id}`)")
            })
            .unwrap_or_default();
        out.push(StaleEntry {
            id: id.clone(),
            sig: hex::encode(h.finalize())[..16].to_string(),
            line: format!("  - {id} ({files} referenced file{} changed){note}", if files == 1 { "" } else { "s" }),
        });
    }
    (total, out)
}

/// `id:sig` → the message number when the session was told (lines are `id:sig msg`).
fn read_told(path: &Path) -> HashMap<String, usize> {
    fs::read_to_string(path)
        .unwrap_or_default()
        .lines()
        .filter(|l| !l.is_empty())
        .map(|l| match l.split_once(' ') {
            Some((k, n)) => (k.to_string(), n.parse().unwrap_or(0)),
            None => (l.to_string(), 0),
        })
        .collect()
}

fn remember_told(path: &Path, items: &[StaleEntry], msg: usize) {
    let mut told = read_told(path);
    for s in items {
        told.insert(format!("{}:{}", s.id, s.sig), msg);
    }
    let mut v: Vec<String> = told.into_iter().map(|(k, n)| format!("{k} {n}")).collect();
    v.sort();
    let _ = fs::write(path, v.join("\n") + "\n");
}

/// Should a stale entry be mentioned on message `n`? Yes if never mentioned for this change,
/// or if it was mentioned REMIND_AFTER or more messages ago and is still stale.
pub(crate) fn due(told_at: Option<usize>, n: usize) -> bool {
    told_at.is_none_or(|at| n >= at + REMIND_AFTER)
}

/// Messages this session has sent, counted by the prompt hook.
fn message_count(brain: &Brain, session: &str) -> usize {
    fs::read_to_string(state_file(brain, session, "count")).ok().and_then(|s| s.trim().parse().ok()).unwrap_or(0)
}

/// SessionStart: tell the agent the brain exists, and what its project needs.
pub fn session_start() {
    let input = stdin_json();
    let Ok(brain) = Brain::open() else { return };
    let project = brain.project_for_dir(&cwd_of(&input));
    let mut lines = vec![format!(
        "Big Brain (`brain` CLI) is the shared memory for all of {}'s agents: how things were built, \
         workflows, fixes, gotchas, decisions, runbooks. Load the `big-brain` skill, `brain search` before \
         building or debugging anything non-trivial, and record what you build or learn.",
        brain.config.identity.user
    )];
    if let Some(p) = &project {
        let (total, stale) = project_stale(&brain, p);
        lines.push(format!("This session's project: `{p}` — {total} entries in the brain."));
        if total == 0 {
            lines.push(
                "It has no entries yet: backfill it (overview, systems, workflows, fixes) — see the skill.".into(),
            );
        }
        if !stale.is_empty() {
            lines.push(format!(
                "{} of its entries reference files that changed since they were written. Check each \
                 (`brain get <id>`), then `brain update <id>` or `brain ack <id>`:",
                stale.len()
            ));
            let shown: Vec<StaleEntry> = stale.into_iter().take(SHOW_MAX).collect();
            lines.extend(shown.iter().map(|s| s.line.clone()));
            if let Some(sid) = session_id(&input) {
                // Only what was shown counts as told; the rest comes up on later messages.
                remember_told(&state_file(&brain, &sid, "told"), &shown, message_count(&brain, &sid));
            }
        }
    }
    emit_context("SessionStart", &lines.join("\n"));
}

/// UserPromptSubmit: mention entries of this session's project that went stale
/// since the session was last told, or that are still stale REMIND_AFTER messages after it was. Silent otherwise.
pub fn prompt() {
    let input = stdin_json();
    let Some(sid) = session_id(&input) else { return };
    let Ok(brain) = Brain::open() else { return };
    let Some(project) = brain.project_for_dir(&cwd_of(&input)) else { return };

    let n = message_count(&brain, &sid) + 1;
    let _ = fs::write(state_file(&brain, &sid, "count"), n.to_string());
    let every = brain.config.hooks.prompt_every.max(1);
    if !n.is_multiple_of(every) {
        return;
    }

    let told_path = state_file(&brain, &sid, "told");
    let told = read_told(&told_path);
    let (_, stale) = project_stale(&brain, &project);
    // New staleness, or told REMIND_AFTER+ messages ago and still not handled.
    let new: Vec<StaleEntry> = stale
        .into_iter()
        .filter(|s| due(told.get(&format!("{}:{}", s.id, s.sig)).copied(), n))
        .collect();
    if new.is_empty() {
        return;
    }
    let total = new.len();
    let shown: Vec<StaleEntry> = new.into_iter().take(SHOW_MAX).collect();
    let more = if total > shown.len() {
        format!("\n({} more will be listed on later messages.)", total - shown.len())
    } else {
        String::new()
    };
    let user = &brain.config.identity.user;
    emit_context("UserPromptSubmit", &format!(
        "Big Brain: {total} of `{project}`'s entries are stale (their referenced files changed). After handling \
         {user}'s message, read the changes and `brain update <id>` if the entry is now wrong, or `brain ack <id>` \
         if it's still accurate:\n{}{more}",
        shown.iter().map(|s| s.line.as_str()).collect::<Vec<_>>().join("\n")
    ));
    // Only what was shown counts as told.
    remember_told(&told_path, &shown, n);
}

/// File edits in a session transcript. Matches each harness's edit tools by name, so it works on
/// Claude Code (Edit/Write/MultiEdit/NotebookEdit), Codex (apply_patch) and Grok (search_replace/write)
/// transcripts without parsing their different formats.
pub(crate) fn count_edits(transcript: &Path) -> usize {
    let Ok(text) = fs::read_to_string(transcript) else { return 0 };
    static EDIT_RE: std::sync::LazyLock<regex::Regex> = std::sync::LazyLock::new(|| {
        regex::Regex::new(r#""name":\s*"(Edit|Write|MultiEdit|NotebookEdit|apply_patch|search_replace|write)""#).unwrap()
    });
    EDIT_RE.find_iter(&text).count()
}

/// Print context for the model in the shape Claude Code and Codex both read.
fn emit_context(event: &str, text: &str) {
    println!("{}", json!({"hookSpecificOutput": {"hookEventName": event, "additionalContext": text}}));
}

/// Stop: once real work has happened, ask the agent (once, then rarely) to record it.
pub fn stop() {
    let input = stdin_json();
    if input["stop_hook_active"].as_bool() == Some(true) {
        return;
    }
    let (Some(session), Some(transcript)) = (session_id(&input), input["transcript_path"].as_str()) else {
        return;
    };
    let Ok(brain) = Brain::open() else { return };
    let edits = count_edits(Path::new(transcript));
    let state = brain.home.join("hook-state").join(&session);
    let last: Option<usize> = fs::read_to_string(&state).ok().and_then(|s| s.trim().parse().ok());
    let due = match last {
        None => edits >= FIRST_NUDGE,
        Some(at) => edits >= at + NEXT_NUDGE,
    };
    if !due {
        return;
    }
    let _ = fs::create_dir_all(brain.home.join("hook-state"));
    let _ = fs::write(&state, edits.to_string());
    let reason = format!(
        "Big Brain check (automatic, from {}'s setup): you've changed files this session. \
        If you built, solved, or learned something another agent could reuse — a system and its key files, \
        a workflow, a fix, a gotcha, a decision — record it now: `brain search` first, then `brain update <id>` \
        if an entry exists, else `brain record` (see the big-brain skill). If nothing is worth recording, \
        say so in one line and finish. Don't repeat work you already recorded.",
        brain.config.identity.user
    );
    println!("{}", json!({"decision": "block", "reason": reason}));
}