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    /// # Errors
75    ///
76    /// Returns [`BrainError::BrainNotFound`] if `db.sqlite` is missing.
77    pub fn open(workspace: impl AsRef<Path>) -> Result<Self> {
78        let workspace = canonicalize_or_owned(workspace.as_ref())?;
79        let brain_dir = workspace.join(".brain");
80        let db_path = brain_dir.join("db.sqlite");
81        if !db_path.exists() {
82            return Err(BrainError::BrainNotFound { path: brain_dir });
83        }
84        let db = Database::open(&db_path)?;
85        Ok(Self {
86            workspace,
87            brain_dir,
88            db,
89        })
90    }
91
92    /// Open an existing brain, or [`Self::create`] if none is present.
93    pub fn open_or_create(workspace: impl AsRef<Path>) -> Result<Self> {
94        let workspace = workspace.as_ref();
95        let db_path = workspace.join(".brain").join("db.sqlite");
96        if db_path.exists() {
97            Self::open(workspace)
98        } else {
99            Self::create(workspace)
100        }
101    }
102
103    /// Absolute path to the workspace root (parent of `.brain`).
104    pub fn workspace(&self) -> &Path {
105        &self.workspace
106    }
107
108    /// Absolute path to the `.brain` directory.
109    pub fn brain_dir(&self) -> &Path {
110        &self.brain_dir
111    }
112
113    /// Borrow the underlying [`Database`] for advanced queries or tooling.
114    pub fn database(&self) -> &Database {
115        &self.db
116    }
117
118    /// Index Markdown (and optional AST / Canvas) and recompile the CSR mmap cache.
119    ///
120    /// Unchanged files (matching `content_hash`) are skipped. WikiLinks and
121    /// `symbol:…` refs that cannot be resolved are stored as pending links and
122    /// retried at the end of the run (and on later syncs).
123    pub fn sync(&mut self) -> Result<SyncStats> {
124        let db_path = self.brain_dir.join("db.sqlite");
125        let db = Database::open(&db_path)?;
126        let indexer = WorkspaceIndexer::new(db, self.workspace.clone());
127        let stats = indexer.index_workspace()?;
128        // Refresh our handle so we see committed WAL state.
129        self.db = Database::open(&db_path)?;
130        Ok(stats)
131    }
132
133    /// Ranked search returning only nodes (score order preserved, scores discarded).
134    ///
135    /// Prefer [`Self::query_ranked`] when scores or match reasons matter.
136    pub fn query(&self, q: &str) -> Result<Vec<Node>> {
137        let hits = self.query_ranked(q, &QueryOptions::default())?;
138        Ok(hits.into_iter().map(|h| h.node).collect())
139    }
140
141    /// Ranked search with scores and human-readable match reasons.
142    ///
143    /// Combines FTS5 BM25 with title, id, tag, and alias boosts. See [`QueryOptions`].
144    pub fn query_ranked(&self, q: &str, opts: &QueryOptions) -> Result<Vec<RankedHit>> {
145        self.db.search_ranked(q, opts)
146    }
147
148    /// Build a graph-aware prompt context with a simple token budget.
149    ///
150    /// Equivalent to [`Self::context_for_prompt_with`] using default hop depth and
151    /// packing limits. Token accounting uses a ~4 characters/token heuristic.
152    pub fn context_for_prompt(&self, prompt: &str, max_tokens: usize) -> Result<ContextBundle> {
153        let opts = ContextOptions {
154            max_tokens,
155            ..ContextOptions::default()
156        };
157        self.context_for_prompt_with(prompt, &opts)
158    }
159
160    /// Context assembly with full control over hops, seed counts, and packing.
161    ///
162    /// When `graph.mmap` exists and `opts.hop_depth > 0`, CSR neighbors of seed
163    /// hits are scored and packed alongside FTS seeds.
164    pub fn context_for_prompt_with(
165        &self,
166        prompt: &str,
167        opts: &ContextOptions,
168    ) -> Result<ContextBundle> {
169        assemble_context(&self.db, &self.brain_dir, prompt, opts)
170    }
171
172    /// Export nodes/edges to a portable `.brainbundle` JSON file.
173    ///
174    /// When `decouple_ast` is true, symbol nodes and file/symbol path fields are
175    /// stripped so the bundle can move across repositories cleanly.
176    pub fn export(&self, out: impl AsRef<Path>, decouple_ast: bool) -> Result<()> {
177        BrainExporter::export_bundle(&self.db, out, decouple_ast)
178    }
179
180    /// Import a `.brainbundle`, upserting nodes and edges.
181    ///
182    /// Recompiles `graph.mmap` when the `mmap` feature is enabled.
183    ///
184    /// # Returns
185    ///
186    /// Number of nodes upserted from the bundle.
187    pub fn import(&mut self, input: impl AsRef<Path>) -> Result<usize> {
188        let n = BrainImporter::import_bundle(&self.db, input)?;
189        #[cfg(feature = "mmap")]
190        {
191            let indexer = WorkspaceIndexer::new(
192                Database::open(self.brain_dir.join("db.sqlite"))?,
193                self.workspace.clone(),
194            );
195            let _ = indexer.compile_mmap(&self.brain_dir.join("graph.mmap"));
196            self.db = Database::open(self.brain_dir.join("db.sqlite"))?;
197        }
198        Ok(n)
199    }
200
201    /// Block the current thread and watch the workspace for changes.
202    ///
203    /// Requires the `watch` Cargo feature. Debounces filesystem events for
204    /// `debounce_ms` milliseconds, then runs a full sync (content-hash skips
205    /// unchanged files) and remaps the CSR cache.
206    pub fn watch(&self, debounce_ms: u64) -> Result<()> {
207        crate::watch::watch_workspace(
208            &self.workspace,
209            crate::watch::WatchConfig {
210                debounce: std::time::Duration::from_millis(debounce_ms),
211                verbose: true,
212            },
213        )
214    }
215
216    /// Deterministic docs/ignore bootstrap for mature repositories.
217    ///
218    /// See [`crate::bootstrap::bootstrap_workspace`]. Does not replace a full
219    /// [`Self::sync`] — call sync afterward to index new files.
220    pub fn bootstrap(
221        workspace: impl AsRef<Path>,
222        opts: crate::bootstrap::BootstrapOptions,
223    ) -> Result<crate::bootstrap::BootstrapReport> {
224        crate::bootstrap::bootstrap_workspace(workspace.as_ref(), opts)
225    }
226
227    /// Health check for this brain (pending links, ratios, schema).
228    pub fn doctor(&self) -> Result<crate::doctor::DoctorReport> {
229        crate::doctor::run_doctor(&self.workspace)
230    }
231
232    /// Create a Markdown note on disk (`docs/…`) without opening a second DB.
233    ///
234    /// Call [`Self::sync`] afterward so FTS/graph pick it up.
235    pub fn note_new(&self, opts: &crate::note::NoteNewOptions) -> Result<crate::note::NoteCreated> {
236        crate::note::create_note(&self.workspace, opts)
237    }
238}
239
240fn canonicalize_or_owned(path: &Path) -> Result<PathBuf> {
241    if path.exists() {
242        Ok(fs_canonicalize(path)?)
243    } else {
244        std::fs::create_dir_all(path)?;
245        Ok(fs_canonicalize(path)?)
246    }
247}
248
249fn fs_canonicalize(path: &Path) -> Result<PathBuf> {
250    std::fs::canonicalize(path).map_err(BrainError::from)
251}
252
253#[cfg(test)]
254mod tests {
255    use super::*;
256    use crate::types::ContextRole;
257    use tempfile::tempdir;
258
259    #[test]
260    fn create_sync_query_context_export() {
261        let dir = tempdir().unwrap();
262        let docs = dir.path().join("docs");
263        std::fs::create_dir_all(&docs).unwrap();
264        std::fs::write(
265            docs.join("raft.md"),
266            "---\ntags: [raft]\nnode_type: concept\n---\n# Raft\nSee [[logcompaction]].\n",
267        )
268        .unwrap();
269        std::fs::write(
270            docs.join("logcompaction.md"),
271            "---\ntags: [log]\nnode_type: concept\n---\n# Log Compaction\nSee [[raft]].\n",
272        )
273        .unwrap();
274
275        let mut brain = Brain::create(dir.path()).unwrap();
276        let stats = brain.sync().unwrap();
277        assert_eq!(stats.markdown_files, 2);
278
279        let hits = brain.query("raft").unwrap();
280        assert!(!hits.is_empty());
281
282        let ranked = brain
283            .query_ranked("raft", &QueryOptions::default())
284            .unwrap();
285        assert!(ranked[0].score > 0.0);
286
287        let ctx = brain.context_for_prompt("raft", 512).unwrap();
288        assert!(
289            !ctx.nodes.is_empty(),
290            "expected FTS hits for 'raft', got none"
291        );
292        assert!(ctx.tokens_used > 0);
293        assert!(ctx.nodes.iter().any(|n| n.role == ContextRole::Seed));
294        let xml = ctx.to_xml();
295        assert!(xml.contains("<rustbrain_context"));
296        assert!(xml.contains("tokens_used="));
297
298        let out = dir.path().join("out.brainbundle");
299        brain.export(&out, true).unwrap();
300        assert!(out.exists());
301
302        let _ = brain.sync().unwrap();
303        assert_eq!(brain.database().count_fts_rows().unwrap(), 2);
304    }
305
306    #[test]
307    fn note_anchors_to_symbol() {
308        let dir = tempdir().unwrap();
309        let docs = dir.path().join("docs");
310        let src = dir.path().join("src");
311        std::fs::create_dir_all(&docs).unwrap();
312        std::fs::create_dir_all(&src).unwrap();
313        std::fs::write(
314            dir.path().join("Cargo.toml"),
315            "[package]\nname = \"demo\"\nversion = \"0.1.0\"\nedition = \"2021\"\n",
316        )
317        .unwrap();
318        std::fs::write(
319            src.join("lib.rs"),
320            "/// Storage engine\npub struct StorageEngine;\nimpl StorageEngine { pub fn open() {} }\n",
321        )
322        .unwrap();
323        std::fs::write(
324            docs.join("design.md"),
325            "---\nnode_type: adr\n---\n# Design\nUses symbol:StorageEngine for persistence.\n",
326        )
327        .unwrap();
328
329        let mut brain = Brain::create(dir.path()).unwrap();
330        let stats = brain.sync().unwrap();
331        assert!(stats.symbol_anchors >= 1);
332        assert!(stats.markdown_files >= 1);
333
334        let _ = brain.sync().unwrap();
335        let edges = brain.database().get_all_edges().unwrap();
336        let anchors: Vec<_> = edges
337            .iter()
338            .filter(|e| e.relation_type == "anchors")
339            .collect();
340        assert!(
341            !anchors.is_empty(),
342            "expected anchors edge note→symbol, edges={edges:?}"
343        );
344    }
345}