rustbrain_core/

brain.rs

//! Primary library façade: [`Brain`].
//!
//! Most applications should open a [`Brain`], call [`Brain::sync`] after note or
//! code edits, then use [`Brain::query_ranked`] / [`Brain::context_for_prompt`].

use crate::context::{assemble_context, ContextOptions};
use crate::error::{BrainError, Result};
use crate::exporter::{BrainExporter, BrainImporter};
use crate::indexer::WorkspaceIndexer;
use crate::query::{QueryOptions, RankedHit};
use crate::storage::Database;
use crate::types::{ContextBundle, Node, SyncStats};
use std::path::{Path, PathBuf};

/// A project-scoped rustbrain instance rooted at `workspace/.brain`.
///
/// # Lifecycle
///
/// ```text
/// Brain::create / open / open_or_create
///         │
///         ▼
///      sync()  ──► index Markdown + AST + resolve links + bake graph.mmap
///         │
///         ├─► query / query_ranked
///         ├─► context_for_prompt / context_for_prompt_with
///         ├─► export / import
///         └─► watch  (feature = "watch")
/// ```
///
/// # Threading
///
/// `Brain` is not `Sync`. Open separate instances (or serialize access) for
/// concurrent use. SQLite is opened with WAL and a busy timeout for multi-process
/// friendliness on the same machine.
pub struct Brain {
    workspace: PathBuf,
    brain_dir: PathBuf,
    db: Database,
}

impl Brain {
    /// Create a new brain directory and empty database under `workspace/.brain`.
    ///
    /// Creates the workspace directory if it does not exist. Writes a minimal
    /// `workspace.json` marker. Does **not** index notes — call [`Self::sync`].
    ///
    /// # Errors
    ///
    /// Returns I/O or SQLite errors if the directory or database cannot be created.
    pub fn create(workspace: impl AsRef<Path>) -> Result<Self> {
        let workspace = canonicalize_or_owned(workspace.as_ref())?;
        let brain_dir = workspace.join(".brain");
        std::fs::create_dir_all(&brain_dir)?;
        let db_path = brain_dir.join("db.sqlite");
        let db = Database::open(&db_path)?;
        let marker = brain_dir.join("workspace.json");
        if !marker.exists() {
            let meta = serde_json::json!({
                "version": 1,
                "workspace": workspace.to_string_lossy(),
            });
            std::fs::write(&marker, serde_json::to_string_pretty(&meta)?)?;
        }
        Ok(Self {
            workspace,
            brain_dir,
            db,
        })
    }

    /// Open an existing brain under `workspace/.brain`.
    ///
    /// # Errors
    ///
    /// Returns [`BrainError::BrainNotFound`] if `db.sqlite` is missing.
    pub fn open(workspace: impl AsRef<Path>) -> Result<Self> {
        let workspace = canonicalize_or_owned(workspace.as_ref())?;
        let brain_dir = workspace.join(".brain");
        let db_path = brain_dir.join("db.sqlite");
        if !db_path.exists() {
            return Err(BrainError::BrainNotFound { path: brain_dir });
        }
        let db = Database::open(&db_path)?;
        Ok(Self {
            workspace,
            brain_dir,
            db,
        })
    }

    /// Open an existing brain, or [`Self::create`] if none is present.
    pub fn open_or_create(workspace: impl AsRef<Path>) -> Result<Self> {
        let workspace = workspace.as_ref();
        let db_path = workspace.join(".brain").join("db.sqlite");
        if db_path.exists() {
            Self::open(workspace)
        } else {
            Self::create(workspace)
        }
    }

    /// Absolute path to the workspace root (parent of `.brain`).
    pub fn workspace(&self) -> &Path {
        &self.workspace
    }

    /// Absolute path to the `.brain` directory.
    pub fn brain_dir(&self) -> &Path {
        &self.brain_dir
    }

    /// Borrow the underlying [`Database`] for advanced queries or tooling.
    pub fn database(&self) -> &Database {
        &self.db
    }

