scryer-mcp 0.2.1

use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;
use std::path::{Path, PathBuf};

use rmcp::model::Tool;
use scryer_db::{DependencyPackage, Project, ProjectDependency, SourceFile, Symbol};
use scryer_engine::EngineService;

use super::admin::{make_tool, read_only};
use crate::context::ProjectContextResolver;

// --- Input & Output Types ---

#[derive(Debug, Clone, Deserialize, JsonSchema)]
pub struct GetCrateOutlineParams {
    /// Name of a crate the project depends on, exactly as in Cargo.toml (e.g. "axum", "serde").
    pub crate_name: String,
    /// Version to pick when several are indexed (e.g. "1.0.200"). Default: the one the project uses.
    pub version: Option<String>,
    /// Optional submodule path to narrow the outline (e.g., 'extract' in axum).
    pub module_path: Option<String>,
    /// Optional project slug or path context.
    pub project: Option<String>,
    /// Custom token budget for this response.
    pub max_tokens: Option<usize>,
    /// If true, bypasses token budget truncation.
    pub no_truncate: Option<bool>,
}

#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
pub struct CrateOutlineResult {
    pub found: bool,
    pub crate_name: String,
    pub version: Option<String>,
    pub modules: Vec<String>,
    pub core_types: Vec<String>,
    pub formatted_markdown: String,
    /// Modules and core types left out to fit the token budget.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub omitted: Option<usize>,
}

// --- Tool Handlers ---

/// A symbol definition matched by name, with where its source lives.
#[derive(Debug, Clone)]
pub(crate) struct ResolvedSymbol {
    pub symbol: Symbol,
    /// Project-relative path (workspace) or package-relative path (dependency).
    pub rel_path: Option<String>,
    /// Absolute path of the defining file on disk.
    pub abs_path: Option<PathBuf>,
    /// BLAKE3 hash of the file contents at index time.
    pub content_hash: Option<String>,
    /// The dependency package for symbols from the project 0 cache.
    pub package: Option<DependencyPackage>,
}

impl ResolvedSymbol {
    pub(crate) fn is_external(&self) -> bool {
        self.symbol.project_id == 0
    }

    /// Path shown to agents: project-relative for workspace symbols, absolute for
    /// dependency symbols (they live in the Cargo cache, outside the project).
    pub(crate) fn display_path(&self) -> String {
        if self.is_external()
            && let Some(abs) = &self.abs_path
        {
            return abs.to_string_lossy().into_owned();
        }
        self.rel_path
            .clone()
            .unwrap_or_else(|| "unknown".to_string())
    }
}

/// Symbol kinds that are types (they have members rather than callers).
pub(crate) const TYPE_KINDS: [&str; 6] = ["struct", "enum", "trait", "type", "class", "interface"];

/// One definition offered when a name is ambiguous.
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
pub struct SymbolCandidate {
    pub qualified_name: String,
    pub kind: String,
    /// `path:line` of the definition.
    pub location: String,
}

impl ResolvedSymbol {
    pub(crate) fn candidate(&self) -> SymbolCandidate {
        SymbolCandidate {
            qualified_name: self.symbol.qualified_name.clone(),
            kind: self.symbol.kind.clone(),
            location: format!("{}:{}", self.display_path(), self.symbol.start_line),
        }
    }

    /// Whether this definition lives in `hint` (exact, or either path a suffix of the other
    /// so both relative and absolute hints work).
    pub(crate) fn in_file(&self, hint: &str) -> bool {
        let hint = hint.trim_start_matches("./");
        if hint.is_empty() {
            return false;
        }
        let rel = self.rel_path.as_deref().unwrap_or("");
        let abs = self
            .abs_path
            .as_ref()
            .map(|p| p.to_string_lossy().into_owned())
            .unwrap_or_default();
        (!rel.is_empty() && (rel == hint || hint.ends_with(&format!("/{rel}"))))
            || abs == hint
            || rel.ends_with(&format!("/{hint}"))
    }
}

/// The candidates defined in any of `hints`, or `None` when none are.
pub(crate) fn narrow_by_file(
    candidates: &[ResolvedSymbol],
    hints: &[&str],
) -> Option<Vec<ResolvedSymbol>> {
    let hinted: Vec<ResolvedSymbol> = candidates
        .iter()
        .filter(|c| hints.iter().any(|h| c.in_file(h)))
        .cloned()
        .collect();
    (!hinted.is_empty()).then_some(hinted)
}

