lean_ctx/core/
bm25_store.rs

1// SPDX-License-Identifier: Apache-2.0
2//! BM25 index lifecycle: per-thread shared cache, load-or-refresh with the
3//! cold-build budget, resident-cache storage.
4
5use std::path::Path;
6
7use crate::core::bm25_index::BM25Index;
8
9std::thread_local! {
10    static BM25_SHARED_CACHE: std::cell::RefCell<Option<crate::core::bm25_cache::SharedBm25Cache>> =
11        const { std::cell::RefCell::new(None) };
12}
13
14/// Set the shared BM25 cache for the current thread (called from the registered handler).
15pub fn set_thread_cache(cache: crate::core::bm25_cache::SharedBm25Cache) {
16    BM25_SHARED_CACHE.with(|c| {
17        *c.borrow_mut() = Some(cache);
18    });
19}
20
21/// Clone the current thread's shared BM25 cache, if any. Lets composer tools
22/// propagate the resident cache into a budgeted worker thread so a slow cold
23/// build warms the *same* cache instead of being wasted work.
24pub fn get_thread_cache() -> Option<crate::core::bm25_cache::SharedBm25Cache> {
25    BM25_SHARED_CACHE.with(|c| c.borrow().clone())
26}
27
28/// Result of BM25 index loading — may indicate background build in progress.
29pub enum Bm25LoadResult {
30    Ready(std::sync::Arc<BM25Index>),
31    Building,
32}
33
34pub fn load_or_refresh_bm25(root: &Path) -> Bm25LoadResult {
35    let cached = BM25_SHARED_CACHE.with(|c| {
36        let borrow = c.borrow();
37        borrow
38            .as_ref()
39            .and_then(|cache| crate::core::bm25_cache::get_or_background(cache, root))
40    });
41    if let Some(idx) = cached {
42        return Bm25LoadResult::Ready(idx);
43    }
44
45    let root_str = root.to_string_lossy().to_string();
46
47    // #1724: the resident cache holds ONE root, so every call for a sibling
48    // project (`ctx_compose(path=...)`) misses it and lands here. Loading the
49    // persisted index unconditionally pinned that project's first, partial tree
50    // forever — nothing on this path ever asked whether it was still current.
51    // Check staleness (one stat per tracked directory, no ingest walk); a stale
52    // index still answers this call, but it now also triggers the refresh that
53    // makes the next call complete, and it is not pinned into the resident
54    // cache so the rebuilt index is picked up as soon as it is persisted.
55    if let Some(idx) = crate::core::index_orchestrator::try_load_bm25_index(&root_str) {
56        if crate::core::bm25_index::bm25_index_looks_stale_fast(&idx, root) {
57            tracing::debug!("[bm25_store: stale on-disk index for {root_str}; refreshing]");
58            crate::core::index_orchestrator::ensure_all_background(&root_str);
59            return Bm25LoadResult::Ready(std::sync::Arc::new(idx));
60        }
61        let idx = std::sync::Arc::new(idx);
62        store_in_thread_cache(root, &idx);
63        return Bm25LoadResult::Ready(idx);
64    }
65
66    if crate::core::index_orchestrator::is_building() {
67        return Bm25LoadResult::Building;
68    }
69
70    // Cold path: kick off the background build (which persists the index to
71    // disk) instead of doing an unbounded synchronous build in the MCP handler.
72    // Wait briefly so small/medium repos still return Ready on the first call;
73    // larger repos return Building and the agent retries against the warm cache
74    // once the worker has persisted the index (#150).
75    crate::core::index_orchestrator::ensure_all_background(&root_str);
76
77    let deadline = std::time::Instant::now() + bm25_cold_build_budget();
78    loop {
79        if let Some(idx) = crate::core::index_orchestrator::try_load_bm25_index(&root_str) {
80            let idx = std::sync::Arc::new(idx);
81            store_in_thread_cache(root, &idx);
82            return Bm25LoadResult::Ready(idx);
83        }
84        if std::time::Instant::now() >= deadline {
85            return Bm25LoadResult::Building;
86        }
87        std::thread::sleep(std::time::Duration::from_millis(50));
88    }
89}
90
91/// Time budget for waiting on a cold BM25 build in the MCP handler before
92/// returning `Building`. Overridable via `LEAN_CTX_BM25_COLD_BUDGET_MS`.
93pub(crate) fn bm25_cold_build_budget() -> std::time::Duration {
94    let ms = std::env::var("LEAN_CTX_BM25_COLD_BUDGET_MS")
95        .ok()
96        .and_then(|v| v.parse::<u64>().ok())
97        .unwrap_or(60_000);
98    std::time::Duration::from_millis(ms)
99}
100
101pub(crate) fn store_in_thread_cache(root: &Path, idx: &std::sync::Arc<BM25Index>) {
102    BM25_SHARED_CACHE.with(|c| {
103        let borrow = c.borrow();
104        if let Some(cache) = borrow.as_ref() {
105            let mut guard = cache
106                .lock()
107                .unwrap_or_else(std::sync::PoisonError::into_inner);
108            *guard = Some(crate::core::bm25_cache::Bm25CacheEntry {
109                root: root.to_path_buf(),
110                index: std::sync::Arc::clone(idx),
111                loaded_at: std::time::Instant::now(),
112                fingerprint: crate::core::bm25_cache::index_fingerprint(root),
113            });
114        }
115    });
116}
117
118pub fn filtered_candidate_k(top_k: usize, filtered: bool) -> usize {
119    if !filtered {
120        return top_k;
121    }
122    let candidates = (top_k.max(10)).saturating_mul(10);
123    candidates.clamp(50, 500)
124}
lean_ctx/core/bm25_store.rs

lean_ctx/core/
bm25_store.rs
