Skip to main content

rustbrain_core/
registry.rs

1//! Global developer registry of rustbrain workspaces (XDG-aware).
2//!
3//! The registry powers `rustbrain query --all-workspaces`: each `rustbrain init`
4//! / `sync` registers the workspace path under the user config directory so
5//! multi-repo developers can search every brain on the machine.
6
7use crate::error::{BrainError, Result};
8use crate::query::{QueryOptions, RankedHit};
9use crate::storage::Database;
10use crate::types::Node;
11use serde::{Deserialize, Serialize};
12use std::fs;
13use std::path::{Path, PathBuf};
14
15/// Global registry of rustbrain workspaces on this machine.
16///
17/// Stored at `$XDG_CONFIG_HOME/rustbrain/registry.json` (Linux), or the platform
18/// equivalent returned by the `dirs` crate.
19#[derive(Debug, Clone, Serialize, Deserialize, Default)]
20pub struct GlobalRegistry {
21    /// Absolute workspace paths (canonicalized when possible).
22    #[serde(default)]
23    pub workspaces: Vec<String>,
24}
25
26/// A ranked hit tagged with its originating workspace.
27#[derive(Debug, Clone, Serialize, Deserialize)]
28pub struct GlobalRankedHit {
29    /// Absolute workspace path.
30    pub workspace: String,
31    /// Hit within that workspace's brain.
32    pub hit: RankedHit,
33}
34
35impl GlobalRegistry {
36    /// Config path: `$XDG_CONFIG_HOME/rustbrain/registry.json` or platform equivalent.
37    pub fn config_path() -> Result<PathBuf> {
38        let base = dirs::config_dir()
39            .ok_or_else(|| BrainError::other("could not determine user config directory"))?;
40        let dir = base.join("rustbrain");
41        fs::create_dir_all(&dir)?;
42        Ok(dir.join("registry.json"))
43    }
44
45    /// Load registry, pruning missing workspace paths (and persisting if anything dropped).
46    pub fn load() -> Result<Self> {
47        let path = Self::config_path()?;
48        if !path.exists() {
49            return Ok(Self::default());
50        }
51        let content = fs::read_to_string(&path)?;
52        let mut reg: Self = serde_json::from_str(&content).unwrap_or_default();
53        let before = reg.workspaces.len();
54        reg.prune_missing();
55        if reg.workspaces.len() != before {
56            let _ = reg.save();
57        }
58        Ok(reg)
59    }
60
61    /// Prune missing paths and save. Returns number of entries removed.
62    pub fn gc(&mut self) -> Result<usize> {
63        let before = self.workspaces.len();
64        self.prune_missing();
65        let removed = before.saturating_sub(self.workspaces.len());
66        if removed > 0 {
67            self.save()?;
68        }
69        Ok(removed)
70    }
71
72    /// Persist registry atomically.
73    pub fn save(&self) -> Result<()> {
74        let path = Self::config_path()?;
75        let tmp = path.with_extension("json.tmp");
76        let json = serde_json::to_string_pretty(self)?;
77        fs::write(&tmp, json)?;
78        fs::rename(&tmp, &path)?;
79        Ok(())
80    }
81
82    /// Drop workspace paths that no longer have a brain database.
83    pub fn prune_missing(&mut self) {
84        self.workspaces
85            .retain(|ws| Path::new(ws).join(".brain").join("db.sqlite").exists());
86    }
87
88    /// Register a workspace directory path.
89    pub fn register<P: AsRef<Path>>(&mut self, workspace_path: P) -> Result<()> {
90        let abs_path = fs::canonicalize(workspace_path.as_ref())
91            .unwrap_or_else(|_| workspace_path.as_ref().to_path_buf())
92            .to_string_lossy()
93            .to_string();
94
95        if !self.workspaces.contains(&abs_path) {
96            self.workspaces.push(abs_path);
97            self.save()?;
98        }
99        Ok(())
100    }
101
102    /// Search FTS across all registered workspaces (unranked grouping).
103    pub fn search_all(&self, query: &str) -> Result<Vec<(String, Vec<Node>)>> {
104        let ranked = self.search_all_ranked(query, &QueryOptions::default())?;
105        let mut by_ws: std::collections::BTreeMap<String, Vec<Node>> =
106            std::collections::BTreeMap::new();
107        for gh in ranked {
108            by_ws.entry(gh.workspace).or_default().push(gh.hit.node);
109        }
110        Ok(by_ws.into_iter().collect())
111    }
112
113    /// Cross-workspace ranked merge: highest score first, workspace tag preserved.
114    pub fn search_all_ranked(
115        &self,
116        query: &str,
117        opts: &QueryOptions,
118    ) -> Result<Vec<GlobalRankedHit>> {
119        let mut merged: Vec<GlobalRankedHit> = Vec::new();
120
121        for ws in &self.workspaces {
122            let db_path = PathBuf::from(ws).join(".brain").join("db.sqlite");
123            if !db_path.exists() {
124                continue;
125            }
126            match Database::open(&db_path) {
127                Ok(db) => match db.search_ranked(query, opts) {
128                    Ok(hits) => {
129                        for hit in hits {
130                            merged.push(GlobalRankedHit {
131                                workspace: ws.clone(),
132                                hit,
133                            });
134                        }
135                    }
136                    Err(e) => {
137                        eprintln!("rustbrain: skip workspace {ws}: {e}");
138                    }
139                },
140                Err(e) => {
141                    eprintln!("rustbrain: cannot open {ws}: {e}");
142                }
143            }
144        }
145
146        merged.sort_by(|a, b| {
147            b.hit
148                .score
149                .partial_cmp(&a.hit.score)
150                .unwrap_or(std::cmp::Ordering::Equal)
151        });
152        merged.truncate(opts.limit);
153        Ok(merged)
154    }
155}