/// Crate name a qualified name starts with (`tokio::sync::Mutex` → `tokio`).
pub(crate) fn crate_root(qualified_name: &str) -> &str {
    qualified_name.split("::").next().unwrap_or(qualified_name)
}

/// Full path a `use` stub imports, from its signature: `use tokio::sync::` named `Mutex`
/// (one item of a grouped import) → `tokio::sync::Mutex`. `None` for non-stubs.
pub(crate) fn import_path(stub: &Symbol) -> Option<String> {
    if stub.kind != "reexport" {
        return None;
    }
    let sig = stub.signature.trim().trim_end_matches(';');
    let sig = sig
        .strip_prefix("pub(crate) ")
        .or_else(|| sig.strip_prefix("pub "))
        .unwrap_or(sig);
    let path = sig.strip_prefix("use ")?.trim();
    let path = path.split(" as ").next().unwrap_or(path);
    let path = path.split('{').next().unwrap_or(path).trim();
    if path.ends_with("::") {
        Some(format!("{path}{}", stub.name))
    } else {
        Some(path.to_string())
    }
}

/// The definition a `use` stub imports, when the index has it: a workspace item
/// (`use crate::engine::step` → `…::engine::step`) or a dependency item, including public
/// paths (`use tokio::sync::Mutex` → `tokio::sync::mutex::Mutex`).
pub(crate) async fn follow_import(
    db: &mut toasty::Db,
    project: &Project,
    stub: &Symbol,
) -> anyhow::Result<Option<Symbol>> {
    let Some(path) = import_path(stub) else {
        return Ok(None);
    };
    // `crate::` / `self::` / `super::` are relative to the importing module; drop them and
    // match the rest as a qualified-name suffix.
    let segments: Vec<&str> = path
        .trim_start_matches("::")
        .split("::")
        .skip_while(|s| matches!(*s, "crate" | "self" | "super"))
        .collect();
    if segments.is_empty() {
        return Ok(None);
    }
    let candidates = lookup_symbol_candidates(db, project, &segments.join("::")).await?;
    Ok(candidates
        .into_iter()
        .map(|c| c.symbol)
        .find(|s| s.id != stub.id))
}

/// Whether workspace `use` stub `stub` imports dependency definition `def`. Public paths often
/// differ from definition paths (`use tokio::sync::Mutex` imports
/// `tokio::sync::mutex::Mutex`), so this matches on crate and item name.
pub(crate) fn imports_symbol(stub: &Symbol, def: &Symbol) -> bool {
    import_path(stub).is_some_and(|path| {
        crate_root(&path) == crate_root(&def.qualified_name) && last_segment(&path) == def.name
    })
}

/// Last path segment of a `::`- or `.`-qualified name.
pub(crate) fn last_segment(symbol: &str) -> &str {
    let after_colons = symbol.rsplit("::").next().unwrap_or(symbol);
    after_colons.rsplit('.').next().unwrap_or(after_colons)
}

/// Whether `sym` matches the user-supplied `query` (bare or qualified name).
fn name_matches(sym: &Symbol, query: &str) -> bool {
    if query.contains("::") || query.contains('.') {
        sym.qualified_name == query
            || sym.qualified_name.ends_with(&format!("::{query}"))
            || sym.qualified_name.ends_with(&format!(".{query}"))
    } else {
        sym.name == query
    }
}

