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