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> {
78 let workspace = canonicalize_or_owned(workspace.as_ref())?;
79 let brain_dir = workspace.join(".brain");
80 let db_path = brain_dir.join("db.sqlite");
81 if !db_path.exists() {
82 return Err(BrainError::BrainNotFound { path: brain_dir });
83 }
84 let db = Database::open(&db_path)?;
85 Ok(Self {
86 workspace,
87 brain_dir,
88 db,
89 })
90 }
91
92 pub fn open_or_create(workspace: impl AsRef<Path>) -> Result<Self> {
94 let workspace = workspace.as_ref();
95 let db_path = workspace.join(".brain").join("db.sqlite");
96 if db_path.exists() {
97 Self::open(workspace)
98 } else {
99 Self::create(workspace)
100 }
101 }
102
103 pub fn workspace(&self) -> &Path {
105 &self.workspace
106 }
107
108 pub fn brain_dir(&self) -> &Path {
110 &self.brain_dir
111 }
112
113 pub fn database(&self) -> &Database {
115 &self.db
116 }
117
118 pub fn sync(&mut self) -> Result<SyncStats> {
124 let db_path = self.brain_dir.join("db.sqlite");
125 let db = Database::open(&db_path)?;
126 let indexer = WorkspaceIndexer::new(db, self.workspace.clone());
127 let stats = indexer.index_workspace()?;
128 self.db = Database::open(&db_path)?;
130 Ok(stats)
131 }
132
133 pub fn query(&self, q: &str) -> Result<Vec<Node>> {
137 let hits = self.query_ranked(q, &QueryOptions::default())?;
138 Ok(hits.into_iter().map(|h| h.node).collect())
139 }
140
141 pub fn query_ranked(&self, q: &str, opts: &QueryOptions) -> Result<Vec<RankedHit>> {
145 self.db.search_ranked(q, opts)
146 }
147
148 pub fn context_for_prompt(&self, prompt: &str, max_tokens: usize) -> Result<ContextBundle> {
153 let opts = ContextOptions {
154 max_tokens,
155 ..ContextOptions::default()
156 };
157 self.context_for_prompt_with(prompt, &opts)
158 }
159
160 pub fn context_for_prompt_with(
165 &self,
166 prompt: &str,
167 opts: &ContextOptions,
168 ) -> Result<ContextBundle> {
169 assemble_context(&self.db, &self.brain_dir, prompt, opts)
170 }
171
172 pub fn export(&self, out: impl AsRef<Path>, decouple_ast: bool) -> Result<()> {
177 BrainExporter::export_bundle(&self.db, out, decouple_ast)
178 }
179
180 pub fn import(&mut self, input: impl AsRef<Path>) -> Result<usize> {
188 let n = BrainImporter::import_bundle(&self.db, input)?;
189 #[cfg(feature = "mmap")]
190 {
191 let indexer = WorkspaceIndexer::new(
192 Database::open(self.brain_dir.join("db.sqlite"))?,
193 self.workspace.clone(),
194 );
195 let _ = indexer.compile_mmap(&self.brain_dir.join("graph.mmap"));
196 self.db = Database::open(self.brain_dir.join("db.sqlite"))?;
197 }
198 Ok(n)
199 }
200
201 pub fn watch(&self, debounce_ms: u64) -> Result<()> {
207 crate::watch::watch_workspace(
208 &self.workspace,
209 crate::watch::WatchConfig {
210 debounce: std::time::Duration::from_millis(debounce_ms),
211 verbose: true,
212 },
213 )
214 }
215
216 pub fn bootstrap(
221 workspace: impl AsRef<Path>,
222 opts: crate::bootstrap::BootstrapOptions,
223 ) -> Result<crate::bootstrap::BootstrapReport> {
224 crate::bootstrap::bootstrap_workspace(workspace.as_ref(), opts)
225 }
226
227 pub fn doctor(&self) -> Result<crate::doctor::DoctorReport> {
229 crate::doctor::run_doctor(&self.workspace)
230 }
231
232 pub fn note_new(&self, opts: &crate::note::NoteNewOptions) -> Result<crate::note::NoteCreated> {
236 crate::note::create_note(&self.workspace, opts)
237 }
238}
239
240fn canonicalize_or_owned(path: &Path) -> Result<PathBuf> {
241 if path.exists() {
242 Ok(fs_canonicalize(path)?)
243 } else {
244 std::fs::create_dir_all(path)?;
245 Ok(fs_canonicalize(path)?)
246 }
247}
248
249fn fs_canonicalize(path: &Path) -> Result<PathBuf> {
250 std::fs::canonicalize(path).map_err(BrainError::from)
251}
252
253#[cfg(test)]
254mod tests {
255 use super::*;
256 use crate::types::ContextRole;
257 use tempfile::tempdir;
258
259 #[test]
260 fn create_sync_query_context_export() {
261 let dir = tempdir().unwrap();
262 let docs = dir.path().join("docs");
263 std::fs::create_dir_all(&docs).unwrap();
264 std::fs::write(
265 docs.join("raft.md"),
266 "---\ntags: [raft]\nnode_type: concept\n---\n# Raft\nSee [[logcompaction]].\n",
267 )
268 .unwrap();
269 std::fs::write(
270 docs.join("logcompaction.md"),
271 "---\ntags: [log]\nnode_type: concept\n---\n# Log Compaction\nSee [[raft]].\n",
272 )
273 .unwrap();
274
275 let mut brain = Brain::create(dir.path()).unwrap();
276 let stats = brain.sync().unwrap();
277 assert_eq!(stats.markdown_files, 2);
278
279 let hits = brain.query("raft").unwrap();
280 assert!(!hits.is_empty());
281
282 let ranked = brain
283 .query_ranked("raft", &QueryOptions::default())
284 .unwrap();
285 assert!(ranked[0].score > 0.0);
286
287 let ctx = brain.context_for_prompt("raft", 512).unwrap();
288 assert!(
289 !ctx.nodes.is_empty(),
290 "expected FTS hits for 'raft', got none"
291 );
292 assert!(ctx.tokens_used > 0);
293 assert!(ctx.nodes.iter().any(|n| n.role == ContextRole::Seed));
294 let xml = ctx.to_xml();
295 assert!(xml.contains("<rustbrain_context"));
296 assert!(xml.contains("tokens_used="));
297
298 let out = dir.path().join("out.brainbundle");
299 brain.export(&out, true).unwrap();
300 assert!(out.exists());
301
302 let _ = brain.sync().unwrap();
303 assert_eq!(brain.database().count_fts_rows().unwrap(), 2);
304 }
305
306 #[test]
307 fn note_anchors_to_symbol() {
308 let dir = tempdir().unwrap();
309 let docs = dir.path().join("docs");
310 let src = dir.path().join("src");
311 std::fs::create_dir_all(&docs).unwrap();
312 std::fs::create_dir_all(&src).unwrap();
313 std::fs::write(
314 dir.path().join("Cargo.toml"),
315 "[package]\nname = \"demo\"\nversion = \"0.1.0\"\nedition = \"2021\"\n",
316 )
317 .unwrap();
318 std::fs::write(
319 src.join("lib.rs"),
320 "/// Storage engine\npub struct StorageEngine;\nimpl StorageEngine { pub fn open() {} }\n",
321 )
322 .unwrap();
323 std::fs::write(
324 docs.join("design.md"),
325 "---\nnode_type: adr\n---\n# Design\nUses symbol:StorageEngine for persistence.\n",
326 )
327 .unwrap();
328
329 let mut brain = Brain::create(dir.path()).unwrap();
330 let stats = brain.sync().unwrap();
331 assert!(stats.symbol_anchors >= 1);
332 assert!(stats.markdown_files >= 1);
333
334 let _ = brain.sync().unwrap();
335 let edges = brain.database().get_all_edges().unwrap();
336 let anchors: Vec<_> = edges
337 .iter()
338 .filter(|e| e.relation_type == "anchors")
339 .collect();
340 assert!(
341 !anchors.is_empty(),
342 "expected anchors edge note→symbol, edges={edges:?}"
343 );
344 }
345}