/// Every definition of `query` in `project`, falling back to the project 0 dependency cache
/// when the workspace has none. Dependency candidates of crates the project depends on sort
/// first. Re-export stubs are skipped.
pub(crate) async fn lookup_symbol_candidates(
    db: &mut toasty::Db,
    project: &Project,
    query: &str,
) -> anyhow::Result<Vec<ResolvedSymbol>> {
    let target_name = last_segment(query);
    let root = Path::new(&project.root_path);

    let ws_symbols = Symbol::filter(
        Symbol::fields()
            .project_id()
            .eq(project.id)
            .and(Symbol::fields().name().eq(target_name)),
    )
    .exec(&mut *db)
    .await?;

    let mut out = Vec::new();
    for sym in ws_symbols
        .into_iter()
        .filter(|s| s.kind != "reexport" && name_matches(s, query))
    {
        let file = SourceFile::filter(
            SourceFile::fields()
                .project_id()
                .eq(project.id)
                .and(SourceFile::fields().id().eq(sym.file_id)),
        )
        .first()
        .exec(&mut *db)
        .await?;
        out.push(ResolvedSymbol {
            abs_path: file.as_ref().map(|f| root.join(&f.path)),
            rel_path: file.as_ref().map(|f| f.path.clone()),
            content_hash: file.map(|f| f.content_hash),
            package: None,
            symbol: sym,
        });
    }
    if !out.is_empty() {
        return Ok(out);
    }

    // Dependency cache, scoped by the crate prefix when the query is qualified.
    let linked_pkg_ids: Vec<u64> =
        ProjectDependency::filter(ProjectDependency::fields().project_id().eq(project.id))
            .exec(&mut *db)
            .await?
            .into_iter()
            .map(|l| l.dependency_package_id)
            .collect();

    let root_crate = query
        .contains("::")
        .then(|| query.split("::").next())
        .flatten();
    let target_pkg = match root_crate {
        Some(crate_name) => {
            DependencyPackage::filter(DependencyPackage::fields().name().eq(crate_name))
                .exec(&mut *db)
                .await?
                .into_iter()
                .max_by_key(|p| linked_pkg_ids.contains(&p.id))
        }
        None => None,
    };

    let by_name = Symbol::fields()
        .project_id()
        .eq(0)
        .and(Symbol::fields().name().eq(target_name));
    let dep_symbols = match &target_pkg {
        Some(pkg) => {
            Symbol::filter(by_name.and(Symbol::fields().dependency_package_id().eq(Some(pkg.id))))
                .exec(&mut *db)
                .await?
        }
        None => Symbol::filter(by_name).exec(&mut *db).await?,
    };

    let definitions: Vec<Symbol> = dep_symbols
        .into_iter()
        .filter(|s| s.kind != "reexport")
        .collect();
    let mut matches: Vec<Symbol> = definitions
        .iter()
        .filter(|s| name_matches(s, query))
        .cloned()
        .collect();
    // Public paths often differ from definition paths (`tokio::sync::Mutex` is defined at
    // `tokio::sync::mutex::Mutex` and re-exported), so within a named crate fall back to
    // the bare name.
    if matches.is_empty() && target_pkg.is_some() {
        matches = definitions;
    }
    matches.sort_by_key(|s| {
        !s.dependency_package_id
            .is_some_and(|id| linked_pkg_ids.contains(&id))
    });

    for sym in matches {
        let package = match sym.dependency_package_id {
            Some(pkg_id) => {
                DependencyPackage::filter(DependencyPackage::fields().id().eq(pkg_id))
                    .first()
                    .exec(&mut *db)
                    .await?
            }
            None => None,
        };
        let file = SourceFile::filter(
            SourceFile::fields()
                .project_id()
                .eq(0)
                .and(SourceFile::fields().id().eq(sym.file_id)),
        )
        .first()
        .exec(&mut *db)
        .await?;
        let abs_path = match (&package, &file) {
            (Some(pkg), Some(f)) => Some(Path::new(&pkg.root_path).join(&f.path)),
            (None, Some(f)) => Some(PathBuf::from(&f.path)),
            _ => None,
        };
        out.push(ResolvedSymbol {
            abs_path,
            rel_path: file.as_ref().map(|f| f.path.clone()),
            content_hash: file.map(|f| f.content_hash),
            package,
            symbol: sym,
        });
    }
    Ok(out)
}

