clawmates/crates/cm-brain/src/lib.rs

//! `cm-brain` — ClawMates' facade over the canonical **`.brain`** (a `claw-brain`
//! / ClawhDF5 *brain-pack*): one HDF5 file holding an agent's full definition
//! (system prompt · personality · skills · tools · runtime · provenance) **and**
//! its memory, so claws can be built, deployed, loaded, upskilled, reloaded —
//! and pulled/pushed to ClawBrainHub + synced with ClawSync — as a single file.
//!
//! This is a thin, ClawMates-shaped wrapper over `claw_brain::BrainHandle` (the
//! brain-pack KV) + its keyword index for recall. Sections live under the
//! conventional brain-pack keys (`identity/system_prompt`, `skills/<name>`, …);
//! conversational memory lives under `memory/<ts>` and is keyword-searchable.

use std::path::{Path, PathBuf};
use std::time::{SystemTime, UNIX_EPOCH};

use claw_brain::{index_entry, keyword_search, BrainHandle};
pub use claw_brain::RevisionInfo;

pub mod hub;

const K_SYSTEM_PROMPT: &str = "identity/system_prompt";
const K_PERSONA: &str = "identity/persona";
const K_SOUL: &str = "identity/soul_md";
const K_AGENT_MD: &str = "identity/agent_md";
const K_RUNTIME: &str = "runtime/clawmates";
const K_PROVENANCE: &str = "provenance/clawmates";
const P_SKILL: &str = "skills/"; // skills/<name>
const P_TOOL: &str = "tools/"; // tools/<name>
const P_MEMORY: &str = "memory/"; // memory/<ts_nanos>
const K_SKILLS_MD: &str = "skills/skills_md"; // brain-pack narrative skills doc

#[derive(Debug, thiserror::Error)]
pub enum BrainError {
    #[error("brain backend: {0}")]
    Backend(String),
}

pub(crate) fn be<E: std::fmt::Display>(e: E) -> BrainError {
    BrainError::Backend(e.to_string())
}

fn now_nanos() -> u128 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .map(|d| d.as_nanos())
        .unwrap_or(0)
}

/// A claw's brain: a single `.brain` (ClawhDF5 brain-pack) file.
pub struct ClawBrain {
    h: BrainHandle,
    path: PathBuf,
}

impl ClawBrain {
    /// Open an existing `.brain`, or create a fresh one. `agent_id` is currently
    /// advisory (the brain is identified by its file path).
    pub fn open_or_create(path: impl AsRef<Path>, _agent_id: &str) -> Result<Self, BrainError> {
        let p = path.as_ref().to_path_buf();
        let h = if p.exists() {
            BrainHandle::open(&p).map_err(be)?
        } else {
            if let Some(dir) = p.parent() {
                std::fs::create_dir_all(dir).map_err(be)?;
            }
            BrainHandle::create(&p).map_err(be)?
        };
        Ok(Self { h, path: p })
    }

    /// Path to the underlying `.brain` file (for registry push / ClawSync).
    pub fn path(&self) -> &Path {
        &self.path
    }

    // ── ClawSync: local revision history + rollback (`.onion` sidecar) ─────────

    /// Commit the current brain state as a new revision in the `.brain.onion`
    /// sidecar. The base `.brain` stays current; history lives in the sidecar.
    /// First commit creates the sidecar; later commits store only changed pages.
    pub fn commit(&self, annotation: Option<&str>) -> Result<u64, BrainError> {
        self.h.flush().map_err(be)?;
        claw_brain::commit_versioned(&self.h, annotation).map_err(be)
    }

    /// List the brain's revision history (empty if never committed).
    pub fn revisions(&self) -> Result<Vec<claw_brain::RevisionInfo>, BrainError> {
        claw_brain::list_brain_revisions(&self.path).map_err(be)
    }

