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 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 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 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 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 pub fn doctor(&self) -> Result<crate::doctor::DoctorReport> {
264 crate::doctor::run_doctor(&self.workspace)
265 }
266
267 pub fn note_new(&self, opts: &crate::note::NoteNewOptions) -> Result<crate::note::NoteCreated> {
271 crate::note::create_note(&self.workspace, opts)
272 }
273
274 pub fn list_orphans(&self) -> Result<Vec<crate::autolink::OrphanNote>> {
276 crate::autolink::list_orphan_notes(&self.db)
277 }
278
279 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 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 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 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
348pub 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}