pub async fn handle_get_crate_outline(
    context: &ProjectContextResolver,
    engine: &EngineService,
    params: GetCrateOutlineParams,
) -> anyhow::Result<CrateOutlineResult> {
    let mut guard = engine.db().lock().await;

    let (project, _) = context
        .resolve_project(None, params.project.as_deref())
        .await?;

    let dep_links =
        ProjectDependency::filter(ProjectDependency::fields().project_id().eq(project.id))
            .exec(&mut *guard)
            .await?;
    let linked_pkg_ids: Vec<u64> = dep_links
        .into_iter()
        .map(|l| l.dependency_package_id)
        .collect();

    let all_pkgs =
        DependencyPackage::filter(DependencyPackage::fields().name().eq(&params.crate_name))
            .exec(&mut *guard)
            .await?;

    let matching_pkg = all_pkgs.into_iter().find(|p| {
        if let Some(v) = &params.version {
            p.version == *v
        } else {
            linked_pkg_ids.contains(&p.id) || true
        }
    });

    let Some(pkg) = matching_pkg else {
        return Ok(CrateOutlineResult {
            found: false,
            crate_name: params.crate_name.clone(),
            version: params.version,
            modules: Vec::new(),
            core_types: Vec::new(),
            formatted_markdown: format!(
                "Crate '{0}' is not indexed in the dependency cache. Only crates the project's Cargo.toml depends on are indexed (run index_workspace after adding one); check the spelling with search_symbols(query: \"{0}\", scope_level: \"dependencies\").",
                params.crate_name
            ),
            omitted: None,
        });
    };

    let symbols = Symbol::filter(
        Symbol::fields()
            .project_id()
            .eq(0)
            .and(Symbol::fields().dependency_package_id().eq(Some(pkg.id))),
    )
    .exec(&mut *guard)
    .await?;

    let mut module_items: BTreeMap<String, Vec<String>> = BTreeMap::new();
    let mut core_types = Vec::new();

    let module_filter = params.module_path.as_deref();

    for s in &symbols {
        if let Some(mf) = module_filter {
            let matches_mod = s.qualified_name.contains(&format!("::{mf}::"))
                || s.qualified_name.ends_with(&format!("::{mf}"))
                || s.qualified_name.starts_with(&format!("{}::{mf}", pkg.name));
            if !matches_mod {
                continue;
            }
        }

        // Group into modules
        if let Some((mod_prefix, _)) = s.qualified_name.rsplit_once("::") {
            let entry = module_items.entry(mod_prefix.to_string()).or_default();
            if entry.len() < 5 && !entry.contains(&s.name) {
                entry.push(s.name.clone());
            }
        }

        // Identify core types
        if matches!(
            s.kind.as_str(),
            "struct" | "enum" | "trait" | "type" | "interface"
        ) {
            let doc_summary = s
                .docstring
                .as_deref()
                .and_then(|d| d.lines().next())
                .unwrap_or("")
                .trim();
            let summary_str = if doc_summary.is_empty() {
                String::new()
            } else {
                format!(": {}", doc_summary.trim_end_matches('.'))
            };
            core_types.push(format!("`{} {}`{}", s.kind, s.qualified_name, summary_str));
        }
    }

    let mut modules_vec = Vec::new();
    for (mod_path, items) in module_items {
        if items.is_empty() {
            modules_vec.push(format!("`{mod_path}`"));
        } else {
            modules_vec.push(format!("`{mod_path}` ({})", items.join(", ")));
        }
    }

    let md = crate_outline_markdown(&pkg.name, &pkg.version, &modules_vec, &core_types);

    Ok(CrateOutlineResult {
        found: true,
        crate_name: pkg.name,
        version: Some(pkg.version),
        modules: modules_vec,
        core_types,
        formatted_markdown: md,
        omitted: None,
    })
}

fn crate_outline_markdown(
    name: &str,
    version: &str,
    modules: &[String],
    core_types: &[String],
) -> String {
    let mut md = format!("# Crate: {name} ({version})\n\n");

    if !modules.is_empty() {
        md.push_str("## Modules\n");
        for m in modules {
            md.push_str(&format!("- {m}\n"));
        }
        md.push('\n');
    }

    if !core_types.is_empty() {
        md.push_str("## Core Types\n");
        for t in core_types {
            md.push_str(&format!("- {t}\n"));
        }
    }
    md
}