    /// Materialize a prior revision back onto the `.brain` file (the sidecar
    /// history is preserved, so a rollback is itself reversible).
    pub fn rollback(&self, revision: u64) -> Result<(), BrainError> {
        claw_brain::rollback_to_revision(&self.path, revision).map_err(be)
    }

    // ── raw key helpers ───────────────────────────────────────────────────────
    fn put(&self, key: &str, content: &str) -> Result<(), BrainError> {
        self.h.write(key, content.as_bytes().to_vec()).map_err(be)?;
        self.h.flush().map_err(be)
    }
    fn get(&self, key: &str) -> Option<String> {
        self.h.read(key).ok().and_then(|b| String::from_utf8(b).ok())
    }
    fn del(&self, key: &str) -> Result<bool, BrainError> {
        if self.h.read(key).is_err() {
            return Ok(false);
        }
        self.h.remove(key).map_err(be)?;
        self.h.flush().map_err(be)?;
        Ok(true)
    }
    /// `(suffix, value)` for every active key under `prefix` (excluding nested).
    fn list(&self, prefix: &str) -> Vec<(String, String)> {
        let mut out = Vec::new();
        for k in self.h.keys() {
            if let Some(name) = k.strip_prefix(prefix) {
                if name.is_empty() || name.contains('/') {
                    continue;
                }
                if let Some(v) = self.get(&k) {
                    out.push((name.to_string(), v));
                }
            }
        }
        out
    }

    // ── identity ──────────────────────────────────────────────────────────────
    pub fn set_system_prompt(&mut self, s: &str) -> Result<(), BrainError> { self.put(K_SYSTEM_PROMPT, s) }
    /// System prompt, falling back to the brain-pack `soul_md` (pulled brains
    /// often carry the identity there with an empty `system_prompt`).
    pub fn system_prompt(&self) -> Option<String> {
        match self.get(K_SYSTEM_PROMPT) {
            Some(s) if !s.trim().is_empty() => Some(s),
            _ => self.get(K_SOUL).filter(|s| !s.trim().is_empty()),
        }
    }
    pub fn set_personality(&mut self, s: &str) -> Result<(), BrainError> { self.put(K_PERSONA, s) }
    pub fn personality(&self) -> Option<String> { self.get(K_PERSONA).filter(|s| !s.trim().is_empty()) }
    /// AGENTS.md — "how I operate" (workflow/rules), prepended to the prompt.
    pub fn set_agent_md(&mut self, s: &str) -> Result<(), BrainError> { self.put(K_AGENT_MD, s) }
    pub fn agent_md(&self) -> Option<String> { self.get(K_AGENT_MD).filter(|s| !s.trim().is_empty()) }

    /// Assemble the agent's full system prompt from the canonical identity files
    /// (mirrors ZeroClaw's personality render): `system_prompt`‖`soul_md`, then
    /// `agent_md` (AGENTS.md), then `persona`.
    pub fn assembled_identity(&self) -> String {
        let mut out = String::new();
        if let Some(sp) = self.system_prompt() {
            out.push_str(&sp);
        }
        if let Some(a) = self.agent_md() {
            if !out.is_empty() { out.push_str("\n\n"); }
            out.push_str("## How I operate\n");
            out.push_str(&a);
        }
        if let Some(p) = self.personality() {
            if !out.is_empty() { out.push_str("\n\n"); }
            out.push_str("## Personality\n");
            out.push_str(&p);
        }
        out
    }

