scryer-mcp 0.3.0

//! What a tool call does about the state of the index before and after it answers.
//!
//! Before: wait (bounded) for a first index that is still building, and refresh the files the
//! call names. After: refresh the files the answer names and, if any changed, answer again.
//! Whatever is still not trustworthy is said in the answer, never left as an empty list.

use std::collections::BTreeSet;
use std::path::{Path, PathBuf};
use std::time::Duration;

use scryer_db::Project;
use scryer_engine::IndexState;
use serde_json::Value;

use crate::context::ProjectContextResolver;

/// How long a query waits for a first index to finish before it answers with an explicit
/// `INDEXING` state. Seconds; the default is 20 (measured first indexes take 2 to 11 s).
pub const FIRST_QUERY_WAIT_ENV: &str = "SCRYER_FIRST_QUERY_WAIT";
const DEFAULT_FIRST_QUERY_WAIT: Duration = Duration::from_secs(20);

/// How long a query waits for the file watcher to finish a batch it already holds.
pub const WATCHER_DRAIN: Duration = Duration::from_millis(50);

/// Tools whose answers come from the index and so need it ready and fresh.
pub fn needs_index(tool: &str) -> bool {
    matches!(
        tool,
        "resolve_definition"
            | "find_references"
            | "get_enclosing_scope"
            | "get_file_outline"
            | "get_type_contract"
            | "trace_call_hierarchy"
            | "calculate_blast_radius"
            | "get_crate_outline"
            | "search_symbols"
            | "inspect_symbol"
    )
}

/// [`FIRST_QUERY_WAIT_ENV`] as a duration.
pub fn first_query_wait() -> Duration {
    std::env::var(FIRST_QUERY_WAIT_ENV)
        .ok()
        .and_then(|v| v.trim().parse::<f64>().ok())
        .filter(|secs| secs.is_finite() && *secs >= 0.0)
        .map_or(DEFAULT_FIRST_QUERY_WAIT, Duration::from_secs_f64)
}

/// The index state of the project a call targets, as the call saw it.
#[derive(Debug, Clone)]
pub struct CallIndex {
    /// The project the call resolved to, if it resolved to one.
    pub project: Option<Project>,
    /// The state after waiting. `Building` here means the wait ran out.
    pub state: IndexState,
    /// Files that were re-indexed to make this answer current.
    pub refreshed: BTreeSet<String>,
    /// Files that no longer exist and were dropped from the index.
    pub removed: BTreeSet<String>,
}

impl CallIndex {
    /// A call that does not depend on a project's index.
    pub fn none() -> Self {
        Self {
            project: None,
            state: IndexState::Ready,
            refreshed: BTreeSet::new(),
            removed: BTreeSet::new(),
        }
    }

    /// The root of the targeted project.
    pub fn root(&self) -> Option<PathBuf> {
        self.project.as_ref().map(|p| PathBuf::from(&p.root_path))
    }

    /// The reason the answer is not trustworthy, if it is not: the index is still building, or
    /// a file the answer names is waiting for a re-index. `touched` are those file paths.
    pub fn caveat(&self, touched: &BTreeSet<String>) -> Option<String> {
        match &self.state {
            IndexState::Ready => None,
            IndexState::Building { .. } => {
                let percent = self.state.percent().unwrap_or(0);
                let eta = self.state.eta_secs().unwrap_or(5);
                Some(format!(
                    "INDEXING {percent}%: results so far are partial; retry in about {eta} s"
                ))
            }
            IndexState::Stale { pending } => {
                let hit: Vec<&String> = touched.iter().filter(|f| pending.contains(*f)).collect();
                if hit.is_empty() {
                    return None;
                }
                let shown: Vec<&str> = hit.iter().take(5).map(|s| s.as_str()).collect();
                let more = hit.len().saturating_sub(shown.len());
                Some(format!(
                    "STALE: {} file(s) pending re-index ({}{}): results for them may be out of date; retry in a few seconds",
                    hit.len(),
                    shown.join(", "),
                    if more > 0 {
                        format!(", +{more} more")
                    } else {
                        String::new()
                    }
                ))
            }
        }
    }

    /// The freshness the answer reports: what was refreshed, or `fresh`.
    pub fn freshness(&self) -> String {
        let n = self.refreshed.len() + self.removed.len();
        if n == 0 {
            "fresh".to_string()
        } else if n == 1 {
            "refreshed 1 file".to_string()
        } else {
            format!("refreshed {n} files")
        }
    }

    /// The header's index line: the caveat when there is one, otherwise `facts` (the
    /// `index: N files, ...` line) with the freshness folded in.
    pub fn header_line(&self, facts: Option<&str>, touched: &BTreeSet<String>) -> Option<String> {
        if let Some(caveat) = self.caveat(touched) {
            return Some(caveat);
        }
        let freshness = self.freshness();
        match facts {
            Some(line) => Some(match line.strip_prefix("index: ") {
                Some(rest) => format!("index: {freshness}, {rest}"),
                None => format!("{line} | index: {freshness}"),
            }),
            None => (self.project.is_some()).then(|| format!("index: {freshness}")),
        }
    }

