Skip to main content

lean_ctx/core/
multi_repo.rs

1use std::collections::HashMap;
2use std::path::{Path, PathBuf};
3
4use serde::{Deserialize, Serialize};
5
6use crate::core::bm25_index::{BM25Index, SearchResult};
7
8/// Default RRF parameter (controls how quickly rank decay affects fusion scores).
9const DEFAULT_RRF_K: f64 = 60.0;
10
11/// Maximum number of repo roots that can be served simultaneously.
12const MAX_ROOTS: usize = 16;
13
14/// A single search result from one repo root.
15#[derive(Debug, Clone)]
16pub struct RepoSearchResult {
17    pub repo_alias: String,
18    pub repo_path: String,
19    pub file_path: String,
20    pub symbol_name: String,
21    pub content: String,
22    pub start_line: usize,
23    pub end_line: usize,
24    pub score: f64,
25}
26
27/// A merged result after RRF fusion across multiple repos.
28#[derive(Debug, Clone)]
29pub struct FusedSearchResult {
30    pub repo_alias: String,
31    pub repo_path: String,
32    pub file_path: String,
33    pub symbol_name: String,
34    pub content: String,
35    pub start_line: usize,
36    pub end_line: usize,
37    pub rrf_score: f64,
38}
39
40/// Configuration for a single repository root in multi-repo mode.
41#[derive(Debug, Clone, Serialize, Deserialize)]
42pub struct RepoRootConfig {
43    pub path: String,
44    #[serde(default)]
45    pub alias: Option<String>,
46}
47
48impl RepoRootConfig {
49    pub fn effective_alias(&self) -> String {
50        self.alias.clone().unwrap_or_else(|| {
51            Path::new(&self.path)
52                .file_name()
53                .and_then(|n| n.to_str())
54                .unwrap_or("unknown")
55                .to_string()
56        })
57    }
58}
59
60/// Multi-repo configuration loaded from `~/.config/lean-ctx/multi-repo.toml`.
61#[derive(Debug, Clone, Serialize, Deserialize, Default)]
62pub struct MultiRepoConfig {
63    #[serde(default)]
64    pub repos: Vec<RepoRootConfig>,
65    #[serde(default)]
66    pub rrf_k: Option<f64>,
67}
68
69impl MultiRepoConfig {
70    pub fn load() -> Self {
71        let config_path = config_file_path();
72        if !config_path.exists() {
73            return Self::default();
74        }
75        match std::fs::read_to_string(&config_path) {
76            Ok(content) => toml::from_str(&content).unwrap_or_default(),
77            Err(_) => Self::default(),
78        }
79    }
80
81    pub fn save(&self) -> Result<(), String> {
82        let config_path = config_file_path();
83        if let Some(parent) = config_path.parent() {
84            std::fs::create_dir_all(parent)
85                .map_err(|e| format!("Failed to create config dir: {e}"))?;
86        }
87        let content =
88            toml::to_string_pretty(self).map_err(|e| format!("Failed to serialize config: {e}"))?;
89        let defaults = toml::to_string_pretty(&Self::default())
90            .map_err(|e| format!("Failed to serialize defaults: {e}"))?;
91        crate::config_io::write_toml_preserving_minimal(&config_path, &content, &defaults)
92            .map_err(|e| format!("Failed to write config: {e}"))?;
93        Ok(())
94    }
95}
96
97/// An active repo root with its loaded BM25 index.
98pub struct ActiveRepoRoot {
99    pub config: RepoRootConfig,
100    pub path: PathBuf,
101    index: Option<BM25Index>,
102}
103
104impl std::fmt::Debug for ActiveRepoRoot {
105    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
106        f.debug_struct("ActiveRepoRoot")
107            .field("config", &self.config)
108            .field("path", &self.path)
109            .field("has_index", &self.index.is_some())
110            .finish()
111    }
112}
113
114impl ActiveRepoRoot {
115    fn new(config: RepoRootConfig) -> Result<Self, String> {
116        let path = PathBuf::from(&config.path);
117        if !path.is_dir() {
118            return Err(format!(
119                "Path does not exist or is not a directory: {}",
120                config.path
121            ));
122        }
123        let path = path
124            .canonicalize()
125            .map_err(|e| format!("Cannot canonicalize {}: {e}", config.path))?;
126        Ok(Self {
127            config,
128            path,
129            index: None,
130        })
131    }
132
133    fn ensure_index(&mut self) {
134        if self.index.is_some() {
135            return;
136        }
137        self.index = Some(BM25Index::load_or_build(&self.path));
138    }
139
140    pub fn alias(&self) -> String {
141        self.config.effective_alias()
142    }
143
144    pub fn search(&mut self, query: &str, max_results: usize) -> Vec<RepoSearchResult> {
145        self.ensure_index();
146        let Some(ref index) = self.index else {
147            return Vec::new();
148        };
149
150        let results: Vec<SearchResult> = index.search(query, max_results);
151        let alias = self.alias();
152        let repo_path = self.path.to_string_lossy().to_string();
153
154        results
155            .into_iter()
156            .enumerate()
157            .map(|(rank, sr)| RepoSearchResult {
158                repo_alias: alias.clone(),
159                repo_path: repo_path.clone(),
160                file_path: sr.file_path,
161                symbol_name: sr.symbol_name,
162                content: sr.snippet,
163                start_line: sr.start_line,
164                end_line: sr.end_line,
165                score: 1.0 / (rank as f64 + 1.0),
166            })
167            .collect()
168    }
169}
170
171/// Manages multiple repository roots and performs cross-repo search with RRF fusion.
172pub struct MultiRepoManager {
173    roots: Vec<ActiveRepoRoot>,
174    rrf_k: f64,
175}
176
177impl MultiRepoManager {
178    pub fn new() -> Self {
179        Self {
180            roots: Vec::new(),
181            rrf_k: DEFAULT_RRF_K,
182        }
183    }
184
185    pub fn with_rrf_k(mut self, k: f64) -> Self {
186        self.rrf_k = k;
187        self
188    }
189
190    pub fn from_config(config: &MultiRepoConfig) -> Result<Self, String> {
191        let mut manager = Self::new();
192        if let Some(k) = config.rrf_k {
193            manager.rrf_k = k;
194        }
195        for repo_config in &config.repos {
196            manager.add_root_config(repo_config.clone())?;
197        }
198        Ok(manager)
199    }
200
201    pub fn add_root(&mut self, path: &str, alias: Option<&str>) -> Result<(), String> {
202        if self.roots.len() >= MAX_ROOTS {
203            return Err(format!("Maximum number of roots ({MAX_ROOTS}) reached"));
204        }
205        let config = RepoRootConfig {
206            path: path.to_string(),
207            alias: alias.map(String::from),
208        };
209        let root = ActiveRepoRoot::new(config)?;
210        if self.roots.iter().any(|r| r.path == root.path) {
211            return Err(format!("Root already exists: {path}"));
212        }
213        self.roots.push(root);
214        Ok(())
215    }
216
217    fn add_root_config(&mut self, config: RepoRootConfig) -> Result<(), String> {
218        if self.roots.len() >= MAX_ROOTS {
219            return Err(format!("Maximum number of roots ({MAX_ROOTS}) reached"));
220        }
221        let root = ActiveRepoRoot::new(config)?;
222        if self.roots.iter().any(|r| r.path == root.path) {
223            return Err(format!(
224                "Root already exists: {}",
225                root.path.to_string_lossy()
226            ));
227        }
228        self.roots.push(root);
229        Ok(())
230    }
231
232    pub fn remove_root(&mut self, path: &str) -> Result<(), String> {
233        let normalized = PathBuf::from(path)
234            .canonicalize()
235            .unwrap_or_else(|_| PathBuf::from(path));
236        let before = self.roots.len();
237        self.roots
238            .retain(|r| r.path != normalized && r.config.path != path);
239        if self.roots.len() == before {
240            return Err(format!("Root not found: {path}"));
241        }
242        Ok(())
243    }
244
245    pub fn list_roots(&self) -> Vec<RootInfo> {
246        self.roots
247            .iter()
248            .map(|r| RootInfo {
249                path: r.path.to_string_lossy().to_string(),
250                alias: r.alias(),
251                has_index: r.index.is_some(),
252            })
253            .collect()
254    }
255
256    pub fn root_count(&self) -> usize {
257        self.roots.len()
258    }
259
260    pub fn is_active(&self) -> bool {
261        self.roots.len() > 1
262    }
263
264    /// Resolve a repo alias or path to the corresponding root index.
265    pub fn resolve_root(&self, repo: &str) -> Option<usize> {
266        self.roots.iter().position(|r| {
267            r.alias() == repo || r.config.path == repo || r.path.to_string_lossy() == repo
268        })
269    }
270
271    /// Search across all roots (or a subset) and merge with Reciprocal Rank Fusion.
272    pub fn search(
273        &mut self,
274        query: &str,
275        max_results: usize,
276        filter_roots: Option<&[String]>,
277    ) -> Vec<FusedSearchResult> {
278        let per_root_max = (max_results * 2).max(20);
279
280        let mut all_results: HashMap<String, FusedSearchResult> = HashMap::new();
281
282        for root in &mut self.roots {
283            if let Some(filter) = filter_roots {
284                let alias = root.alias();
285                let path = root.path.to_string_lossy().to_string();
286                if !filter.iter().any(|f| f == &alias || f == &path) {
287                    continue;
288                }
289            }
290
291            let results = root.search(query, per_root_max);
292
293            for (rank, result) in results.iter().enumerate() {
294                let rrf_contribution = 1.0 / (self.rrf_k + rank as f64 + 1.0);
295                let key = format!(
296                    "{}:{}:{}",
297                    result.repo_alias, result.file_path, result.start_line
298                );
299
300                all_results
301                    .entry(key)
302                    .and_modify(|existing| {
303                        existing.rrf_score += rrf_contribution;
304                    })
305                    .or_insert_with(|| FusedSearchResult {
306                        repo_alias: result.repo_alias.clone(),
307                        repo_path: result.repo_path.clone(),
308                        file_path: result.file_path.clone(),
309                        symbol_name: result.symbol_name.clone(),
310                        content: result.content.clone(),
311                        start_line: result.start_line,
312                        end_line: result.end_line,
313                        rrf_score: rrf_contribution,
314                    });
315            }
316        }
317
318        let mut fused: Vec<FusedSearchResult> = all_results.into_values().collect();
319        fused.sort_by(|a, b| {
320            b.rrf_score
321                .partial_cmp(&a.rrf_score)
322                .unwrap_or(std::cmp::Ordering::Equal)
323        });
324        fused.truncate(max_results);
325        fused
326    }
327
328    /// Search within a specific repo root (no RRF, single-repo query).
329    pub fn search_single_repo(
330        &mut self,
331        repo: &str,
332        query: &str,
333        max_results: usize,
334    ) -> Result<Vec<RepoSearchResult>, String> {
335        let idx = self
336            .resolve_root(repo)
337            .ok_or_else(|| format!("Unknown repo: {repo}"))?;
338        Ok(self.roots[idx].search(query, max_results))
339    }
340}
341
342impl Default for MultiRepoManager {
343    fn default() -> Self {
344        Self::new()
345    }
346}
347
348/// Summary info about a registered root.
349#[derive(Debug, Clone, Serialize)]
350pub struct RootInfo {
351    pub path: String,
352    pub alias: String,
353    pub has_index: bool,
354}
355
356/// Returns the path to the multi-repo config file.
357///
358/// #594: resolved through the unified config base so it matches `config.toml`
359/// (on macOS this no longer diverges to `~/Library/Application Support`); any
360/// legacy copy at the old location is adopted on first access.
361pub fn config_file_path() -> PathBuf {
362    crate::core::paths::config_dir_member("multi-repo.toml")
363        .unwrap_or_else(|_| PathBuf::from("~/.config/lean-ctx/multi-repo.toml"))
364}
365
366/// Global multi-repo manager instance (lazily initialized).
367static GLOBAL_MANAGER: std::sync::OnceLock<std::sync::Mutex<MultiRepoManager>> =
368    std::sync::OnceLock::new();
369
370pub fn global_manager() -> &'static std::sync::Mutex<MultiRepoManager> {
371    GLOBAL_MANAGER.get_or_init(|| {
372        let config = MultiRepoConfig::load();
373        let manager = MultiRepoManager::from_config(&config).unwrap_or_default();
374        std::sync::Mutex::new(manager)
375    })
376}
377
378/// Initialize the global manager with explicit roots (e.g. from CLI `--root` flags).
379pub fn init_with_roots(
380    roots: &[(String, Option<String>)],
381    rrf_k: Option<f64>,
382) -> Result<(), String> {
383    let mut manager = MultiRepoManager::new();
384    if let Some(k) = rrf_k {
385        manager.rrf_k = k;
386    }
387    for (path, alias) in roots {
388        manager.add_root(path, alias.as_deref())?;
389    }
390    GLOBAL_MANAGER
391        .set(std::sync::Mutex::new(manager))
392        .map_err(|_| "Multi-repo manager already initialized".to_string())
393}
394
395/// Resolve a `repo` alias/path to the actual filesystem root.
396/// Used by existing tools (ctx_read, ctx_search, etc.) when a `repo` param is provided.
397/// Returns the absolute path to the repo root, or None if multi-repo is inactive or repo not found.
398pub fn resolve_repo_root(repo: &str) -> Option<String> {
399    let manager = global_manager();
400    let mgr = manager.lock().ok()?;
401    let idx = mgr.resolve_root(repo)?;
402    Some(mgr.roots[idx].path.to_string_lossy().to_string())
403}
404
405/// Check if multi-repo mode is active (more than 1 root configured).
406pub fn is_multi_repo_active() -> bool {
407    let manager = global_manager();
408    manager.lock().is_ok_and(|mgr| mgr.is_active())
409}
410
411/// Get all configured repo root paths (for tools that need to iterate).
412pub fn all_root_paths() -> Vec<String> {
413    let manager = global_manager();
414    let Ok(mgr) = manager.lock() else {
415        return Vec::new();
416    };
417    mgr.roots
418        .iter()
419        .map(|r| r.path.to_string_lossy().to_string())
420        .collect()
421}
422
423/// Format search results for MCP output.
424pub fn format_fused_results(results: &[FusedSearchResult]) -> String {
425    if results.is_empty() {
426        return "No results found across repos.".to_string();
427    }
428
429    let mut out = String::with_capacity(results.len() * 200);
430    out.push_str(&format!(
431        "Cross-repo results ({} matches):\n\n",
432        results.len()
433    ));
434
435    for (i, result) in results.iter().enumerate() {
436        out.push_str(&format!(
437            "{}. [{}] {}:{}-{} ({})\n   RRF: {:.4}\n",
438            i + 1,
439            result.repo_alias,
440            result.file_path,
441            result.start_line,
442            result.end_line,
443            result.symbol_name,
444            result.rrf_score,
445        ));
446        let preview: String = result
447            .content
448            .lines()
449            .take(3)
450            .collect::<Vec<_>>()
451            .join("\n");
452        if !preview.is_empty() {
453            out.push_str(&format!("   {}\n", preview.replace('\n', "\n   ")));
454        }
455        out.push('\n');
456    }
457    out
458}
459
460#[cfg(test)]
461mod tests {
462    use super::*;
463
464    #[test]
465    fn repo_root_config_effective_alias() {
466        let cfg = RepoRootConfig {
467            path: "/home/user/projects/backend".to_string(),
468            alias: None,
469        };
470        assert_eq!(cfg.effective_alias(), "backend");
471
472        let cfg_with_alias = RepoRootConfig {
473            path: "/home/user/projects/backend".to_string(),
474            alias: Some("api".to_string()),
475        };
476        assert_eq!(cfg_with_alias.effective_alias(), "api");
477    }
478
479    #[test]
480    fn multi_repo_config_default_is_empty() {
481        let cfg = MultiRepoConfig::default();
482        assert!(cfg.repos.is_empty());
483        assert!(cfg.rrf_k.is_none());
484    }
485
486    #[test]
487    fn multi_repo_config_deserialize() {
488        let toml_str = r#"
489rrf_k = 45.0
490
491[[repos]]
492path = "/home/user/backend"
493alias = "backend"
494
495[[repos]]
496path = "/home/user/frontend"
497"#;
498        let cfg: MultiRepoConfig = toml::from_str(toml_str).unwrap();
499        assert_eq!(cfg.repos.len(), 2);
500        assert_eq!(cfg.rrf_k, Some(45.0));
501        assert_eq!(cfg.repos[0].alias, Some("backend".to_string()));
502        assert_eq!(cfg.repos[1].alias, None);
503    }
504
505    #[test]
506    fn manager_max_roots_enforced() {
507        let mut manager = MultiRepoManager::new();
508        for i in 0..MAX_ROOTS {
509            let dir = std::env::temp_dir().join(format!("multi_repo_test_{i}"));
510            let _ = std::fs::create_dir_all(&dir);
511            let _ = manager.add_root(&dir.to_string_lossy(), Some(&format!("repo{i}")));
512        }
513        let extra = std::env::temp_dir().join("multi_repo_test_extra");
514        let _ = std::fs::create_dir_all(&extra);
515        let result = manager.add_root(&extra.to_string_lossy(), None);
516        assert!(result.is_err());
517
518        for i in 0..=MAX_ROOTS {
519            let dir = std::env::temp_dir().join(format!("multi_repo_test_{i}"));
520            let _ = std::fs::remove_dir_all(&dir);
521        }
522        let _ = std::fs::remove_dir_all(&extra);
523    }
524
525    #[test]
526    fn manager_duplicate_root_rejected() {
527        let dir = std::env::temp_dir().join("multi_repo_dup_test");
528        let _ = std::fs::create_dir_all(&dir);
529        let mut manager = MultiRepoManager::new();
530        assert!(
531            manager
532                .add_root(&dir.to_string_lossy(), Some("first"))
533                .is_ok()
534        );
535        assert!(
536            manager
537                .add_root(&dir.to_string_lossy(), Some("second"))
538                .is_err()
539        );
540        let _ = std::fs::remove_dir_all(&dir);
541    }
542
543    #[test]
544    fn rrf_fusion_basic() {
545        let manager = MultiRepoManager::new().with_rrf_k(60.0);
546        // RRF score for rank 0: 1/(60+0+1) = 1/61 ≈ 0.01639
547        let score: f64 = 1.0 / (60.0 + 0.0 + 1.0);
548        assert!((score - 0.01639).abs() < 0.001);
549
550        assert_eq!(manager.rrf_k, 60.0);
551    }
552
553    #[test]
554    fn remove_root_works() {
555        let dir = std::env::temp_dir().join("multi_repo_remove_test");
556        let _ = std::fs::create_dir_all(&dir);
557        let mut manager = MultiRepoManager::new();
558        manager
559            .add_root(&dir.to_string_lossy(), Some("removable"))
560            .unwrap();
561        assert_eq!(manager.root_count(), 1);
562        manager.remove_root(&dir.to_string_lossy()).unwrap();
563        assert_eq!(manager.root_count(), 0);
564        let _ = std::fs::remove_dir_all(&dir);
565    }
566
567    #[test]
568    fn list_roots_returns_info() {
569        let dir = std::env::temp_dir().join("multi_repo_list_test");
570        let _ = std::fs::create_dir_all(&dir);
571        let mut manager = MultiRepoManager::new();
572        manager
573            .add_root(&dir.to_string_lossy(), Some("myrepo"))
574            .unwrap();
575        let roots = manager.list_roots();
576        assert_eq!(roots.len(), 1);
577        assert_eq!(roots[0].alias, "myrepo");
578        let _ = std::fs::remove_dir_all(&dir);
579    }
580
581    #[test]
582    fn format_empty_results() {
583        let results: Vec<FusedSearchResult> = Vec::new();
584        let output = format_fused_results(&results);
585        assert!(output.contains("No results"));
586    }
587}