lean_ctx/core/

index_orchestrator.rs

use std::collections::HashMap;
use std::path::Path;
use std::sync::{Arc, Mutex, OnceLock};
use std::time::{SystemTime, UNIX_EPOCH};

use serde::Serialize;

use crate::core::bm25_index::BM25Index;
use crate::core::graph_index::{self, ProjectIndex};

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum State {
    Idle,
    Building,
    Ready,
    Failed,
}

#[derive(Debug, Clone)]
struct Component {
    state: State,
    started_ms: Option<u64>,
    finished_ms: Option<u64>,
    duration_ms: Option<u64>,
    last_error: Option<String>,
    /// Human-readable outcome detail surfaced to operators (e.g. doc count +
    /// persisted size, or the "not persisted: too large …" remedy). Independent
    /// of `last_error` so a *successful* build can still carry a warning note.
    note: Option<String>,
}

impl Component {
    fn new() -> Self {
        Self {
            state: State::Idle,
            started_ms: None,
            finished_ms: None,
            duration_ms: None,
            last_error: None,
            note: None,
        }
    }
}

#[derive(Debug)]
struct ProjectBuild {
    worker_running: bool,
    /// Set the first time a heavy-index tool lazily pre-warms this root (#152).
    /// Prevents re-triggering a full rebuild on every subsequent dispatch — the
    /// tools' own `load_or_build` paths handle staleness from then on.
    warm_triggered: bool,
    graph: Component,
    bm25: Component,
}

impl ProjectBuild {
    fn new() -> Self {
        Self {
            worker_running: false,
            warm_triggered: false,
            graph: Component::new(),
            bm25: Component::new(),
        }
    }
}

// Lock ordering (see rust/LOCK_ORDERING.md):
//   L1 = REGISTRY outer Mutex  (the HashMap guard)
//   L2 = per-project Arc<Mutex<ProjectBuild>>  (inner guard)
//
// Invariant: L1 must NEVER be held while locking L2.
// `entry_for()` enforces this by cloning the Arc and dropping L1 before
// the caller acquires L2.
static REGISTRY: OnceLock<Mutex<HashMap<String, Arc<Mutex<ProjectBuild>>>>> = OnceLock::new();

fn registry() -> &'static Mutex<HashMap<String, Arc<Mutex<ProjectBuild>>>> {
    REGISTRY.get_or_init(|| Mutex::new(HashMap::new()))
}

fn entry_for(project_root: &str) -> Arc<Mutex<ProjectBuild>> {
    let mut map = registry()
        .lock()
        .unwrap_or_else(std::sync::PoisonError::into_inner);
    map.entry(project_root.to_string())
        .or_insert_with(|| Arc::new(Mutex::new(ProjectBuild::new())))
        .clone()
}

fn now_ms() -> u64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default()
        .as_millis() as u64
}

fn start_component(c: &mut Component) {
    c.state = State::Building;
    c.started_ms = Some(now_ms());
    c.finished_ms = None;
    c.duration_ms = None;
    c.last_error = None;
    c.note = None;
}

fn finish_ok(c: &mut Component) {
    c.state = State::Ready;
    let end = now_ms();
    c.finished_ms = Some(end);
    c.duration_ms = c.started_ms.map(|s| end.saturating_sub(s));
}

fn finish_err(c: &mut Component, e: String) {
    c.state = State::Failed;
    let end = now_ms();
    c.finished_ms = Some(end);
    c.duration_ms = c.started_ms.map(|s| end.saturating_sub(s));
    c.last_error = Some(e);
}

/// The index warmth a tool benefits from. Drives lazy, demand-driven warming
/// (issue #152) so the server no longer scans the whole project eagerly on every
/// `initialize` — a session that only uses `ctx_read`/`ctx_shell`/`ctx_tree`
/// pays zero indexing cost.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum WarmNeed {
    /// No prebuilt index needed.
    None,
    /// Only the resident line-search (trigram) index — cheap, used by `ctx_search`.
    Search,
    /// Full project indices (graph + BM25; this also warms the search index).
    Heavy,
}

