Skip to main content

lean_ctx/core/bm25_index/
mod.rs

1use std::collections::HashMap;
2use std::path::{Path, PathBuf};
3use std::time::UNIX_EPOCH;
4
5use serde::{Deserialize, Serialize};
6mod chunking;
7pub use chunking::*;
8#[cfg(test)]
9mod tests;
10
11const MAX_BM25_FILES: usize = 5000;
12const CHUNK_COUNT_WARNING: usize = 50_000;
13const ZSTD_LEVEL: i32 = 9;
14
15const DEFAULT_BM25_IGNORES: &[&str] = &[
16    "vendor/**",
17    "dist/**",
18    "build/**",
19    "public/vendor/**",
20    "public/js/**",
21    "public/css/**",
22    "public/build/**",
23    ".next/**",
24    ".nuxt/**",
25    "__pycache__/**",
26    "*.min.js",
27    "*.min.css",
28    "*.bundle.js",
29    "*.chunk.js",
30];
31
32fn max_bm25_cache_bytes() -> u64 {
33    // Single source of truth: `Config::bm25_max_cache_mb_effective` (env override
34    // › explicit config › disk-budget › generous default). Decoupled from the RAM
35    // profile so large repos persist instead of rebuilding forever (issue #249).
36    let mb = std::env::var("LEAN_CTX_BM25_MAX_CACHE_MB")
37        .ok()
38        .and_then(|v| v.parse::<u64>().ok())
39        .unwrap_or_else(|| crate::core::config::Config::load().bm25_max_cache_mb_effective());
40    mb * 1024 * 1024
41}
42
43/// Effective on-disk ceiling (bytes) for the persisted BM25 index. Single source
44/// of truth shared with `doctor` so its "oversized index" warning matches what
45/// `save`/`load` actually enforce.
46pub fn persist_ceiling_bytes() -> u64 {
47    max_bm25_cache_bytes()
48}
49
50/// Outcome of persisting a BM25 index to disk. Distinguishes a real write from a
51/// size-capped refusal so callers never mistake "refused to persist" for
52/// success (the bug behind the perpetual "index warming" report, issue #249).
53#[derive(Debug, Clone, Copy, PartialEq, Eq)]
54pub enum SaveOutcome {
55    /// Written to disk. Carries the compressed (zstd) size in bytes.
56    Persisted { compressed_bytes: u64 },
57    /// Built fine but NOT written — the compressed size exceeds the disk
58    /// ceiling. The in-memory index is still usable for this process; callers
59    /// should surface the remedy (raise the cap / add ignore patterns) instead
60    /// of silently rebuilding on every call.
61    SkippedTooLarge {
62        compressed_bytes: u64,
63        limit_bytes: u64,
64    },
65}
66
67#[derive(Debug, Clone, Serialize, Deserialize)]
68pub struct CodeChunk {
69    pub file_path: String,
70    pub symbol_name: String,
71    pub kind: ChunkKind,
72    pub start_line: usize,
73    pub end_line: usize,
74    pub content: String,
75    #[serde(default)]
76    pub tokens: Vec<String>,
77    pub token_count: usize,
78}
79
80#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
81pub enum ChunkKind {
82    Function,
83    Struct,
84    Impl,
85    Module,
86    Class,
87    Method,
88    Other,
89    // -- External source kinds (Context Engine) --
90    Issue,
91    PullRequest,
92    WikiPage,
93    DbSchema,
94    ApiEndpoint,
95    Ticket,
96    ExternalOther,
97}
98
99#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
100pub struct IndexedFileState {
101    pub mtime_ms: u64,
102    pub size_bytes: u64,
103}
104
105impl IndexedFileState {
106    fn from_path(path: &Path) -> Option<Self> {
107        let meta = path.metadata().ok()?;
108        let size_bytes = meta.len();
109        let mtime_ms = meta
110            .modified()
111            .ok()
112            .and_then(|t| t.duration_since(UNIX_EPOCH).ok())
113            .map(|d| d.as_millis() as u64)?;
114        Some(Self {
115            mtime_ms,
116            size_bytes,
117        })
118    }
119}
120
121#[derive(Debug, Clone, Serialize, Deserialize)]
122pub struct BM25Index {
123    pub chunks: Vec<CodeChunk>,
124    pub inverted: HashMap<String, Vec<(usize, f64)>>,
125    pub avg_doc_len: f64,
126    pub doc_count: usize,
127    pub doc_freqs: HashMap<String, usize>,
128    #[serde(default)]
129    pub files: HashMap<String, IndexedFileState>,
130    /// True once `shrink_resident_content_to_snippet` has trimmed each chunk's
131    /// `content` down to the snippet lines. Resident-only RAM-saving state: never
132    /// persisted (`skip`) so the on-disk index keeps full content, and a reload
133    /// always starts as `false`. Guards the embedding pass against re-embedding
134    /// truncated bodies (see `ensure_embeddings`).
135    #[serde(default, skip)]
136    pub content_truncated: bool,
137}
138
139#[derive(Debug, Clone, Serialize, Deserialize)]
140pub struct SearchResult {
141    pub chunk_idx: usize,
142    pub score: f64,
143    pub file_path: String,
144    pub symbol_name: String,
145    pub kind: ChunkKind,
146    pub start_line: usize,
147    pub end_line: usize,
148    pub snippet: String,
149}
150
151const BM25_K1: f64 = 1.2;
152const BM25_B: f64 = 0.75;
153
154impl Default for BM25Index {
155    fn default() -> Self {
156        Self::new()
157    }
158}
159
160impl BM25Index {
161    pub fn new() -> Self {
162        Self {
163            chunks: Vec::new(),
164            inverted: HashMap::new(),
165            avg_doc_len: 0.0,
166            doc_count: 0,
167            doc_freqs: HashMap::new(),
168            files: HashMap::new(),
169            content_truncated: false,
170        }
171    }
172
173    /// Approximate heap memory used by this index in bytes.
174    pub fn memory_usage_bytes(&self) -> usize {
175        let chunks_size: usize = self
176            .chunks
177            .iter()
178            .map(|c| {
179                c.content.len()
180                    + c.file_path.len()
181                    + c.symbol_name.len()
182                    + c.tokens.iter().map(String::len).sum::<usize>()
183                    + 64
184            })
185            .sum();
186        let inverted_size: usize = self
187            .inverted
188            .iter()
189            .map(|(k, v)| k.len() + v.len() * 16 + 32)
190            .sum();
191        let files_size: usize = self.files.keys().map(|k| k.len() + 24).sum();
192        let freqs_size: usize = self.doc_freqs.keys().map(|k| k.len() + 16).sum();
193        chunks_size + inverted_size + files_size + freqs_size
194    }
195
196    /// Drops all in-memory data, effectively freeing heap. Index can be re-loaded from disk.
197    pub fn unload(&mut self) {
198        let usage = self.memory_usage_bytes();
199        self.chunks = Vec::new();
200        self.inverted = HashMap::new();
201        self.doc_freqs = HashMap::new();
202        self.files = HashMap::new();
203        self.avg_doc_len = 0.0;
204        self.doc_count = 0;
205        tracing::info!(
206            "[bm25] unloaded index, freed ~{:.1}MB",
207            usage as f64 / 1_048_576.0
208        );
209    }
210
211    /// Shrinks each resident chunk's `content` to its first `keep_lines` lines,
212    /// reclaiming the RAM held by the full source bodies once the embedding pass
213    /// has already consumed them. The search path only ever reads
214    /// `content.lines().take(5)` for snippets, so the trimmed copy is functionally
215    /// complete for BM25/dense/hybrid result rendering.
216    ///
217    /// RESIDENT-ONLY: this mutates the in-memory copy. The persisted `.bin.zst`
218    /// keeps full content (truncation never runs before `save`), so a reload
219    /// restores complete bodies. Sets `content_truncated` so a later
220    /// `ensure_embeddings` against this same resident index skips re-embedding
221    /// (which would otherwise feed truncated bodies to the embedder).
222    ///
223    /// Idempotent and only ever shrinks: chunks shorter than `keep_lines` are
224    /// left untouched.
225    pub fn shrink_resident_content_to_snippet(&mut self, keep_lines: usize) {
226        let before = self.memory_usage_bytes();
227        for chunk in &mut self.chunks {
228            // Cheap line-count gate: only allocate a new string when the body is
229            // actually longer than the snippet window.
230            if chunk.content.lines().nth(keep_lines).is_some() {
231                let trimmed: String = chunk
232                    .content
233                    .lines()
234                    .take(keep_lines)
235                    .collect::<Vec<_>>()
236                    .join("\n");
237                chunk.content = trimmed;
238                // Reclaim the spare capacity left by the larger original body.
239                chunk.content.shrink_to_fit();
240            }
241        }
242        self.content_truncated = true;
243        let after = self.memory_usage_bytes();
244        tracing::debug!(
245            "[bm25] shrank resident content to {keep_lines} lines/chunk, freed ~{:.1}MB",
246            before.saturating_sub(after) as f64 / 1_048_576.0
247        );
248    }
249
250    /// Builds an index from explicit chunks (unit tests; avoids filesystem walking).
251    #[cfg(test)]
252    pub(crate) fn from_chunks_for_test(chunks: Vec<CodeChunk>) -> Self {
253        let mut index = Self::new();
254        for mut chunk in chunks {
255            if chunk.token_count == 0 {
256                chunk.token_count = tokenize(&chunk.content).len();
257            }
258            index.add_chunk(chunk);
259        }
260        index.finalize();
261        index
262    }
263
264    pub fn build_from_directory(root: &Path) -> Self {
265        Self::build_from_directory_inner(root, &HashMap::new())
266    }
267
268    /// Like `build_from_directory` but reuses file content from a prior scan
269    /// (e.g. the graph index walk) to avoid redundant disk reads.
270    pub fn build_with_content_hint(root: &Path, content_hint: &HashMap<String, String>) -> Self {
271        Self::build_from_directory_inner(root, content_hint)
272    }
273
274    fn build_from_directory_inner(root: &Path, content_hint: &HashMap<String, String>) -> Self {
275        let root_str = root.to_string_lossy();
276        if !super::graph_index::is_safe_scan_root_public(&root_str) {
277            tracing::warn!("[bm25: scan aborted for unsafe root {root_str}]");
278            return Self::new();
279        }
280        let mut index = Self::new();
281        let files = list_code_files(root);
282        const MAX_FILE_SIZE_BYTES: u64 = 2 * 1024 * 1024;
283        let mut cache_hits = 0usize;
284
285        for (i, rel) in files.iter().enumerate() {
286            if i.is_multiple_of(500) && crate::core::memory_guard::is_under_pressure() {
287                tracing::warn!(
288                    "[bm25: stopping build at file {i}/{} due to memory pressure]",
289                    files.len()
290                );
291                break;
292            }
293            if crate::core::memory_guard::abort_requested() {
294                tracing::warn!("[bm25: aborting build due to critical memory pressure]");
295                break;
296            }
297
298            let abs = root.join(rel);
299            let Some(state) = IndexedFileState::from_path(&abs) else {
300                continue;
301            };
302            if state.size_bytes > MAX_FILE_SIZE_BYTES {
303                continue;
304            }
305
306            // Content sources, cheapest first: an explicit per-build hint, then
307            // the shared resident content cache (populated by the search-index
308            // build / ctx_search, issue #148) validated by `(mtime, size)`, then
309            // a one-time disk read that also publishes into the shared cache.
310            let cache_state = crate::core::content_cache::FileState {
311                mtime_ms: state.mtime_ms,
312                size_bytes: state.size_bytes,
313            };
314            let content = if crate::core::extractors::is_binary_document(&abs) {
315                // Binary document (PDF, …): extract clean text from raw bytes.
316                // Skipped if extraction yields nothing (e.g. scanned/image-only).
317                // Never populates the shared UTF-8 content cache (not text).
318                match std::fs::read(&abs) {
319                    Ok(bytes) => {
320                        let text = crate::core::extractors::extract(&abs, &bytes).text;
321                        if text.is_empty() {
322                            continue;
323                        }
324                        std::borrow::Cow::Owned(text)
325                    }
326                    Err(_) => continue,
327                }
328            } else if let Some(cached) = content_hint.get(rel) {
329                cache_hits += 1;
330                std::borrow::Cow::Borrowed(cached.as_str())
331            } else if let Some(arc) = crate::core::content_cache::get(&abs, cache_state) {
332                cache_hits += 1;
333                std::borrow::Cow::Owned(arc.to_string())
334            } else {
335                match std::fs::read_to_string(&abs) {
336                    Ok(c) => {
337                        crate::core::content_cache::insert(
338                            &abs,
339                            cache_state,
340                            std::sync::Arc::from(c.as_str()),
341                        );
342                        std::borrow::Cow::Owned(c)
343                    }
344                    Err(_) => continue,
345                }
346            };
347
348            let mut chunks = extract_chunks(rel, &content);
349            chunks.sort_by(|a, b| {
350                a.start_line
351                    .cmp(&b.start_line)
352                    .then_with(|| a.end_line.cmp(&b.end_line))
353                    .then_with(|| a.symbol_name.cmp(&b.symbol_name))
354            });
355            for chunk in chunks {
356                index.add_chunk(chunk);
357            }
358            index.files.insert(rel.clone(), state);
359        }
360
361        if cache_hits > 0 {
362            tracing::info!(
363                "[bm25: reused {cache_hits}/{} file contents from graph scan cache]",
364                files.len()
365            );
366        }
367
368        index.finalize();
369        index
370    }
371
372    pub fn rebuild_incremental(root: &Path, prev: &BM25Index) -> Self {
373        let mut old_by_file: HashMap<String, Vec<CodeChunk>> = HashMap::new();
374        for c in &prev.chunks {
375            old_by_file
376                .entry(c.file_path.clone())
377                .or_default()
378                .push(c.clone());
379        }
380        for v in old_by_file.values_mut() {
381            v.sort_by(|a, b| {
382                a.start_line
383                    .cmp(&b.start_line)
384                    .then_with(|| a.end_line.cmp(&b.end_line))
385                    .then_with(|| a.symbol_name.cmp(&b.symbol_name))
386            });
387        }
388
389        let mut index = Self::new();
390        let files = list_code_files(root);
391        const MAX_FILE_SIZE_BYTES: u64 = 2 * 1024 * 1024;
392
393        for (i, rel) in files.iter().enumerate() {
394            if i.is_multiple_of(500) && crate::core::memory_guard::is_under_pressure() {
395                tracing::warn!(
396                    "[bm25: stopping incremental rebuild at file {i}/{} due to memory pressure]",
397                    files.len()
398                );
399                break;
400            }
401
402            let abs = root.join(rel);
403            let Some(state) = IndexedFileState::from_path(&abs) else {
404                continue;
405            };
406
407            let unchanged = prev.files.get(rel).is_some_and(|old| *old == state);
408            if unchanged {
409                if let Some(chunks) = old_by_file.get(rel) {
410                    if chunks.first().is_some_and(|c| !c.content.is_empty()) {
411                        for chunk in chunks {
412                            index.add_chunk(chunk.clone());
413                        }
414                        index.files.insert(rel.clone(), state);
415                        continue;
416                    }
417                }
418            }
419
420            if state.size_bytes > MAX_FILE_SIZE_BYTES {
421                continue;
422            }
423            let content = if crate::core::extractors::is_binary_document(&abs) {
424                match std::fs::read(&abs) {
425                    Ok(bytes) => crate::core::extractors::extract(&abs, &bytes).text,
426                    Err(_) => continue,
427                }
428            } else {
429                match std::fs::read_to_string(&abs) {
430                    Ok(c) => c,
431                    Err(_) => continue,
432                }
433            };
434            if content.is_empty() {
435                continue;
436            }
437            let mut chunks = extract_chunks(rel, &content);
438            chunks.sort_by(|a, b| {
439                a.start_line
440                    .cmp(&b.start_line)
441                    .then_with(|| a.end_line.cmp(&b.end_line))
442                    .then_with(|| a.symbol_name.cmp(&b.symbol_name))
443            });
444            for chunk in chunks {
445                index.add_chunk(chunk);
446            }
447            index.files.insert(rel.clone(), state);
448        }
449
450        index.finalize();
451        index
452    }
453
454    fn add_chunk(&mut self, chunk: CodeChunk) {
455        let idx = self.chunks.len();
456
457        let enriched = enrich_for_bm25(&chunk);
458        let tokens = tokenize(&enriched);
459        for token in &tokens {
460            let lower = token.to_lowercase();
461            let postings = self.inverted.entry(lower.clone()).or_default();
462            if postings.last().map(|(last_idx, _)| *last_idx) != Some(idx) {
463                *self.doc_freqs.entry(lower).or_insert(0) += 1;
464            }
465            postings.push((idx, 1.0));
466        }
467
468        self.chunks.push(CodeChunk {
469            token_count: tokens.len(),
470            tokens: Vec::new(),
471            ..chunk
472        });
473    }
474
475    fn finalize(&mut self) {
476        self.doc_count = self.chunks.len();
477        if self.doc_count == 0 {
478            return;
479        }
480
481        let total_len: usize = self.chunks.iter().map(|c| c.token_count).sum();
482        self.avg_doc_len = total_len as f64 / self.doc_count as f64;
483    }
484
485    pub fn search(&self, query: &str, top_k: usize) -> Vec<SearchResult> {
486        let query_tokens = tokenize(query);
487        if query_tokens.is_empty() || self.doc_count == 0 {
488            return Vec::new();
489        }
490
491        // Pre-allocated score array: O(1) per-access vs HashMap overhead.
492        // Kolmogorov-optimal: minimal allocation for the scoring operation.
493        let n = self.chunks.len();
494        let mut scores = vec![0.0f64; n];
495        let mut touched = Vec::with_capacity(n.min(256));
496
497        for token in &query_tokens {
498            let lower = token.to_lowercase();
499            let df = *self.doc_freqs.get(&lower).unwrap_or(&0) as f64;
500            if df == 0.0 {
501                continue;
502            }
503
504            let idf = ((self.doc_count as f64 - df + 0.5) / (df + 0.5) + 1.0).ln();
505
506            if let Some(postings) = self.inverted.get(&lower) {
507                for &(idx, weight) in postings {
508                    let doc_len = self.chunks[idx].token_count as f64;
509                    let norm_len = doc_len / self.avg_doc_len.max(1.0);
510                    let bm25 = idf * (weight * (BM25_K1 + 1.0))
511                        / (weight + BM25_K1 * (1.0 - BM25_B + BM25_B * norm_len));
512
513                    if scores[idx] == 0.0 {
514                        touched.push(idx);
515                    }
516                    scores[idx] += bm25;
517                }
518            }
519        }
520
521        let mut results: Vec<SearchResult> = touched
522            .iter()
523            .filter(|&&idx| scores[idx] > 0.0)
524            .map(|&idx| {
525                let chunk = &self.chunks[idx];
526                let snippet = chunk.content.lines().take(5).collect::<Vec<_>>().join("\n");
527                SearchResult {
528                    chunk_idx: idx,
529                    score: scores[idx],
530                    file_path: chunk.file_path.clone(),
531                    symbol_name: chunk.symbol_name.clone(),
532                    kind: chunk.kind.clone(),
533                    start_line: chunk.start_line,
534                    end_line: chunk.end_line,
535                    snippet,
536                }
537            })
538            .collect();
539
540        results.sort_by(|a, b| {
541            b.score
542                .partial_cmp(&a.score)
543                .unwrap_or(std::cmp::Ordering::Equal)
544                .then_with(|| a.file_path.cmp(&b.file_path))
545                .then_with(|| a.symbol_name.cmp(&b.symbol_name))
546                .then_with(|| a.start_line.cmp(&b.start_line))
547                .then_with(|| a.end_line.cmp(&b.end_line))
548        });
549        results.truncate(top_k);
550        results
551    }
552
553    pub fn save(&self, root: &Path) -> std::io::Result<SaveOutcome> {
554        if self.chunks.len() > CHUNK_COUNT_WARNING {
555            tracing::warn!(
556                "[bm25] index has {} chunks (threshold {}), consider adding extra_ignore_patterns",
557                self.chunks.len(),
558                CHUNK_COUNT_WARNING
559            );
560        }
561
562        let dir = index_dir(root);
563        std::fs::create_dir_all(&dir)?;
564        let data = bincode::serde::encode_to_vec(self, bincode::config::standard())
565            .map_err(|e| std::io::Error::other(e.to_string()))?;
566
567        let compressed = zstd::encode_all(data.as_slice(), ZSTD_LEVEL)
568            .map_err(|e| std::io::Error::other(format!("zstd compress: {e}")))?;
569        let compressed_bytes = compressed.len() as u64;
570
571        let max_bytes = max_bm25_cache_bytes();
572        if compressed_bytes > max_bytes {
573            // Do NOT pretend success: a silent `Ok(())` here made `load` return
574            // `None` forever and the index rebuild on every call (issue #249).
575            // Report the refusal so the orchestrator can record an actionable
576            // note and the agent-facing tools can stop claiming the index will
577            // be "ready next call".
578            tracing::warn!(
579                "[bm25] compressed index too large ({:.1} MB, limit {:.0} MB), refusing to persist: {}",
580                compressed_bytes as f64 / 1_048_576.0,
581                max_bytes / (1024 * 1024),
582                dir.display()
583            );
584            return Ok(SaveOutcome::SkippedTooLarge {
585                compressed_bytes,
586                limit_bytes: max_bytes,
587            });
588        }
589
590        tracing::info!(
591            "[bm25] index: {:.1} MB bincode → {:.1} MB zstd ({:.0}% saved)",
592            data.len() as f64 / 1_048_576.0,
593            compressed_bytes as f64 / 1_048_576.0,
594            (1.0 - compressed_bytes as f64 / data.len().max(1) as f64) * 100.0
595        );
596
597        let target = dir.join("bm25_index.bin.zst");
598        let tmp = dir.join("bm25_index.bin.zst.tmp");
599        std::fs::write(&tmp, &compressed)?;
600        std::fs::rename(&tmp, &target)?;
601
602        let _ = std::fs::remove_file(dir.join("bm25_index.bin"));
603        let _ = std::fs::remove_file(dir.join("bm25_index.json"));
604
605        let _ = std::fs::write(
606            dir.join("project_root.txt"),
607            root.to_string_lossy().as_bytes(),
608        );
609
610        Ok(SaveOutcome::Persisted { compressed_bytes })
611    }
612
613    pub fn load(root: &Path) -> Option<Self> {
614        let dir = index_dir(root);
615        let max_bytes = max_bm25_cache_bytes();
616
617        let zst_path = dir.join("bm25_index.bin.zst");
618        if zst_path.exists() {
619            let meta = std::fs::metadata(&zst_path).ok()?;
620            if meta.len() > max_bytes {
621                tracing::warn!(
622                    "[bm25] compressed index too large ({:.1} GB, limit {:.0} MB), quarantining: {}",
623                    meta.len() as f64 / 1_073_741_824.0,
624                    max_bytes / (1024 * 1024),
625                    zst_path.display()
626                );
627                let quarantined = zst_path.with_extension("zst.quarantined");
628                let _ = std::fs::rename(&zst_path, &quarantined);
629                return None;
630            }
631            let compressed = std::fs::read(&zst_path).ok()?;
632            let max_decompressed = max_bytes * 20; // allow 20x expansion ratio
633            let data = bounded_zstd_decode(&compressed, max_decompressed)?;
634            let (idx, _): (Self, _) =
635                bincode::serde::decode_from_slice(&data, bincode::config::standard()).ok()?;
636            return Some(idx);
637        }
638
639        let bin_path = dir.join("bm25_index.bin");
640        if bin_path.exists() {
641            let meta = std::fs::metadata(&bin_path).ok()?;
642            if meta.len() > max_bytes {
643                tracing::warn!(
644                    "[bm25] index too large ({:.1} GB, limit {:.0} MB), quarantining: {}",
645                    meta.len() as f64 / 1_073_741_824.0,
646                    max_bytes / (1024 * 1024),
647                    bin_path.display()
648                );
649                let quarantined = bin_path.with_extension("bin.quarantined");
650                let _ = std::fs::rename(&bin_path, &quarantined);
651                return None;
652            }
653            let data = std::fs::read(&bin_path).ok()?;
654            let (idx, _): (Self, _) =
655                bincode::serde::decode_from_slice(&data, bincode::config::standard()).ok()?;
656            // Auto-migrate: compress legacy .bin to .bin.zst
657            if let Ok(compressed) = zstd::encode_all(data.as_slice(), ZSTD_LEVEL) {
658                let zst_tmp = zst_path.with_extension("zst.tmp");
659                if std::fs::write(&zst_tmp, &compressed).is_ok()
660                    && std::fs::rename(&zst_tmp, &zst_path).is_ok()
661                {
662                    tracing::info!(
663                        "[bm25] migrated {:.1} MB → {:.1} MB zstd",
664                        data.len() as f64 / 1_048_576.0,
665                        compressed.len() as f64 / 1_048_576.0
666                    );
667                    let _ = std::fs::remove_file(&bin_path);
668                }
669            }
670            return Some(idx);
671        }
672
673        let json_path = dir.join("bm25_index.json");
674        if json_path.exists() {
675            let meta = std::fs::metadata(&json_path).ok()?;
676            if meta.len() > max_bytes {
677                tracing::warn!(
678                    "[bm25] index too large ({:.1} GB, limit {:.0} MB), quarantining: {}",
679                    meta.len() as f64 / 1_073_741_824.0,
680                    max_bytes / (1024 * 1024),
681                    json_path.display()
682                );
683                let quarantined = json_path.with_extension("json.quarantined");
684                let _ = std::fs::rename(&json_path, &quarantined);
685                return None;
686            }
687            let data = std::fs::read_to_string(&json_path).ok()?;
688            return serde_json::from_str(&data).ok();
689        }
690
691        None
692    }
693
694    pub fn load_or_build(root: &Path) -> Self {
695        Self::load_or_build_inner(root, false)
696    }
697
698    /// Like `load_or_build` but uses a fast sentinel-sampling staleness check
699    /// that skips the expensive full directory walk for new-file detection.
700    pub fn load_or_build_fast(root: &Path) -> Self {
701        Self::load_or_build_inner(root, true)
702    }
703
704    fn load_or_build_inner(root: &Path, fast_stale: bool) -> Self {
705        if !is_safe_bm25_root(root) {
706            return Self::default();
707        }
708        if let Some(idx) = Self::load(root) {
709            let stale = if fast_stale {
710                bm25_index_looks_stale_fast(&idx, root)
711            } else {
712                bm25_index_looks_stale(&idx, root)
713            };
714            if !stale {
715                return idx;
716            }
717            tracing::debug!(
718                "[bm25_index: stale index detected for {}; rebuilding]",
719                root.display()
720            );
721            let rebuilt = if idx.files.is_empty() {
722                Self::build_from_directory(root)
723            } else {
724                Self::rebuild_incremental(root, &idx)
725            };
726            let _ = rebuilt.save(root);
727            return rebuilt;
728        }
729
730        let built = Self::build_from_directory(root);
731        let _ = built.save(root);
732        built
733    }
734
735    pub fn index_file_path(root: &Path) -> PathBuf {
736        let dir = index_dir(root);
737        let zst = dir.join("bm25_index.bin.zst");
738        if zst.exists() {
739            return zst;
740        }
741        let bin = dir.join("bm25_index.bin");
742        if bin.exists() {
743            return bin;
744        }
745        dir.join("bm25_index.json")
746    }
747
748    /// Ingest external `ContentChunk`s into the BM25 index.
749    /// Converts each chunk to a `CodeChunk` (backward-compatible) and
750    /// rebuilds the inverted index. Returns the number of chunks ingested.
751    pub fn ingest_content_chunks(
752        &mut self,
753        chunks: impl IntoIterator<Item = super::content_chunk::ContentChunk>,
754    ) -> usize {
755        let mut count = 0usize;
756        for cc in chunks {
757            self.add_chunk(cc.into());
758            count += 1;
759        }
760        if count > 0 {
761            self.finalize();
762        }
763        count
764    }
765
766    /// Number of chunks originating from external providers.
767    pub fn external_chunk_count(&self) -> usize {
768        self.chunks
769            .iter()
770            .filter(|c| c.file_path.contains("://"))
771            .count()
772    }
773}
774
775fn is_safe_bm25_root(root: &Path) -> bool {
776    super::graph_index::is_safe_scan_root_public(&root.to_string_lossy())
777}
778
779fn bm25_index_looks_stale(index: &BM25Index, root: &Path) -> bool {
780    bm25_index_looks_stale_inner(index, root, false)
781}
782
783/// Fast staleness check: samples a subset of tracked files and skips the
784/// expensive `list_code_files()` walk for new-file detection.
785pub fn bm25_index_looks_stale_fast(index: &BM25Index, root: &Path) -> bool {
786    bm25_index_looks_stale_inner(index, root, true)
787}
788
789fn bm25_index_looks_stale_inner(index: &BM25Index, root: &Path, fast: bool) -> bool {
790    if index.chunks.is_empty() {
791        return false;
792    }
793
794    if index.files.is_empty() {
795        let mut seen = std::collections::HashSet::<&str>::new();
796        for chunk in &index.chunks {
797            let rel = chunk.file_path.trim_start_matches(['/', '\\']);
798            if rel.is_empty() {
799                continue;
800            }
801            if !seen.insert(rel) {
802                continue;
803            }
804            if !root.join(rel).exists() {
805                return true;
806            }
807        }
808        return false;
809    }
810
811    if fast {
812        let sample_size = index.files.len().min(SENTINEL_SAMPLE_SIZE);
813        let step = if index.files.len() > sample_size {
814            index.files.len() / sample_size
815        } else {
816            1
817        };
818        for (i, (rel, old_state)) in index.files.iter().enumerate() {
819            if i % step != 0 {
820                continue;
821            }
822            let abs = root.join(rel);
823            if !abs.exists() {
824                return true;
825            }
826            let Some(cur) = IndexedFileState::from_path(&abs) else {
827                return true;
828            };
829            if &cur != old_state {
830                return true;
831            }
832        }
833        return false;
834    }
835
836    for (rel, old_state) in &index.files {
837        let abs = root.join(rel);
838        if !abs.exists() {
839            return true;
840        }
841        let Some(cur) = IndexedFileState::from_path(&abs) else {
842            return true;
843        };
844        if &cur != old_state {
845            return true;
846        }
847    }
848
849    for rel in list_code_files(root) {
850        if !index.files.contains_key(&rel) {
851            return true;
852        }
853    }
854
855    false
856}
857
858const SENTINEL_SAMPLE_SIZE: usize = 10;
859
860fn bounded_zstd_decode(compressed: &[u8], max_bytes: u64) -> Option<Vec<u8>> {
861    use std::io::Read;
862    let mut decoder = zstd::Decoder::new(compressed).ok()?;
863    let mut buf = Vec::new();
864    let mut chunk = vec![0u8; 65536];
865    let mut total = 0u64;
866    loop {
867        let n = decoder.read(&mut chunk).ok()?;
868        if n == 0 {
869            break;
870        }
871        total += n as u64;
872        if total > max_bytes {
873            tracing::warn!(
874                "[bm25] decompressed index exceeds limit ({:.0} MB > {:.0} MB), aborting load",
875                total as f64 / (1024.0 * 1024.0),
876                max_bytes as f64 / (1024.0 * 1024.0)
877            );
878            return None;
879        }
880        buf.extend_from_slice(&chunk[..n]);
881    }
882    Some(buf)
883}
884
885fn index_dir(root: &Path) -> PathBuf {
886    crate::core::index_namespace::vectors_dir(root)
887}
888
889fn list_code_files(root: &Path) -> Vec<String> {
890    let walker = ignore::WalkBuilder::new(root)
891        .hidden(true)
892        .git_ignore(true)
893        .git_global(true)
894        .git_exclude(true)
895        .max_depth(Some(20))
896        .filter_entry(crate::core::cloud_files::keep_entry)
897        .build();
898
899    let cfg = crate::core::config::Config::load();
900    let mut ignore_patterns: Vec<glob::Pattern> = DEFAULT_BM25_IGNORES
901        .iter()
902        .filter_map(|p| glob::Pattern::new(p).ok())
903        .collect();
904    ignore_patterns.extend(
905        cfg.extra_ignore_patterns
906            .iter()
907            .filter_map(|p| glob::Pattern::new(p).ok()),
908    );
909
910    let mut files: Vec<String> = Vec::new();
911    for entry in walker.flatten() {
912        let path = entry.path();
913        if !path.is_file() {
914            continue;
915        }
916        if !crate::core::ingestion::is_ingestible(path) {
917            continue;
918        }
919        let rel = path
920            .strip_prefix(root)
921            .unwrap_or(path)
922            .to_string_lossy()
923            .to_string();
924        if rel.is_empty() {
925            continue;
926        }
927        if ignore_patterns.iter().any(|p| p.matches(&rel)) {
928            continue;
929        }
930        if files.len() >= MAX_BM25_FILES {
931            tracing::warn!(
932                "[bm25] file cap reached ({MAX_BM25_FILES}), skipping remaining files in {}",
933                root.display()
934            );
935            break;
936        }
937        files.push(rel);
938    }
939
940    files.sort();
941    files.dedup();
942    files
943}
944
945pub fn is_code_file(path: &Path) -> bool {
946    let ext = path
947        .extension()
948        .and_then(|e| e.to_str())
949        .unwrap_or("")
950        .to_lowercase();
951    matches!(
952        ext.as_str(),
953        "rs" | "ts"
954            | "tsx"
955            | "js"
956            | "jsx"
957            | "py"
958            | "go"
959            | "java"
960            | "c"
961            | "cc"
962            | "cpp"
963            | "h"
964            | "hpp"
965            | "rb"
966            | "cs"
967            | "kt"
968            | "swift"
969            | "php"
970            | "scala"
971            | "sql"
972            | "ex"
973            | "exs"
974            | "zig"
975            | "lua"
976            | "dart"
977            | "vue"
978            | "svelte"
979    )
980}