    // ── skills ────────────────────────────────────────────────────────────────
    pub fn set_skill(&mut self, name: &str, body: &str) -> Result<(), BrainError> { self.put(&format!("{P_SKILL}{name}"), body) }
    /// Set the brain-pack narrative skills doc (`skills/skills_md`); `skills()`
    /// parses its `## <name>` sections.
    pub fn set_skills_md(&mut self, md: &str) -> Result<(), BrainError> { self.put(K_SKILLS_MD, md) }
    pub fn remove_skill(&mut self, name: &str) -> Result<bool, BrainError> { self.del(&format!("{P_SKILL}{name}")) }
    /// All skills as `(name, body)` — both per-skill keys and, if present, the
    /// brain-pack narrative `skills/skills_md` parsed into `## <name>` sections.
    pub fn skills(&self) -> Vec<(String, String)> {
        let mut out: Vec<(String, String)> = self
            .list(P_SKILL)
            .into_iter()
            .filter(|(n, _)| format!("{P_SKILL}{n}") != K_SKILLS_MD)
            .collect();
        if let Some(md) = self.get(K_SKILLS_MD) {
            for (n, b) in parse_skills_md(&md) {
                if !out.iter().any(|(en, _)| *en == n) {
                    out.push((n, b));
                }
            }
        }
        out
    }

    // ── tools / doors (value = state, e.g. "gated" | "blocked") ────────────────
    pub fn set_tool(&mut self, name: &str, state: &str) -> Result<(), BrainError> { self.put(&format!("{P_TOOL}{name}"), state) }
    pub fn remove_tool(&mut self, name: &str) -> Result<bool, BrainError> { self.del(&format!("{P_TOOL}{name}")) }
    pub fn tools(&self) -> Vec<(String, String)> { self.list(P_TOOL) }

    // ── runtime + provenance (opaque JSON blobs) ───────────────────────────────
    pub fn set_runtime(&mut self, json: &str) -> Result<(), BrainError> { self.put(K_RUNTIME, json) }
    pub fn runtime(&self) -> Option<String> { self.get(K_RUNTIME) }
    pub fn set_provenance(&mut self, json: &str) -> Result<(), BrainError> { self.put(K_PROVENANCE, json) }
    pub fn provenance(&self) -> Option<String> { self.get(K_PROVENANCE) }

    // ── conversational memory ──────────────────────────────────────────────────
    /// Append a turn to the brain's memory (keyword-indexed, recallable later).
    pub fn remember(&mut self, role: &str, text: &str, _session_id: &str) -> Result<(), BrainError> {
        let key = format!("{P_MEMORY}{:020}", now_nanos());
        let chunk = format!("{role}: {text}");
        self.h.write(&key, chunk.into_bytes()).map_err(be)?;
        index_entry(&self.h, &key, text).map_err(be)?;
        self.h.flush().map_err(be)
    }

    /// Recall up to `k` past memory chunks most relevant to `query` (BM25 over
    /// the keyword index; only `memory/` entries). Best-effort.
    pub fn recall(&self, query: &str, k: usize) -> Vec<String> {
        let hits = match keyword_search(&self.h, query, k.saturating_mul(2).max(k)) {
            Ok(h) => h,
            Err(_) => return Vec::new(),
        };
        hits.into_iter()
            .filter(|r| r.key.starts_with(P_MEMORY))
            .filter_map(|r| self.get(&r.key))
            .take(k)
            .collect()
    }

    /// Most recent memory chunks, newest first, as `(timestamp_secs, text)`.
    pub fn recent_memory(&self, k: usize) -> Vec<(f64, String)> {
        let mut keys: Vec<(u128, String)> = self
            .h
            .keys()
            .into_iter()
            .filter_map(|key| key.strip_prefix(P_MEMORY).and_then(|s| s.parse::<u128>().ok()).map(|n| (n, key)))
            .collect();
        keys.sort_by(|a, b| b.0.cmp(&a.0));
        keys.into_iter()
            .take(k)
            .filter_map(|(n, key)| self.get(&key).map(|t| (n as f64 / 1e9, t)))
            .collect()
    }

    pub fn memory_count(&self) -> usize {
        self.h.keys().iter().filter(|k| k.starts_with(P_MEMORY)).count()
    }