/// Classify a tool by the index warmth it benefits from. Unknown tools default
/// to [`WarmNeed::None`]; a heavy tool mis-classified as `None` still works — it
/// just builds its index synchronously on first use instead of being pre-warmed.
#[must_use]
pub fn warm_need_for_tool(tool: &str) -> WarmNeed {
    match tool {
        "ctx_search" => WarmNeed::Search,
        // Tools that build/consume the graph, call-graph, BM25 or artifact index.
        "ctx_graph"
        | "ctx_callgraph"
        | "ctx_routes"
        | "ctx_repomap"
        | "ctx_impact"
        | "ctx_artifacts"
        | "ctx_semantic_search"
        | "ctx_provider"
        | "ctx_compose"
        | "ctx_review" => WarmNeed::Heavy,
        _ => WarmNeed::None,
    }
}

/// Lazily warm the indices a tool needs, deduped per root. Never blocks (all
/// work is spawned in the background) and is safe to call on every dispatch.
///
/// Returns `true` only when this call is the *first* heavy pre-warm for `root`
/// in this process — the caller can use that signal to warm secondary roots once
/// without re-reading session state on every dispatch.
pub fn ensure_warm_for_tool(project_root: &str, tool: &str) -> bool {
    if project_root.is_empty() {
        return false;
    }
    match warm_need_for_tool(tool) {
        WarmNeed::None => false,
        WarmNeed::Search => {
            // The search index has its own TTL + background-rebuild dedup, so it
            // is safe (and cheap) to nudge on every `ctx_search`.
            crate::core::search_index::ensure_background(project_root, true, false);
            false
        }
        WarmNeed::Heavy => {
            let entry = entry_for(project_root);
            let first_warm = {
                let mut s = entry
                    .lock()
                    .unwrap_or_else(std::sync::PoisonError::into_inner);
                if s.warm_triggered {
                    false
                } else {
                    s.warm_triggered = true;
                    true
                }
            };
            if first_warm {
                ensure_all_background(project_root);
            }
            first_warm
        }
    }
}

pub fn ensure_all_background(project_root: &str) {
    let state = entry_for(project_root);
    let should_spawn = {
        let mut s = state
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        if s.worker_running {
            false
        } else {
            s.worker_running = true;
            true
        }
    };

    if !should_spawn {
        return;
    }

    let root = project_root.to_string();
    std::thread::spawn(move || {
        let state = entry_for(&root);

        // Pre-warm the resident line-search index in parallel (own thread,
        // deduped internally) so the first ctx_search hits the fast path.
        crate::core::search_index::ensure_background(&root, true, false);

        // Phase 1: Graph index — may produce a content cache from the file walk
        {
            let mut s = state
                .lock()
                .unwrap_or_else(std::sync::PoisonError::into_inner);
            start_component(&mut s.graph);
        }
        let graph_result = std::panic::catch_unwind(|| {
            let (idx, content_cache) = graph_index::scan_with_content_cache(&root);
            // JSON index write is kept for backward compatibility with remaining
            // direct ProjectIndex consumers. Will be removed when all consumers
            // are migrated to GraphProvider/PropertyGraph. (OPT-14/15 Phase 6)
            let _ = idx.save();
            (idx, content_cache)
        });
        let content_cache = if let Ok((_idx, cache)) = graph_result {
            let mut s = state
                .lock()
                .unwrap_or_else(std::sync::PoisonError::into_inner);
            finish_ok(&mut s.graph);
            cache
        } else {
            let mut s = state
                .lock()
                .unwrap_or_else(std::sync::PoisonError::into_inner);
            finish_err(&mut s.graph, "graph index build panicked".to_string());
            HashMap::new()
        };

        // Phase 2: BM25 index — reuses content from graph scan when available
        {
            let mut s = state
                .lock()
                .unwrap_or_else(std::sync::PoisonError::into_inner);
            start_component(&mut s.bm25);
        }
        let bm = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
            let root_pb = Path::new(&root);
            let idx = if content_cache.is_empty() {
                BM25Index::load_or_build(root_pb)
            } else {
                BM25Index::build_with_content_hint(root_pb, &content_cache)
            };
            let outcome = idx.save(root_pb);
            (idx.doc_count, outcome)
        }));
        if let Ok((doc_count, save_res)) = bm {
            let mut s = state
                .lock()
                .unwrap_or_else(std::sync::PoisonError::into_inner);
            finish_ok(&mut s.bm25);
            s.bm25.note = Some(bm25_build_note(doc_count, &save_res));
        } else {
            let mut s = state
                .lock()
                .unwrap_or_else(std::sync::PoisonError::into_inner);
            finish_err(&mut s.bm25, "bm25 build panicked".to_string());
        }

        let mut s = state
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        s.worker_running = false;
    });
}