    /// Index Markdown (and optional AST / Canvas) and recompile the CSR mmap cache.
    ///
    /// Unchanged files (matching `content_hash`) are skipped. WikiLinks and
    /// `symbol:…` refs that cannot be resolved are stored as pending links and
    /// retried at the end of the run (and on later syncs).
    pub fn sync(&mut self) -> Result<SyncStats> {
        let db_path = self.brain_dir.join("db.sqlite");
        let db = Database::open(&db_path)?;
        let indexer = WorkspaceIndexer::new(db, self.workspace.clone());
        let stats = indexer.index_workspace()?;
        // Refresh our handle so we see committed WAL state.
        self.db = Database::open(&db_path)?;
        Ok(stats)
    }

    /// Ranked search returning only nodes (score order preserved, scores discarded).
    ///
    /// Prefer [`Self::query_ranked`] when scores or match reasons matter.
    pub fn query(&self, q: &str) -> Result<Vec<Node>> {
        let hits = self.query_ranked(q, &QueryOptions::default())?;
        Ok(hits.into_iter().map(|h| h.node).collect())
    }

    /// Ranked search with scores and human-readable match reasons.
    ///
    /// Combines FTS5 BM25 with title, id, tag, and alias boosts. See [`QueryOptions`].
    pub fn query_ranked(&self, q: &str, opts: &QueryOptions) -> Result<Vec<RankedHit>> {
        self.db.search_ranked(q, opts)
    }

    /// Build a graph-aware prompt context with a simple token budget.
    ///
    /// Equivalent to [`Self::context_for_prompt_with`] using default hop depth and
    /// packing limits. Token accounting uses a ~4 characters/token heuristic.
    pub fn context_for_prompt(&self, prompt: &str, max_tokens: usize) -> Result<ContextBundle> {
        let opts = ContextOptions {
            max_tokens,
            ..ContextOptions::default()
        };
        self.context_for_prompt_with(prompt, &opts)
    }

    /// Context assembly with full control over hops, seed counts, and packing.
    ///
    /// When `graph.mmap` exists and `opts.hop_depth > 0`, CSR neighbors of seed
    /// hits are scored and packed alongside FTS seeds.
    pub fn context_for_prompt_with(
        &self,
        prompt: &str,
        opts: &ContextOptions,
    ) -> Result<ContextBundle> {
        assemble_context(&self.db, &self.brain_dir, prompt, opts)
    }

    /// Export nodes/edges to a portable `.brainbundle` JSON file.
    ///
    /// When `decouple_ast` is true, symbol nodes and file/symbol path fields are
    /// stripped so the bundle can move across repositories cleanly.
    pub fn export(&self, out: impl AsRef<Path>, decouple_ast: bool) -> Result<()> {
        BrainExporter::export_bundle(&self.db, out, decouple_ast)
    }

    /// Import a `.brainbundle`, upserting nodes and edges.
    ///
    /// Recompiles `graph.mmap` when the `mmap` feature is enabled.
    ///
    /// # Returns
    ///
    /// Number of nodes upserted from the bundle.
    pub fn import(&mut self, input: impl AsRef<Path>) -> Result<usize> {
        let n = BrainImporter::import_bundle(&self.db, input)?;
        #[cfg(feature = "mmap")]
        {
            let indexer = WorkspaceIndexer::new(
                Database::open(self.brain_dir.join("db.sqlite"))?,
                self.workspace.clone(),
            );
            let _ = indexer.compile_mmap(&self.brain_dir.join("graph.mmap"));
            self.db = Database::open(self.brain_dir.join("db.sqlite"))?;
        }
        Ok(n)
    }

