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