/// Ensure background indexing for all extra roots (in addition to the primary).
/// Each extra root that is not a subdirectory of `primary_root` gets its own
/// graph + BM25 index. Capped at `MAX_EXTRA_ROOT_BUILDS` to prevent runaway.
const MAX_EXTRA_ROOT_BUILDS: usize = 8;

pub fn ensure_extra_roots_background(primary_root: &str, extra_roots: &[String]) {
    let primary = Path::new(primary_root);
    let mut built = 0;
    for root in extra_roots {
        if built >= MAX_EXTRA_ROOT_BUILDS {
            break;
        }
        let rp = Path::new(root);
        if !rp.is_dir() {
            continue;
        }
        // Skip if extra_root is inside primary (already indexed by the primary scan)
        if rp.starts_with(primary) {
            continue;
        }
        // Skip if primary is inside this extra_root (avoid double-indexing the parent)
        if primary.starts_with(rp) {
            continue;
        }
        ensure_all_background(root);
        built += 1;
    }
}

/// Build a human-readable outcome note for a finished BM25 build, including the
/// indexed chunk count and whether the index was persisted to disk. A
/// "too large" refusal carries the exact remedy so the operator (or agent) is
/// never left guessing why search/ranking stays cold (issue #249).
fn bm25_build_note(
    doc_count: usize,
    save: &std::io::Result<crate::core::bm25_index::SaveOutcome>,
) -> String {
    use crate::core::bm25_index::SaveOutcome;
    match save {
        Ok(SaveOutcome::Persisted { compressed_bytes }) => format!(
            "indexed {doc_count} chunks, {:.1} MB persisted",
            *compressed_bytes as f64 / 1_048_576.0
        ),
        Ok(SaveOutcome::SkippedTooLarge {
            compressed_bytes,
            limit_bytes,
        }) => format!(
            "indexed {doc_count} chunks but NOT persisted to disk: compressed {:.1} MB exceeds the {:.0} MB cap. \
             Raise it via LEAN_CTX_BM25_MAX_CACHE_MB (or bm25_max_cache_mb in config) or add extra_ignore_patterns, \
             then run `lean-ctx reindex`. Until then the index is rebuilt from scratch on every cold start.",
            *compressed_bytes as f64 / 1_048_576.0,
            *limit_bytes as f64 / 1_048_576.0
        ),
        Err(e) => format!("indexed {doc_count} chunks but persisting failed: {e}"),
    }
}

/// Lightweight, allocation-frugal snapshot of the BM25 component for the
/// in-call composer/search messaging. Avoids the heavier [`disk_status`] walk.
#[derive(Debug, Clone)]
pub struct Bm25Summary {
    pub state: &'static str,
    /// While building: elapsed so far. Otherwise: last build duration.
    pub elapsed_ms: Option<u64>,
    pub note: Option<String>,
    pub last_error: Option<String>,
}