    /// Block the current thread and watch the workspace for changes.
    ///
    /// Requires the `watch` Cargo feature. Debounces filesystem events for
    /// `debounce_ms` milliseconds, then runs a full sync (content-hash skips
    /// unchanged files) and remaps the CSR cache.
    pub fn watch(&self, debounce_ms: u64) -> Result<()> {
        crate::watch::watch_workspace(
            &self.workspace,
            crate::watch::WatchConfig {
                debounce: std::time::Duration::from_millis(debounce_ms),
                verbose: true,
            },
        )
    }
}

fn canonicalize_or_owned(path: &Path) -> Result<PathBuf> {
    if path.exists() {
        Ok(fs_canonicalize(path)?)
    } else {
        std::fs::create_dir_all(path)?;
        Ok(fs_canonicalize(path)?)
    }
}

fn fs_canonicalize(path: &Path) -> Result<PathBuf> {
    std::fs::canonicalize(path).map_err(BrainError::from)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::types::ContextRole;
    use tempfile::tempdir;

    #[test]
    fn create_sync_query_context_export() {
        let dir = tempdir().unwrap();
        let docs = dir.path().join("docs");
        std::fs::create_dir_all(&docs).unwrap();
        std::fs::write(
            docs.join("raft.md"),
            "---\ntags: [raft]\nnode_type: concept\n---\n# Raft\nSee [[logcompaction]].\n",
        )
        .unwrap();
        std::fs::write(
            docs.join("logcompaction.md"),
            "---\ntags: [log]\nnode_type: concept\n---\n# Log Compaction\nSee [[raft]].\n",
        )
        .unwrap();

        let mut brain = Brain::create(dir.path()).unwrap();
        let stats = brain.sync().unwrap();
        assert_eq!(stats.markdown_files, 2);

        let hits = brain.query("raft").unwrap();
        assert!(!hits.is_empty());

        let ranked = brain
            .query_ranked("raft", &QueryOptions::default())
            .unwrap();
        assert!(ranked[0].score > 0.0);

        let ctx = brain.context_for_prompt("raft", 512).unwrap();
        assert!(
            !ctx.nodes.is_empty(),
            "expected FTS hits for 'raft', got none"
        );
        assert!(ctx.tokens_used > 0);
        assert!(ctx.nodes.iter().any(|n| n.role == ContextRole::Seed));
        let xml = ctx.to_xml();
        assert!(xml.contains("<rustbrain_context"));
        assert!(xml.contains("tokens_used="));

        let out = dir.path().join("out.brainbundle");
        brain.export(&out, true).unwrap();
        assert!(out.exists());

        let _ = brain.sync().unwrap();
        assert_eq!(brain.database().count_fts_rows().unwrap(), 2);
    }

    #[test]
    fn note_anchors_to_symbol() {
        let dir = tempdir().unwrap();
        let docs = dir.path().join("docs");
        let src = dir.path().join("src");
        std::fs::create_dir_all(&docs).unwrap();
        std::fs::create_dir_all(&src).unwrap();
        std::fs::write(
            dir.path().join("Cargo.toml"),
            "[package]\nname = \"demo\"\nversion = \"0.1.0\"\nedition = \"2021\"\n",
        )
        .unwrap();
        std::fs::write(
            src.join("lib.rs"),
            "/// Storage engine\npub struct StorageEngine;\nimpl StorageEngine { pub fn open() {} }\n",
        )
        .unwrap();
        std::fs::write(
            docs.join("design.md"),
            "---\nnode_type: adr\n---\n# Design\nUses symbol:StorageEngine for persistence.\n",
        )
        .unwrap();

        let mut brain = Brain::create(dir.path()).unwrap();
        let stats = brain.sync().unwrap();
        assert!(stats.symbol_anchors >= 1);
        assert!(stats.markdown_files >= 1);

        let _ = brain.sync().unwrap();
        let edges = brain.database().get_all_edges().unwrap();
        let anchors: Vec<_> = edges
            .iter()
            .filter(|e| e.relation_type == "anchors")
            .collect();
        assert!(
            !anchors.is_empty(),
            "expected anchors edge note→symbol, edges={edges:?}"
        );
    }
}