Skip to main content

rustbrain_core/
brain.rs

1//! Primary library façade: [`Brain`].
2//!
3//! Most applications should open a [`Brain`], call [`Brain::sync`] after note or
4//! code edits, then use [`Brain::query_ranked`] / [`Brain::context_for_prompt`].
5
6use crate::context::{assemble_context, ContextOptions};
7use crate::error::{BrainError, Result};
8use crate::exporter::{BrainExporter, BrainImporter};
9use crate::indexer::WorkspaceIndexer;
10use crate::query::{QueryOptions, RankedHit};
11use crate::storage::Database;
12use crate::types::{ContextBundle, Node, SyncStats};
13use std::path::{Path, PathBuf};
14
15/// A project-scoped rustbrain instance rooted at `workspace/.brain`.
16///
17/// # Lifecycle
18///
19/// ```text
20/// Brain::create / open / open_or_create
21///         │
22///         ▼
23///      sync()  ──► index Markdown + AST + resolve links + bake graph.mmap
24///         │
25///         ├─► query / query_ranked
26///         ├─► context_for_prompt / context_for_prompt_with
27///         ├─► export / import
28///         └─► watch  (feature = "watch")
29/// ```
30///
31/// # Threading
32///
33/// `Brain` is not `Sync`. Open separate instances (or serialize access) for
34/// concurrent use. SQLite is opened with WAL and a busy timeout for multi-process
35/// friendliness on the same machine.
36pub struct Brain {
37    workspace: PathBuf,
38    brain_dir: PathBuf,
39    db: Database,
40}
41
42impl Brain {
43    /// Create a new brain directory and empty database under `workspace/.brain`.
44    ///
45    /// Creates the workspace directory if it does not exist. Writes a minimal
46    /// `workspace.json` marker. Does **not** index notes — call [`Self::sync`].
47    ///
48    /// # Errors
49    ///
50    /// Returns I/O or SQLite errors if the directory or database cannot be created.
51    pub fn create(workspace: impl AsRef<Path>) -> Result<Self> {
52        let workspace = canonicalize_or_owned(workspace.as_ref())?;
53        let brain_dir = workspace.join(".brain");
54        std::fs::create_dir_all(&brain_dir)?;
55        let db_path = brain_dir.join("db.sqlite");
56        let db = Database::open(&db_path)?;
57        let marker = brain_dir.join("workspace.json");
58        if !marker.exists() {
59            let meta = serde_json::json!({
60                "version": 1,
61                "workspace": workspace.to_string_lossy(),
62            });
63            std::fs::write(&marker, serde_json::to_string_pretty(&meta)?)?;
64        }
65        Ok(Self {
66            workspace,
67            brain_dir,
68            db,
69        })
70    }
71
72    /// Open an existing brain under `workspace/.brain`.
73    ///
74    /// If `workspace/.brain/db.sqlite` is missing, walks **parent directories**
75    /// (like git) until a brain is found or the filesystem root is reached.
76    ///
77    /// # Errors
78    ///
79    /// Returns [`BrainError::BrainNotFound`] if no brain is found.
80    pub fn open(workspace: impl AsRef<Path>) -> Result<Self> {
81        let start = canonicalize_or_owned(workspace.as_ref())?;
82        if let Some((ws, brain_dir)) = find_brain_dir(&start) {
83            let db = Database::open(brain_dir.join("db.sqlite"))?;
84            return Ok(Self {
85                workspace: ws,
86                brain_dir,
87                db,
88            });
89        }
90        Err(BrainError::BrainNotFound {
91            path: start.join(".brain"),
92        })
93    }
94
95    /// Open a brain **only** at `workspace` (no parent walk).
96    ///
97    /// Prefer [`Self::open`] for CLI tooling.
98    pub fn open_exact(workspace: impl AsRef<Path>) -> Result<Self> {
99        let workspace = canonicalize_or_owned(workspace.as_ref())?;
100        let brain_dir = workspace.join(".brain");
101        let db_path = brain_dir.join("db.sqlite");
102        if !db_path.exists() {
103            return Err(BrainError::BrainNotFound { path: brain_dir });
104        }
105        let db = Database::open(&db_path)?;
106        Ok(Self {
107            workspace,
108            brain_dir,
109            db,
110        })
111    }
112
113    /// Open an existing brain, or [`Self::create`] if none is present at `workspace`
114    /// (does not walk parents for create — create is always local).
115    pub fn open_or_create(workspace: impl AsRef<Path>) -> Result<Self> {
116        let workspace = workspace.as_ref();
117        let db_path = workspace.join(".brain").join("db.sqlite");
118        if db_path.exists() {
119            Self::open_exact(workspace)
120        } else if let Ok(b) = Self::open(workspace) {
121            // Found a parent brain when CWD is a subdir.
122            Ok(b)
123        } else {
124            Self::create(workspace)
125        }
126    }
127
128    /// Absolute path to the workspace root (parent of `.brain`).
129    pub fn workspace(&self) -> &Path {
130        &self.workspace
131    }
132
133    /// Absolute path to the `.brain` directory.
134    pub fn brain_dir(&self) -> &Path {
135        &self.brain_dir
136    }
137
138    /// Borrow the underlying [`Database`] for advanced queries or tooling.
139    pub fn database(&self) -> &Database {
140        &self.db
141    }
142
143    /// Index Markdown (and optional AST / Canvas) and recompile the CSR mmap cache.
144    ///
145    /// Unchanged files (matching `content_hash`) are skipped. WikiLinks and
146    /// `symbol:…` refs that cannot be resolved are stored as pending links and
147    /// retried at the end of the run (and on later syncs).
148    pub fn sync(&mut self) -> Result<SyncStats> {
149        let db_path = self.brain_dir.join("db.sqlite");
150        let db = Database::open(&db_path)?;
151        let indexer = WorkspaceIndexer::new(db, self.workspace.clone());
152        let stats = indexer.index_workspace()?;
153        // Refresh our handle so we see committed WAL state.
154        self.db = Database::open(&db_path)?;
155        Ok(stats)
156    }
157
158    /// Ranked search returning only nodes (score order preserved, scores discarded).
159    ///
160    /// Prefer [`Self::query_ranked`] when scores or match reasons matter.
161    pub fn query(&self, q: &str) -> Result<Vec<Node>> {
162        let hits = self.query_ranked(q, &QueryOptions::default())?;
163        Ok(hits.into_iter().map(|h| h.node).collect())
164    }
165
166    /// Ranked search with scores and human-readable match reasons.
167    ///
168    /// Combines FTS5 BM25 with title, id, tag, and alias boosts. See [`QueryOptions`].
169    pub fn query_ranked(&self, q: &str, opts: &QueryOptions) -> Result<Vec<RankedHit>> {
170        self.db.search_ranked(q, opts)
171    }
172
173    /// Build a graph-aware prompt context with a simple token budget.
174    ///
175    /// Equivalent to [`Self::context_for_prompt_with`] using default hop depth and
176    /// packing limits. Token accounting uses a ~4 characters/token heuristic.
177    pub fn context_for_prompt(&self, prompt: &str, max_tokens: usize) -> Result<ContextBundle> {
178        let opts = ContextOptions {
179            max_tokens,
180            ..ContextOptions::default()
181        };
182        self.context_for_prompt_with(prompt, &opts)
183    }
184
185    /// Context assembly with full control over hops, seed counts, and packing.
186    ///
187    /// When `graph.mmap` exists and `opts.hop_depth > 0`, CSR neighbors of seed
188    /// hits are scored and packed alongside FTS seeds.
189    pub fn context_for_prompt_with(
190        &self,
191        prompt: &str,
192        opts: &ContextOptions,
193    ) -> Result<ContextBundle> {
194        assemble_context(&self.db, &self.brain_dir, prompt, opts)
195    }
196
197    /// Export nodes/edges to a portable `.brainbundle` JSON file.
198    ///
199    /// When `decouple_ast` is true, symbol nodes and file/symbol path fields are
200    /// stripped so the bundle can move across repositories cleanly.
201    pub fn export(&self, out: impl AsRef<Path>, decouple_ast: bool) -> Result<()> {
202        BrainExporter::export_bundle(&self.db, out, decouple_ast)
203    }
204
205    /// Export optionally limited to one SubBrain scope (plus hubs).
206    pub fn export_scope(
207        &self,
208        out: impl AsRef<Path>,
209        decouple_ast: bool,
210        scope: Option<&str>,
211    ) -> Result<()> {
212        BrainExporter::export_bundle_filtered(&self.db, out, decouple_ast, scope)
213    }
214
215    /// Import a `.brainbundle`, upserting nodes and edges.
216    ///
217    /// Recompiles `graph.mmap` when the `mmap` feature is enabled.
218    ///
219    /// # Returns
220    ///
221    /// Number of nodes upserted from the bundle.
222    pub fn import(&mut self, input: impl AsRef<Path>) -> Result<usize> {
223        let n = BrainImporter::import_bundle(&self.db, input)?;
224        #[cfg(feature = "mmap")]
225        {
226            let indexer = WorkspaceIndexer::new(
227                Database::open(self.brain_dir.join("db.sqlite"))?,
228                self.workspace.clone(),
229            );
230            let _ = indexer.compile_mmap(&self.brain_dir.join("graph.mmap"));
231            self.db = Database::open(self.brain_dir.join("db.sqlite"))?;
232        }
233        Ok(n)
234    }
235
236    /// Block the current thread and watch the workspace for changes.
237    ///
238    /// Requires the `watch` Cargo feature. Debounces filesystem events for
239    /// `debounce_ms` milliseconds, then runs a full sync (content-hash skips
240    /// unchanged files) and remaps the CSR cache.
241    pub fn watch(&self, debounce_ms: u64) -> Result<()> {
242        crate::watch::watch_workspace(
243            &self.workspace,
244            crate::watch::WatchConfig {
245                debounce: std::time::Duration::from_millis(debounce_ms),
246                verbose: true,
247            },
248        )
249    }
250
251    /// Deterministic docs/ignore bootstrap for mature repositories.
252    ///
253    /// See [`crate::bootstrap::bootstrap_workspace`]. Does not replace a full
254    /// [`Self::sync`] — call sync afterward to index new files.
255    pub fn bootstrap(
256        workspace: impl AsRef<Path>,
257        opts: crate::bootstrap::BootstrapOptions,
258    ) -> Result<crate::bootstrap::BootstrapReport> {
259        crate::bootstrap::bootstrap_workspace(workspace.as_ref(), opts)
260    }
261
262    /// Health check for this brain (pending links, ratios, schema).
263    pub fn doctor(&self) -> Result<crate::doctor::DoctorReport> {
264        crate::doctor::run_doctor(&self.workspace)
265    }
266
267    /// Create a Markdown note on disk (`docs/…`) without opening a second DB.
268    ///
269    /// Call [`Self::sync`] afterward so FTS/graph pick it up.
270    pub fn note_new(&self, opts: &crate::note::NoteNewOptions) -> Result<crate::note::NoteCreated> {
271        crate::note::create_note(&self.workspace, opts)
272    }
273
274    /// Notes with no explicit graph edges (soft `auto_*` edges do not count).
275    pub fn list_orphans(&self) -> Result<Vec<crate::autolink::OrphanNote>> {
276        crate::autolink::list_orphan_notes(&self.db)
277    }
278
279    /// Create soft auto-links (filename stem + shared tags). Optional path focuses one note.
280    ///
281    /// Recompiles `graph.mmap` when the `mmap` feature is enabled.
282    pub fn auto_link(
283        &mut self,
284        target: Option<&std::path::Path>,
285    ) -> Result<crate::autolink::AutoLinkReport> {
286        let report = crate::autolink::run_auto_link(&self.db, target)?;
287        #[cfg(feature = "mmap")]
288        {
289            let indexer = WorkspaceIndexer::new(
290                Database::open(self.brain_dir.join("db.sqlite"))?,
291                self.workspace.clone(),
292            );
293            let _ = indexer.compile_mmap(&self.brain_dir.join("graph.mmap"));
294            self.db = Database::open(self.brain_dir.join("db.sqlite"))?;
295        }
296        Ok(report)
297    }
298
299    /// Plan/apply pending WikiLink normalizations and optional AC discovery.
300    ///
301    /// See [`crate::apply_links::apply_links`]. When `opts.write` is true and
302    /// `opts.sync_after`, call [`Self::sync`] after a successful write.
303    ///
304    /// Sets `cache_dir` to `.brain/` when the caller left it unset so discover
305    /// can reuse the LinkLexicon fingerprint cache.
306    pub fn apply_links(
307        &self,
308        opts: &crate::apply_links::ApplyOptions,
309    ) -> Result<crate::apply_links::ApplyReport> {
310        let mut opts = opts.clone();
311        if opts.cache_dir.is_none() {
312            opts.cache_dir = Some(self.brain_dir.clone());
313        }
314        crate::apply_links::apply_links(&self.workspace, &self.db, &opts)
315    }
316
317    /// k-hop neighborhood around a node (path, id, title, or `symbol:…`).
318    ///
319    /// Uses SQLite edges (relation types preserved). Prefer for inspection;
320    /// use [`Self::context_for_prompt`] when packing content for agents.
321    pub fn graph_neighborhood(
322        &self,
323        target: &str,
324        opts: &crate::graph::GraphOptions,
325    ) -> Result<crate::graph::GraphNeighborhood> {
326        crate::graph::neighborhood(&self.db, target, opts)
327    }
328
329    /// Aggregate node/edge counts, relation breakdown, and high-degree hubs.
330    pub fn graph_stats(&self) -> Result<crate::graph::GraphStats> {
331        crate::graph::graph_stats(&self.db)
332    }
333}
334
335fn canonicalize_or_owned(path: &Path) -> Result<PathBuf> {
336    if path.exists() {
337        Ok(fs_canonicalize(path)?)
338    } else {
339        std::fs::create_dir_all(path)?;
340        Ok(fs_canonicalize(path)?)
341    }
342}
343
344fn fs_canonicalize(path: &Path) -> Result<PathBuf> {
345    std::fs::canonicalize(path).map_err(BrainError::from)
346}
347
348/// Walk `start` and its parents looking for `.brain/db.sqlite`.
349///
350/// Returns `(workspace_root, brain_dir)`.
351pub fn find_brain_dir(start: &Path) -> Option<(PathBuf, PathBuf)> {
352    let mut cur = start.to_path_buf();
353    loop {
354        let brain_dir = cur.join(".brain");
355        if brain_dir.join("db.sqlite").is_file() {
356            return Some((cur, brain_dir));
357        }
358        if !cur.pop() {
359            break;
360        }
361    }
362    None
363}
364
365#[cfg(test)]
366mod tests {
367    use super::*;
368    use crate::types::ContextRole;
369    use tempfile::tempdir;
370
371    #[test]
372    fn open_walks_parent_for_brain() {
373        let dir = tempdir().unwrap();
374        let root = dir.path();
375        let sub = root.join("src").join("nested");
376        std::fs::create_dir_all(&sub).unwrap();
377        Brain::create(root).unwrap();
378        let opened = Brain::open(&sub).unwrap();
379        assert_eq!(
380            opened.workspace().canonicalize().unwrap(),
381            root.canonicalize().unwrap()
382        );
383    }
384
385    #[test]
386    fn create_sync_query_context_export() {
387        let dir = tempdir().unwrap();
388        let docs = dir.path().join("docs");
389        std::fs::create_dir_all(&docs).unwrap();
390        std::fs::write(
391            docs.join("raft.md"),
392            "---\ntags: [raft]\nnode_type: concept\n---\n# Raft\nSee [[logcompaction]].\n",
393        )
394        .unwrap();
395        std::fs::write(
396            docs.join("logcompaction.md"),
397            "---\ntags: [log]\nnode_type: concept\n---\n# Log Compaction\nSee [[raft]].\n",
398        )
399        .unwrap();
400
401        let mut brain = Brain::create(dir.path()).unwrap();
402        let stats = brain.sync().unwrap();
403        assert_eq!(stats.markdown_files, 2);
404
405        let hits = brain.query("raft").unwrap();
406        assert!(!hits.is_empty());
407
408        let ranked = brain
409            .query_ranked("raft", &QueryOptions::default())
410            .unwrap();
411        assert!(ranked[0].score > 0.0);
412
413        let ctx = brain.context_for_prompt("raft", 512).unwrap();
414        assert!(
415            !ctx.nodes.is_empty(),
416            "expected FTS hits for 'raft', got none"
417        );
418        assert!(ctx.tokens_used > 0);
419        assert!(ctx.nodes.iter().any(|n| n.role == ContextRole::Seed));
420        let xml = ctx.to_xml();
421        assert!(xml.contains("<rustbrain_context"));
422        assert!(xml.contains("tokens_used="));
423
424        let out = dir.path().join("out.brainbundle");
425        brain.export(&out, true).unwrap();
426        assert!(out.exists());
427
428        let _ = brain.sync().unwrap();
429        assert_eq!(brain.database().count_fts_rows().unwrap(), 2);
430    }
431
432    #[test]
433    fn note_anchors_to_symbol() {
434        let dir = tempdir().unwrap();
435        let docs = dir.path().join("docs");
436        let src = dir.path().join("src");
437        std::fs::create_dir_all(&docs).unwrap();
438        std::fs::create_dir_all(&src).unwrap();
439        std::fs::write(
440            dir.path().join("Cargo.toml"),
441            "[package]\nname = \"demo\"\nversion = \"0.1.0\"\nedition = \"2021\"\n",
442        )
443        .unwrap();
444        std::fs::write(
445            src.join("lib.rs"),
446            "/// Storage engine\npub struct StorageEngine;\nimpl StorageEngine { pub fn open() {} }\n",
447        )
448        .unwrap();
449        std::fs::write(
450            docs.join("design.md"),
451            "---\nnode_type: adr\n---\n# Design\nUses symbol:StorageEngine for persistence.\n",
452        )
453        .unwrap();
454
455        let mut brain = Brain::create(dir.path()).unwrap();
456        let stats = brain.sync().unwrap();
457        assert!(stats.symbol_anchors >= 1);
458        assert!(stats.markdown_files >= 1);
459
460        let _ = brain.sync().unwrap();
461        let edges = brain.database().get_all_edges().unwrap();
462        let anchors: Vec<_> = edges
463            .iter()
464            .filter(|e| e.relation_type == "anchors")
465            .collect();
466        assert!(
467            !anchors.is_empty(),
468            "expected anchors edge note→symbol, edges={edges:?}"
469        );
470    }
471}