    /// Fold this call's index state into a JSON answer: drop what the partial index makes
    /// misleading (`complete`, "not found" notes) and add the caveat or the refreshed files.
    pub fn annotate_json(&self, value: &mut Value, touched: &BTreeSet<String>) {
        let Some(object) = value.as_object_mut() else {
            return;
        };
        let caveat = self.caveat(touched);
        let mut add = Vec::new();
        if let Some(caveat) = &caveat {
            object.remove("complete");
            add.push(caveat.clone());
        }
        if !self.refreshed.is_empty() || !self.removed.is_empty() {
            let files: Vec<&str> = self
                .refreshed
                .iter()
                .chain(&self.removed)
                .take(5)
                .map(|s| s.as_str())
                .collect();
            add.push(format!("{}: {}", self.freshness(), files.join(", ")));
        }
        let notes = object
            .entry("notes")
            .or_insert_with(|| Value::Array(Vec::new()));
        let Some(notes) = notes.as_array_mut() else {
            return;
        };
        if caveat.is_some() {
            notes.retain(|n| !n.as_str().is_some_and(is_not_found_note));
        }
        notes.extend(add.into_iter().map(Value::String));
        if notes.is_empty() {
            object.remove("notes");
        }
    }

    /// Drop the "not found" notes of a typed result whose answer is partial.
    pub fn scrub_notes(notes: &mut Vec<String>) {
        notes.retain(|n| !is_not_found_note(n));
    }
}

/// A note that says a name or location was not found in the index. Wrong whenever the index is
/// not complete, because the thing may simply not be indexed yet.
pub fn is_not_found_note(note: &str) -> bool {
    const PREFIXES: [&str; 4] = [
        "no indexed symbol named",
        "no symbol or scope contains",
        "no symbol contains",
        "no definition named",
    ];
    PREFIXES.iter().any(|p| note.starts_with(p))
}

/// Every file path a result names: each `file_path` string and each `location` (`path:line`).
pub fn collect_files(value: &Value, out: &mut BTreeSet<String>) {
    match value {
        Value::Object(map) => {
            for (key, v) in map {
                match (key.as_str(), v) {
                    ("file_path", Value::String(path)) if !path.is_empty() => {
                        out.insert(path.clone());
                    }
                    ("location", Value::String(loc)) => {
                        out.insert(strip_line(loc).to_string());
                    }
                    // `name — path:line`, as `inspect_symbol` lists 1-hop callers and callees.
                    ("callers" | "callees", Value::Array(items)) => {
                        for item in items.iter().filter_map(Value::as_str) {
                            if let Some((_, loc)) = item.rsplit_once(" — ") {
                                out.insert(strip_line(loc).to_string());
                            }
                        }
                        // `find_references` lists callers as objects with their own `file_path`.
                        collect_files(v, out);
                    }
                    _ => collect_files(v, out),
                }
            }
        }
        Value::Array(items) => items.iter().for_each(|v| collect_files(v, out)),
        _ => {}
    }
}

/// `path:12` or `path:12:3` to `path`.
fn strip_line(location: &str) -> &str {
    let mut path = location;
    for _ in 0..2 {
        match path.rsplit_once(':') {
            Some((head, tail)) if !tail.is_empty() && tail.bytes().all(|b| b.is_ascii_digit()) => {
                path = head
            }
            _ => break,
        }
    }
    path
}

/// A file path as a path relative to the project `root`: relative paths are taken as they
/// are, absolute ones must lie under the root (dependency sources do not, and are skipped).
pub fn relative_to_root(root: &Path, path: &str) -> Option<PathBuf> {
    let p = Path::new(path);
    if p.is_absolute() {
        p.strip_prefix(root).ok().map(Path::to_path_buf)
    } else {
        Some(p.to_path_buf())
    }
}

/// Wait for the project a call targets and bring the files it names up to date.
///
/// Registers the session's working directory first when no project matches and it is a
/// project root, so the first call in a fresh checkout works without any setup.
pub async fn prepare_call(
    context: &ProjectContextResolver,
    engine: &scryer_engine::EngineService,
    arguments: &Value,
    project_arg: Option<&str>,
    wait: Duration,
) -> CallIndex {
    let file_arg = arguments.get("file_path").and_then(Value::as_str);
    let resolve = || async {
        context
            .resolve_project(file_arg.map(Path::new), project_arg)
            .await
    };
    let resolved = match resolve().await {
        Ok(r) => Some(r),
        Err(_) => match context.session_cwd().await {
            Some(cwd)
                if crate::registration::ensure_cwd_registered(context.registry(), engine, &cwd)
                    .await
                    .ok()
                    .flatten()
                    .is_some() =>
            {
                resolve().await.ok()
            }
            _ => None,
        },
    };
    let Some((project, rel)) = resolved else {
        return CallIndex::none();
    };

    let mut state = engine.index_state(project.id);
    if state.is_building() {
        state = engine.wait_until_ready(project.id, wait).await;
    }
    engine.wait_watcher_idle(WATCHER_DRAIN).await;

    let mut call = CallIndex {
        project: Some(project),
        state,
        refreshed: BTreeSet::new(),
        removed: BTreeSet::new(),
    };
    if let Some(rel) = rel {
        refresh(engine, &mut call, &[rel]).await;
    }
    call
}

