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 apply_links(
297 &self,
298 opts: &crate::apply_links::ApplyOptions,
299 ) -> Result<crate::apply_links::ApplyReport> {
300 let mut opts = opts.clone();
301 if opts.cache_dir.is_none() {
302 opts.cache_dir = Some(self.brain_dir.clone());
303 }
304 crate::apply_links::apply_links(&self.workspace, &self.db, &opts)
305 }
306
307 pub fn graph_neighborhood(
312 &self,
313 target: &str,
314 opts: &crate::graph::GraphOptions,
315 ) -> Result<crate::graph::GraphNeighborhood> {
316 crate::graph::neighborhood(&self.db, target, opts)
317 }
318
319 pub fn graph_stats(&self) -> Result<crate::graph::GraphStats> {
321 crate::graph::graph_stats(&self.db)
322 }
323}
324
325fn canonicalize_or_owned(path: &Path) -> Result<PathBuf> {
326 if path.exists() {
327 Ok(fs_canonicalize(path)?)
328 } else {
329 std::fs::create_dir_all(path)?;
330 Ok(fs_canonicalize(path)?)
331 }
332}
333
334fn fs_canonicalize(path: &Path) -> Result<PathBuf> {
335 std::fs::canonicalize(path).map_err(BrainError::from)
336}
337
338pub fn find_brain_dir(start: &Path) -> Option<(PathBuf, PathBuf)> {
342 let mut cur = start.to_path_buf();
343 loop {
344 let brain_dir = cur.join(".brain");
345 if brain_dir.join("db.sqlite").is_file() {
346 return Some((cur, brain_dir));
347 }
348 if !cur.pop() {
349 break;
350 }
351 }
352 None
353}
354
355#[cfg(test)]
356mod tests {
357 use super::*;
358 use crate::types::ContextRole;
359 use tempfile::tempdir;
360
361 #[test]
362 fn open_walks_parent_for_brain() {
363 let dir = tempdir().unwrap();
364 let root = dir.path();
365 let sub = root.join("src").join("nested");
366 std::fs::create_dir_all(&sub).unwrap();
367 Brain::create(root).unwrap();
368 let opened = Brain::open(&sub).unwrap();
369 assert_eq!(
370 opened.workspace().canonicalize().unwrap(),
371 root.canonicalize().unwrap()
372 );
373 }
374
375 #[test]
376 fn create_sync_query_context_export() {
377 let dir = tempdir().unwrap();
378 let docs = dir.path().join("docs");
379 std::fs::create_dir_all(&docs).unwrap();
380 std::fs::write(
381 docs.join("raft.md"),
382 "---\ntags: [raft]\nnode_type: concept\n---\n# Raft\nSee [[logcompaction]].\n",
383 )
384 .unwrap();
385 std::fs::write(
386 docs.join("logcompaction.md"),
387 "---\ntags: [log]\nnode_type: concept\n---\n# Log Compaction\nSee [[raft]].\n",
388 )
389 .unwrap();
390
391 let mut brain = Brain::create(dir.path()).unwrap();
392 let stats = brain.sync().unwrap();
393 assert_eq!(stats.markdown_files, 2);
394
395 let hits = brain.query("raft").unwrap();
396 assert!(!hits.is_empty());
397
398 let ranked = brain
399 .query_ranked("raft", &QueryOptions::default())
400 .unwrap();
401 assert!(ranked[0].score > 0.0);
402
403 let ctx = brain.context_for_prompt("raft", 512).unwrap();
404 assert!(
405 !ctx.nodes.is_empty(),
406 "expected FTS hits for 'raft', got none"
407 );
408 assert!(ctx.tokens_used > 0);
409 assert!(ctx.nodes.iter().any(|n| n.role == ContextRole::Seed));
410 let xml = ctx.to_xml();
411 assert!(xml.contains("<rustbrain_context"));
412 assert!(xml.contains("tokens_used="));
413
414 let out = dir.path().join("out.brainbundle");
415 brain.export(&out, true).unwrap();
416 assert!(out.exists());
417
418 let _ = brain.sync().unwrap();
419 assert_eq!(brain.database().count_fts_rows().unwrap(), 2);
420 }
421
422 #[test]
423 fn note_anchors_to_symbol() {
424 let dir = tempdir().unwrap();
425 let docs = dir.path().join("docs");
426 let src = dir.path().join("src");
427 std::fs::create_dir_all(&docs).unwrap();
428 std::fs::create_dir_all(&src).unwrap();
429 std::fs::write(
430 dir.path().join("Cargo.toml"),
431 "[package]\nname = \"demo\"\nversion = \"0.1.0\"\nedition = \"2021\"\n",
432 )
433 .unwrap();
434 std::fs::write(
435 src.join("lib.rs"),
436 "/// Storage engine\npub struct StorageEngine;\nimpl StorageEngine { pub fn open() {} }\n",
437 )
438 .unwrap();
439 std::fs::write(
440 docs.join("design.md"),
441 "---\nnode_type: adr\n---\n# Design\nUses symbol:StorageEngine for persistence.\n",
442 )
443 .unwrap();
444
445 let mut brain = Brain::create(dir.path()).unwrap();
446 let stats = brain.sync().unwrap();
447 assert!(stats.symbol_anchors >= 1);
448 assert!(stats.markdown_files >= 1);
449
450 let _ = brain.sync().unwrap();
451 let edges = brain.database().get_all_edges().unwrap();
452 let anchors: Vec<_> = edges
453 .iter()
454 .filter(|e| e.relation_type == "anchors")
455 .collect();
456 assert!(
457 !anchors.is_empty(),
458 "expected anchors edge note→symbol, edges={edges:?}"
459 );
460 }
461}