pub fn bm25_summary(project_root: &str) -> Bm25Summary {
    let entry = entry_for(project_root);
    let s = entry
        .lock()
        .unwrap_or_else(std::sync::PoisonError::into_inner);
    let c = &s.bm25;
    let elapsed_ms = if matches!(c.state, State::Building) {
        c.started_ms.map(|start| now_ms().saturating_sub(start))
    } else {
        c.duration_ms
    };
    Bm25Summary {
        state: match c.state {
            State::Idle => "idle",
            State::Building => "building",
            State::Ready => "ready",
            State::Failed => "failed",
        },
        elapsed_ms,
        note: c.note.clone(),
        last_error: c.last_error.clone(),
    }
}

pub fn try_load_graph_index(project_root: &str) -> Option<ProjectIndex> {
    // Resident cache: avoids re-reading + zstd-decompressing + serde-parsing the
    // on-disk index on every graph-touching query. Returns an in-memory clone.
    crate::core::graph_cache::get_cached(project_root).map(|arc| (*arc).clone())
}

pub fn try_load_bm25_index(project_root: &str) -> Option<BM25Index> {
    BM25Index::load(Path::new(project_root))
}

/// Returns true if any project is currently building its indices.
pub fn is_building() -> bool {
    let map = registry()
        .lock()
        .unwrap_or_else(std::sync::PoisonError::into_inner);
    map.values().any(|entry| {
        let s = entry
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        matches!(s.bm25.state, State::Building) || matches!(s.graph.state, State::Building)
    })
}

#[derive(Debug, Serialize)]
struct ComponentStatus<'a> {
    state: &'a str,
    started_ms: Option<u64>,
    finished_ms: Option<u64>,
    duration_ms: Option<u64>,
    last_error: Option<&'a str>,
    #[serde(skip_serializing_if = "Option::is_none")]
    note: Option<&'a str>,
}

fn component_status(c: &Component) -> ComponentStatus<'_> {
    ComponentStatus {
        state: match c.state {
            State::Idle => "idle",
            State::Building => "building",
            State::Ready => "ready",
            State::Failed => "failed",
        },
        started_ms: c.started_ms,
        finished_ms: c.finished_ms,
        duration_ms: c.duration_ms,
        last_error: c.last_error.as_deref(),
        note: c.note.as_deref(),
    }
}

#[derive(Debug, Serialize)]
struct StatusResponse<'a> {
    project_root: &'a str,
    graph_index: ComponentStatus<'a>,
    bm25_index: ComponentStatus<'a>,
    disk: DiskStatusAll,
}

#[derive(Debug, Serialize, Default)]
pub struct DiskStatus {
    pub exists: bool,
    pub size_bytes: Option<u64>,
    pub file_count: Option<u64>,
    pub modified_at: Option<String>,
}

#[derive(Debug, Serialize, Default)]
pub struct DiskStatusAll {
    pub graph_index: DiskStatus,
    pub bm25_index: DiskStatus,
    pub code_graph: DiskStatus,
}

fn disk_status_for_graph(project_root: &str) -> DiskStatus {
    let Some(dir) = graph_index::ProjectIndex::index_dir(project_root) else {
        return DiskStatus::default();
    };
    let zst = dir.join("index.json.zst");
    let json = dir.join("index.json");
    let path = if zst.exists() {
        zst
    } else if json.exists() {
        json
    } else {
        return DiskStatus::default();
    };
    let meta = std::fs::metadata(&path).ok();
    let file_count =
        graph_index::ProjectIndex::load(project_root).map(|idx| idx.files.len() as u64);
    DiskStatus {
        exists: true,
        size_bytes: meta.as_ref().map(std::fs::Metadata::len),
        file_count,
        modified_at: meta.and_then(|m| m.modified().ok()).map(format_time),
    }
}