/// Refresh `rel_paths` (project-relative) and record what changed in `call`. Returns whether
/// the index changed.
pub async fn refresh(
    engine: &scryer_engine::EngineService,
    call: &mut CallIndex,
    rel_paths: &[PathBuf],
) -> bool {
    let (Some(project), Some(root)) = (call.project.as_ref(), call.root()) else {
        return false;
    };
    if rel_paths.is_empty() {
        return false;
    }
    match engine.refresh_files(project.id, &root, rel_paths).await {
        Ok(report) => {
            let changed = report.changed();
            call.refreshed.extend(report.refreshed);
            call.removed.extend(report.removed);
            changed
        }
        Err(e) => {
            tracing::warn!("Could not refresh files before answering: {e:#}");
            false
        }
    }
}

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

    fn building(done: usize, total: usize) -> CallIndex {
        CallIndex {
            state: IndexState::Building {
                done,
                total,
                elapsed: Duration::from_secs(2),
            },
            ..CallIndex::none()
        }
    }

    #[test]
    fn files_are_collected_from_paths_and_locations() {
        let value = json!({
            "file_path": "src/a.rs",
            "references": [{"file_path": "src/b.rs", "line": 3}],
            "definitions": [{"location": "src/c.rs:12"}, {"location": "src/d.rs:4:9"}],
            "callers": ["run — src/e.rs:7", {"file_path": "src/f.rs"}],
            "notes": ["file_path: not a key"],
        });
        let mut files = BTreeSet::new();
        collect_files(&value, &mut files);
        let files: Vec<_> = files.into_iter().collect();
        assert_eq!(
            files,
            [
                "src/a.rs", "src/b.rs", "src/c.rs", "src/d.rs", "src/e.rs", "src/f.rs"
            ]
        );
    }

    #[test]
    fn a_building_index_says_indexing_and_drops_not_found_notes() {
        let call = building(30, 100);
        let touched = BTreeSet::new();
        let caveat = call.caveat(&touched).unwrap();
        assert!(caveat.starts_with("INDEXING 30%: results so far are partial"));

        let mut value = json!({
            "complete": true,
            "notes": ["no indexed symbol named 'X'; this is not the same as \"no usages\"", "keep me"],
        });
        call.annotate_json(&mut value, &touched);
        assert!(value.get("complete").is_none());
        let notes: Vec<&str> = value["notes"]
            .as_array()
            .unwrap()
            .iter()
            .filter_map(Value::as_str)
            .collect();
        assert_eq!(notes.len(), 2);
        assert_eq!(notes[0], "keep me");
        assert!(notes[1].starts_with("INDEXING"));
    }

    #[test]
    fn a_stale_index_only_matters_for_the_files_it_has_pending() {
        let call = CallIndex {
            state: IndexState::Stale {
                pending: ["src/a.rs".to_string()].into(),
            },
            ..CallIndex::none()
        };
        let other: BTreeSet<String> = ["src/z.rs".to_string()].into();
        assert!(call.caveat(&other).is_none());
        let hit: BTreeSet<String> = ["src/a.rs".to_string(), "src/z.rs".to_string()].into();
        let caveat = call.caveat(&hit).unwrap();
        assert!(caveat.starts_with("STALE: 1 file(s) pending re-index (src/a.rs)"));
    }

    #[test]
    fn the_header_line_reports_freshness() {
        let mut call = CallIndex::none();
        let touched = BTreeSet::new();
        let facts = "index: 12 files, 0 parse failures, indexed T";
        assert_eq!(
            call.header_line(Some(facts), &touched).unwrap(),
            "index: fresh, 12 files, 0 parse failures, indexed T"
        );
        call.refreshed.insert("src/a.rs".to_string());
        call.refreshed.insert("src/b.rs".to_string());
        assert_eq!(
            call.header_line(Some(facts), &touched).unwrap(),
            "index: refreshed 2 files, 12 files, 0 parse failures, indexed T"
        );
    }

    #[test]
    fn relative_paths_stay_and_absolute_ones_must_be_inside_the_root() {
        let root = Path::new("/repo");
        assert_eq!(
            relative_to_root(root, "src/a.rs"),
            Some(PathBuf::from("src/a.rs"))
        );
        assert_eq!(
            relative_to_root(root, "/repo/src/a.rs"),
            Some(PathBuf::from("src/a.rs"))
        );
        assert_eq!(relative_to_root(root, "/cargo/registry/x.rs"), None);
    }

    #[test]
    fn the_first_query_wait_defaults_to_twenty_seconds() {
        assert_eq!(DEFAULT_FIRST_QUERY_WAIT, Duration::from_secs(20));
    }
}