1use 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
15pub struct Brain {
37 workspace: PathBuf,
38 brain_dir: PathBuf,
39 db: Database,
40}
41
42impl Brain {
43 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 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 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 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 Ok(b)
123 } else {
124 Self::create(workspace)
125 }
126 }
127
128 pub fn workspace(&self) -> &Path {
130 &self.workspace
131 }
132
133 pub fn brain_dir(&self) -> &Path {
135 &self.brain_dir
136 }
137
138 pub fn database(&self) -> &Database {
140 &self.db
141 }
142
143 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 self.db = Database::open(&db_path)?;
155 Ok(stats)
156 }
157
158 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 pub fn query_ranked(&self, q: &str, opts: &QueryOptions) -> Result<Vec<RankedHit>> {
170 self.db.search_ranked(q, opts)
171 }
172
173 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 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 pub fn export(&self, out: impl AsRef<Path>, decouple_ast: bool) -> Result<()> {
202 BrainExporter::export_bundle(&self.db, out, decouple_ast)
203 }
204
205 pub fn import(&mut self, input: impl AsRef<Path>) -> Result<usize> {
213 let n = BrainImporter::import_bundle(&self.db, input)?;
214 #[cfg(feature = "mmap")]
215 {
216 let indexer = WorkspaceIndexer::new(
217 Database::open(self.brain_dir.join("db.sqlite"))?,
218 self.workspace.clone(),
219 );
220 let _ = indexer.compile_mmap(&self.brain_dir.join("graph.mmap"));
221 self.db = Database::open(self.brain_dir.join("db.sqlite"))?;
222 }
223 Ok(n)
224 }
225
226 pub fn watch(&self, debounce_ms: u64) -> Result<()> {
232 crate::watch::watch_workspace(
233 &self.workspace,
234 crate::watch::WatchConfig {
235 debounce: std::time::Duration::from_millis(debounce_ms),
236 verbose: true,
237 },
238 )
239 }
240
241 pub fn bootstrap(
246 workspace: impl AsRef<Path>,
247 opts: crate::bootstrap::BootstrapOptions,
248 ) -> Result<crate::bootstrap::BootstrapReport> {
249 crate::bootstrap::bootstrap_workspace(workspace.as_ref(), opts)
250 }
251
252 pub fn doctor(&self) -> Result<crate::doctor::DoctorReport> {
254 crate::doctor::run_doctor(&self.workspace)
255 }
256
257 pub fn note_new(&self, opts: &crate::note::NoteNewOptions) -> Result<crate::note::NoteCreated> {
261 crate::note::create_note(&self.workspace, opts)
262 }
263
264 pub fn list_orphans(&self) -> Result<Vec<crate::autolink::OrphanNote>> {
266 crate::autolink::list_orphan_notes(&self.db)
267 }
268
269 pub fn auto_link(
273 &mut self,
274 target: Option<&std::path::Path>,
275 ) -> Result<crate::autolink::AutoLinkReport> {
276 let report = crate::autolink::run_auto_link(&self.db, target)?;
277 #[cfg(feature = "mmap")]
278 {
279 let indexer = WorkspaceIndexer::new(
280 Database::open(self.brain_dir.join("db.sqlite"))?,
281 self.workspace.clone(),
282 );
283 let _ = indexer.compile_mmap(&self.brain_dir.join("graph.mmap"));
284 self.db = Database::open(self.brain_dir.join("db.sqlite"))?;
285 }
286 Ok(report)
287 }
288
289 pub fn graph_neighborhood(
294 &self,
295 target: &str,
296 opts: &crate::graph::GraphOptions,
297 ) -> Result<crate::graph::GraphNeighborhood> {
298 crate::graph::neighborhood(&self.db, target, opts)
299 }
300
301 pub fn graph_stats(&self) -> Result<crate::graph::GraphStats> {
303 crate::graph::graph_stats(&self.db)
304 }
305}
306
307fn canonicalize_or_owned(path: &Path) -> Result<PathBuf> {
308 if path.exists() {
309 Ok(fs_canonicalize(path)?)
310 } else {
311 std::fs::create_dir_all(path)?;
312 Ok(fs_canonicalize(path)?)
313 }
314}
315
316fn fs_canonicalize(path: &Path) -> Result<PathBuf> {
317 std::fs::canonicalize(path).map_err(BrainError::from)
318}
319
320pub fn find_brain_dir(start: &Path) -> Option<(PathBuf, PathBuf)> {
324 let mut cur = start.to_path_buf();
325 loop {
326 let brain_dir = cur.join(".brain");
327 if brain_dir.join("db.sqlite").is_file() {
328 return Some((cur, brain_dir));
329 }
330 if !cur.pop() {
331 break;
332 }
333 }
334 None
335}
336
337#[cfg(test)]
338mod tests {
339 use super::*;
340 use crate::types::ContextRole;
341 use tempfile::tempdir;
342
343 #[test]
344 fn open_walks_parent_for_brain() {
345 let dir = tempdir().unwrap();
346 let root = dir.path();
347 let sub = root.join("src").join("nested");
348 std::fs::create_dir_all(&sub).unwrap();
349 Brain::create(root).unwrap();
350 let opened = Brain::open(&sub).unwrap();
351 assert_eq!(
352 opened.workspace().canonicalize().unwrap(),
353 root.canonicalize().unwrap()
354 );
355 }
356
357 #[test]
358 fn create_sync_query_context_export() {
359 let dir = tempdir().unwrap();
360 let docs = dir.path().join("docs");
361 std::fs::create_dir_all(&docs).unwrap();
362 std::fs::write(
363 docs.join("raft.md"),
364 "---\ntags: [raft]\nnode_type: concept\n---\n# Raft\nSee [[logcompaction]].\n",
365 )
366 .unwrap();
367 std::fs::write(
368 docs.join("logcompaction.md"),
369 "---\ntags: [log]\nnode_type: concept\n---\n# Log Compaction\nSee [[raft]].\n",
370 )
371 .unwrap();
372
373 let mut brain = Brain::create(dir.path()).unwrap();
374 let stats = brain.sync().unwrap();
375 assert_eq!(stats.markdown_files, 2);
376
377 let hits = brain.query("raft").unwrap();
378 assert!(!hits.is_empty());
379
380 let ranked = brain
381 .query_ranked("raft", &QueryOptions::default())
382 .unwrap();
383 assert!(ranked[0].score > 0.0);
384
385 let ctx = brain.context_for_prompt("raft", 512).unwrap();
386 assert!(
387 !ctx.nodes.is_empty(),
388 "expected FTS hits for 'raft', got none"
389 );
390 assert!(ctx.tokens_used > 0);
391 assert!(ctx.nodes.iter().any(|n| n.role == ContextRole::Seed));
392 let xml = ctx.to_xml();
393 assert!(xml.contains("<rustbrain_context"));
394 assert!(xml.contains("tokens_used="));
395
396 let out = dir.path().join("out.brainbundle");
397 brain.export(&out, true).unwrap();
398 assert!(out.exists());
399
400 let _ = brain.sync().unwrap();
401 assert_eq!(brain.database().count_fts_rows().unwrap(), 2);
402 }
403
404 #[test]
405 fn note_anchors_to_symbol() {
406 let dir = tempdir().unwrap();
407 let docs = dir.path().join("docs");
408 let src = dir.path().join("src");
409 std::fs::create_dir_all(&docs).unwrap();
410 std::fs::create_dir_all(&src).unwrap();
411 std::fs::write(
412 dir.path().join("Cargo.toml"),
413 "[package]\nname = \"demo\"\nversion = \"0.1.0\"\nedition = \"2021\"\n",
414 )
415 .unwrap();
416 std::fs::write(
417 src.join("lib.rs"),
418 "/// Storage engine\npub struct StorageEngine;\nimpl StorageEngine { pub fn open() {} }\n",
419 )
420 .unwrap();
421 std::fs::write(
422 docs.join("design.md"),
423 "---\nnode_type: adr\n---\n# Design\nUses symbol:StorageEngine for persistence.\n",
424 )
425 .unwrap();
426
427 let mut brain = Brain::create(dir.path()).unwrap();
428 let stats = brain.sync().unwrap();
429 assert!(stats.symbol_anchors >= 1);
430 assert!(stats.markdown_files >= 1);
431
432 let _ = brain.sync().unwrap();
433 let edges = brain.database().get_all_edges().unwrap();
434 let anchors: Vec<_> = edges
435 .iter()
436 .filter(|e| e.relation_type == "anchors")
437 .collect();
438 assert!(
439 !anchors.is_empty(),
440 "expected anchors edge note→symbol, edges={edges:?}"
441 );
442 }
443}