lean_ctx/tools/ctx_semantic_search/
bm25_store.rs

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

lean_ctx/tools/ctx_semantic_search/
bm25_store.rs