fn disk_status_for_bm25(project_root: &str) -> DiskStatus {
    let root = Path::new(project_root);
    let path = BM25Index::index_file_path(root);
    if !path.exists() {
        return DiskStatus::default();
    }
    let meta = std::fs::metadata(&path).ok();
    DiskStatus {
        exists: true,
        size_bytes: meta.as_ref().map(std::fs::Metadata::len),
        file_count: None,
        modified_at: meta.and_then(|m| m.modified().ok()).map(format_time),
    }
}

fn disk_status_for_code_graph(project_root: &str) -> DiskStatus {
    let dir = crate::core::property_graph::graph_dir(project_root);
    let db_path = dir.join("graph.db");
    if !db_path.exists() {
        return DiskStatus::default();
    }
    let meta = std::fs::metadata(&db_path).ok();
    let node_count = crate::core::property_graph::CodeGraph::open(project_root)
        .ok()
        .and_then(|g| {
            g.connection()
                .query_row("SELECT count(*) FROM nodes", [], |r| r.get::<_, i64>(0))
                .ok()
                .map(|c| c as u64)
        });
    DiskStatus {
        exists: true,
        size_bytes: meta.as_ref().map(std::fs::Metadata::len),
        file_count: node_count,
        modified_at: meta.and_then(|m| m.modified().ok()).map(format_time),
    }
}

fn format_time(t: SystemTime) -> String {
    let secs = t.duration_since(UNIX_EPOCH).unwrap_or_default().as_secs();
    let dt = chrono::DateTime::from_timestamp(secs as i64, 0);
    dt.map_or_else(
        || format!("{secs}"),
        |d| d.format("%Y-%m-%d %H:%M:%S UTC").to_string(),
    )
}

pub fn disk_status(project_root: &str) -> DiskStatusAll {
    DiskStatusAll {
        graph_index: disk_status_for_graph(project_root),
        bm25_index: disk_status_for_bm25(project_root),
        code_graph: disk_status_for_code_graph(project_root),
    }
}