    /// Render identity + skills as Markdown (for ZeroClaw workspace hydration).
    pub fn export_markdown(&self) -> String {
        let mut s = String::new();
        if let Some(sp) = self.system_prompt() {
            s.push_str("# System Prompt\n\n");
            s.push_str(&sp);
            s.push_str("\n\n");
        }
        if let Some(p) = self.personality() {
            s.push_str("# Personality\n\n");
            s.push_str(&p);
            s.push_str("\n\n");
        }
        let skills = self.skills();
        if !skills.is_empty() {
            s.push_str("# Skills\n\n");
            for (name, body) in skills {
                s.push_str(&format!("## {name}\n\n{body}\n\n"));
            }
        }
        s
    }
}

/// Split a `skills_md` doc into `(name, body)` by its `## <name>` headings.
fn parse_skills_md(md: &str) -> Vec<(String, String)> {
    let mut out = Vec::new();
    let mut name: Option<String> = None;
    let mut body = String::new();
    for line in md.lines() {
        if let Some(h) = line.strip_prefix("## ") {
            if let Some(n) = name.take() {
                out.push((n, body.trim().to_string()));
                body.clear();
            }
            name = Some(h.trim().to_string());
        } else if name.is_some() {
            body.push_str(line);
            body.push('\n');
        }
    }
    if let Some(n) = name.take() {
        out.push((n, body.trim().to_string()));
    }
    out
}

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

    fn temp_path(tag: &str) -> PathBuf {
        std::env::temp_dir().join(format!("cm-brain-{}-{}.brain", std::process::id(), tag))
    }

    #[test]
    fn brain_roundtrip_definition_and_memory() {
        let p = temp_path("rt");
        let _ = std::fs::remove_file(&p);
        {
            let mut b = ClawBrain::open_or_create(&p, "agent-x").expect("create");
            b.set_system_prompt("You are Atlas, a meticulous planner.").unwrap();
            b.set_personality("calm, precise, terse").unwrap();
            b.set_skill("python", "Write idiomatic, tested Python.").unwrap();
            b.set_tool("browser", "gated").unwrap();
            b.set_runtime(r#"{"model":"claude","risk":"toolfree"}"#).unwrap();
            b.remember("user", "My favorite color is teal.", "s1").unwrap();
            b.remember("user", "I live in Boston.", "s1").unwrap();
        }
        let b = ClawBrain::open_or_create(&p, "agent-x").expect("open");
        assert_eq!(b.system_prompt().as_deref(), Some("You are Atlas, a meticulous planner."));
        assert_eq!(b.personality().as_deref(), Some("calm, precise, terse"));
        assert_eq!(b.skills(), vec![("python".to_string(), "Write idiomatic, tested Python.".to_string())]);
        assert_eq!(b.tools(), vec![("browser".to_string(), "gated".to_string())]);
        assert!(b.runtime().unwrap().contains("toolfree"));
        assert_eq!(b.memory_count(), 2);
        let hits = b.recall("what is my favorite color", 3);
        assert!(hits.iter().any(|h| h.contains("teal")), "recall should surface teal; got {hits:?}");
        let _ = std::fs::remove_file(&p);
    }

    #[test]
    fn upsert_and_remove_skill() {
        let p = temp_path("skill");
        let _ = std::fs::remove_file(&p);
        let mut b = ClawBrain::open_or_create(&p, "a").unwrap();
        b.set_skill("py", "v1").unwrap();
        b.set_skill("py", "v2").unwrap();
        assert_eq!(b.skills(), vec![("py".to_string(), "v2".to_string())]);
        assert!(b.remove_skill("py").unwrap());
        assert!(b.skills().is_empty());
        let _ = std::fs::remove_file(&p);
    }

    #[test]
    fn parses_skills_md() {
        let md = "# Skills\n\n## web_search\n\nSearch the web.\n\n## read_file\n\nRead a file.\n";
        let s = parse_skills_md(md);
        assert_eq!(s.len(), 2);
        assert_eq!(s[0].0, "web_search");
        assert!(s[0].1.contains("Search the web"));
    }
}