/// Shrink a crate outline so the whole response fits the token budget.
///
/// The outline carries every module and type twice (as lists and in `formatted_markdown`)
/// and the markdown is one JSON line, so a big crate would otherwise be cut mid-list by the
/// line-based payload limiter. This keeps the first entries of each list that fit, rebuilds
/// the markdown from them and reports how many were left out; `module_path` narrows to the
/// rest. `count_tokens` is the limiter's tokenizer.
pub fn fit_crate_outline_to_budget(
    out: &mut CrateOutlineResult,
    max_tokens: Option<usize>,
    no_truncate: bool,
    count_tokens: impl Fn(&str) -> usize,
) {
    let budget = max_tokens.unwrap_or(crate::telemetry::MAX_PAYLOAD_TOKENS);
    if no_truncate || budget == 0 || !out.found {
        return;
    }
    let fits = |r: &CrateOutlineResult| {
        serde_json::to_string_pretty(r).is_ok_and(|json| count_tokens(&json) <= budget)
    };
    if fits(out) {
        return;
    }

    let total = out.modules.len() + out.core_types.len();
    let version = out.version.clone().unwrap_or_default();
    let with_share = |n: usize| {
        let mut r = out.clone();
        // Keep the same fraction of both lists so neither crowds out the other.
        let keep = |len: usize| (len * n).div_ceil(total.max(1)).min(len);
        r.modules.truncate(keep(out.modules.len()));
        r.core_types.truncate(keep(out.core_types.len()));
        r.formatted_markdown =
            crate_outline_markdown(&r.crate_name, &version, &r.modules, &r.core_types);
        r.omitted = Some(total - r.modules.len() - r.core_types.len());
        r
    };

    // Largest number of entries (out of `total`) that fits; zero entries is the floor.
    let (mut lo, mut hi) = (0, total);
    while lo < hi {
        let mid = (lo + hi).div_ceil(2);
        if fits(&with_share(mid)) {
            lo = mid;
        } else {
            hi = mid - 1;
        }
    }
    *out = with_share(lo);
}

// --- Tool Definitions ---

pub fn tool_definitions() -> Vec<Tool> {
    vec![make_tool::<GetCrateOutlineParams>(
        "get_crate_outline",
        "Use when exploring a third-party crate for the first time, deciding which module or item to import, or discovering available types and traits, without browsing the Cargo cache or web docs. Pass `crate_name` exactly as in Cargo.toml and narrow large crates with `module_path`; if the crate isn't indexed the reply says how to find the right name.",
        read_only(),
    )]
}

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

    fn outline(modules: usize, types: usize) -> CrateOutlineResult {
        let modules: Vec<String> = (0..modules)
            .map(|i| format!("`krate::m{i}` (a, b)"))
            .collect();
        let core_types: Vec<String> = (0..types)
            .map(|i| format!("`struct krate::T{i}`: does a thing"))
            .collect();
        let md = crate_outline_markdown("krate", "1.0.0", &modules, &core_types);
        CrateOutlineResult {
            found: true,
            crate_name: "krate".into(),
            version: Some("1.0.0".into()),
            modules,
            core_types,
            formatted_markdown: md,
            omitted: None,
        }
    }

    /// Roughly one token per four characters.
    fn tokens(s: &str) -> usize {
        s.len().div_ceil(4)
    }

    #[test]
    fn small_outline_is_untouched() {
        let mut out = outline(2, 3);
        let before = serde_json::to_string(&out).unwrap();
        fit_crate_outline_to_budget(&mut out, None, false, tokens);
        assert_eq!(serde_json::to_string(&out).unwrap(), before);
        assert!(out.omitted.is_none());
    }

    #[test]
    fn large_outline_shrinks_to_fit_with_consistent_markdown() {
        let mut out = outline(200, 200);
        fit_crate_outline_to_budget(&mut out, Some(600), false, tokens);

        let json = serde_json::to_string_pretty(&out).unwrap();
        assert!(tokens(&json) <= 600, "{} tokens", tokens(&json));
        let omitted = out.omitted.expect("entries were left out");
        assert_eq!(omitted, 400 - out.modules.len() - out.core_types.len());
        assert!(!out.modules.is_empty() && !out.core_types.is_empty());
        // The markdown lists exactly what is kept.
        assert_eq!(
            out.formatted_markdown,
            crate_outline_markdown("krate", "1.0.0", &out.modules, &out.core_types)
        );
    }

    #[test]
    fn no_truncate_and_zero_budget_leave_it_alone() {
        for (max, no_truncate) in [(Some(10), true), (Some(0), false)] {
            let mut out = outline(200, 200);
            fit_crate_outline_to_budget(&mut out, max, no_truncate, tokens);
            assert_eq!(out.modules.len(), 200);
            assert!(out.omitted.is_none());
        }
    }
}