pub fn status_json(project_root: &str) -> String {
    let state = entry_for(project_root);
    let s = state
        .lock()
        .unwrap_or_else(std::sync::PoisonError::into_inner);
    let res = StatusResponse {
        project_root,
        graph_index: component_status(&s.graph),
        bm25_index: component_status(&s.bm25),
        disk: disk_status(project_root),
    };
    serde_json::to_string(&res).unwrap_or_else(|_| "{}".to_string())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn status_json_is_valid_json() {
        let s = status_json("/tmp");
        let _: serde_json::Value = serde_json::from_str(&s).unwrap();
    }

    #[test]
    fn warm_need_classifies_tools() {
        // Lightweight tools must never trigger a project scan (#152).
        for light in [
            "ctx_read",
            "ctx_shell",
            "ctx_tree",
            "ctx_knowledge",
            "unknown_tool",
        ] {
            assert_eq!(warm_need_for_tool(light), WarmNeed::None, "{light}");
        }
        // ctx_search only needs the cheap trigram index.
        assert_eq!(warm_need_for_tool("ctx_search"), WarmNeed::Search);
        // Graph / BM25 consumers need the full warm.
        for heavy in [
            "ctx_graph",
            "ctx_callgraph",
            "ctx_routes",
            "ctx_repomap",
            "ctx_impact",
            "ctx_artifacts",
            "ctx_semantic_search",
            "ctx_provider",
            "ctx_compose",
            "ctx_review",
        ] {
            assert_eq!(warm_need_for_tool(heavy), WarmNeed::Heavy, "{heavy}");
        }
    }

    #[test]
    fn ensure_warm_lightweight_and_search_never_signal_first_warm() {
        // None and Search must return false (no heavy pre-warm), and an empty
        // root is always a no-op.
        assert!(!ensure_warm_for_tool("", "ctx_graph"));
        let tmp = tempfile::tempdir().unwrap();
        let root = tmp.path().to_string_lossy().to_string();
        assert!(!ensure_warm_for_tool(&root, "ctx_read"));
        assert!(!ensure_warm_for_tool(&root, "ctx_search"));
    }

    #[test]
    fn ensure_warm_heavy_is_once_per_root() {
        // The first heavy pre-warm signals `true` (so the caller warms extra
        // roots once); every subsequent call is a no-op `false`, preventing a
        // rebuild-on-every-dispatch storm.
        let tmp = tempfile::tempdir().unwrap();
        let root = tmp.path().to_string_lossy().to_string();
        assert!(
            ensure_warm_for_tool(&root, "ctx_callgraph"),
            "first heavy warm must signal true"
        );
        assert!(
            !ensure_warm_for_tool(&root, "ctx_callgraph"),
            "second heavy warm must be deduped to false"
        );
        assert!(
            !ensure_warm_for_tool(&root, "ctx_semantic_search"),
            "any later heavy tool on the same root is also deduped"
        );
    }

    #[test]
    fn build_note_persisted_reports_size() {
        let note = bm25_build_note(
            42,
            &Ok(crate::core::bm25_index::SaveOutcome::Persisted {
                compressed_bytes: 3 * 1024 * 1024,
            }),
        );
        assert!(
            note.contains("42 chunks"),
            "note should report chunk count: {note}"
        );
        assert!(
            note.contains("persisted"),
            "note should report persistence: {note}"
        );
    }

    #[test]
    fn build_note_too_large_carries_remedy() {
        let note = bm25_build_note(
            1000,
            &Ok(crate::core::bm25_index::SaveOutcome::SkippedTooLarge {
                compressed_bytes: 600 * 1024 * 1024,
                limit_bytes: 512 * 1024 * 1024,
            }),
        );
        assert!(
            note.contains("NOT persisted"),
            "must flag non-persistence: {note}"
        );
        assert!(
            note.contains("LEAN_CTX_BM25_MAX_CACHE_MB") && note.contains("reindex"),
            "too-large note must carry an actionable remedy: {note}"
        );
    }

    #[test]
    fn build_note_persist_error_is_reported() {
        let note = bm25_build_note(7, &Err(std::io::Error::other("disk full")));
        assert!(note.contains("persisting failed"), "note: {note}");
        assert!(
            note.contains("disk full"),
            "note should include the io error: {note}"
        );
    }

    #[test]
    fn bm25_summary_unknown_project_is_idle() {
        let tmp = tempfile::tempdir().unwrap();
        let summary = bm25_summary(tmp.path().to_string_lossy().as_ref());
        assert_eq!(summary.state, "idle");
        assert!(summary.note.is_none());
        assert!(summary.last_error.is_none());
    }

    #[test]
    fn extra_roots_skips_subdirs_of_primary() {
        let tmp = tempfile::tempdir().unwrap();
        let primary = tmp.path().join("primary");
        std::fs::create_dir_all(&primary).unwrap();
        let sub = primary.join("subdir");
        std::fs::create_dir_all(&sub).unwrap();
        let external = tmp.path().join("external");
        std::fs::create_dir_all(&external).unwrap();

        let primary_str = primary.to_string_lossy().to_string();
        let extra = vec![
            sub.to_string_lossy().to_string(),
            external.to_string_lossy().to_string(),
        ];

        // Should not panic; subdirs are skipped, external is attempted
        ensure_extra_roots_background(&primary_str, &extra);
    }

    #[test]
    fn extra_roots_caps_at_max() {
        let tmp = tempfile::tempdir().unwrap();
        let primary = tmp.path().join("primary");
        std::fs::create_dir_all(&primary).unwrap();

        let mut extra = Vec::new();
        for i in 0..20 {
            let d = tmp.path().join(format!("ext-{i}"));
            std::fs::create_dir_all(&d).unwrap();
            extra.push(d.to_string_lossy().to_string());
        }

        let primary_str = primary.to_string_lossy().to_string();
        // Should not spawn more than MAX_EXTRA_ROOT_BUILDS threads
        ensure_extra_roots_background(&primary_str, &extra);
    }
}