scryer-mcp 0.3.0

//! Tool input schemas as agents see them, and argument coercion against them.
//!
//! `schemars` output is correct JSON Schema but not what every MCP client accepts: optional
//! fields become `type: ["string", "null"]`, enums arrive as `$ref`s into `$defs`, and
//! integers carry `format: "uint"`. Gemini-style function declarations and strict OpenAI
//! validators reject or mangle those. [`simplify_schema`] reduces a schema to the portable
//! subset (plain `type`, inline `enum`, no `$ref`/`$defs`/`format`), and [`coerce_arguments`]
//! makes the server forgiving in the other direction: agents routinely send numbers and
//! booleans as strings.

use std::collections::HashMap;
use std::sync::OnceLock;

use serde_json::{Map, Value};

/// Parameters every tool accepts. Their descriptions are defined once here so they stay
/// identical (and short); [`publish_schema`] decides which tools advertise them.
const MAX_TOKENS_DESC: &str = "Token budget for this response (0 disables).";
const NO_TRUNCATE_DESC: &str = "If true, ignores the token budget.";
const PROJECT_DESC: &str = "Project slug, ID or root path. Default: the session's active project (from the working directory). An unknown value is an error (see list_projects).";

/// The only tools that publish `max_tokens` and `no_truncate`: the ones whose output can
/// outgrow the default budget. Every tool still *accepts* them (and coerces them), and the
/// truncation notice names them, so an agent that learns them from the notice can use them.
const BUDGET_PARAM_TOOLS: &[&str] = &[
    "find_references",
    "get_file_outline",
    "search_symbols",
    "trace_call_hierarchy",
    "get_crate_outline",
];

/// The only tools that publish `project`: those that manage projects. Every other tool acts
/// on the session's project (see `scryer serve --project`) but still accepts the argument.
const PROJECT_PARAM_TOOLS: &[&str] = &[
    "list_projects",
    "switch_active_project",
    "index_workspace",
    "get_indexing_status",
    "delete_project",
];

/// Reduce a `schemars` schema to the portable subset described in the module docs.
pub fn simplify_schema(tool_name: &str, schema: &mut Map<String, Value>) {
    let defs = schema
        .remove("$defs")
        .or_else(|| schema.remove("definitions"))
        .and_then(|v| match v {
            Value::Object(m) => Some(m),
            _ => None,
        })
        .unwrap_or_default();

    schema.remove("$schema");
    schema.remove("title");
    // The schema never allows extra keys a client could use, so don't advertise either way.
    schema.remove("additionalProperties");

    if let Some(Value::Object(props)) = schema.get_mut("properties") {
        for (key, prop) in props.iter_mut() {
            simplify_node(prop, &defs);
            match key.as_str() {
                "max_tokens" => set_description(prop, MAX_TOKENS_DESC),
                "no_truncate" => set_description(prop, NO_TRUNCATE_DESC),
                // These three describe `project` themselves (required, or with special values).
                "project"
                    if !matches!(
                        tool_name,
                        "switch_active_project" | "delete_project" | "get_indexing_status"
                    ) =>
                {
                    set_description(prop, PROJECT_DESC)
                }
                _ => {}
            }
        }
    }
}

/// Reduce an accepted (simplified) schema to what the tool advertises: the shared
/// parameters appear only on the tools that need them. Dispatch and coercion keep using the
/// accepted schema, so a call that still sends a hidden parameter works.
pub fn publish_schema(tool_name: &str, schema: &mut Map<String, Value>) {
    // `format` (text or json) is accepted by the list tools but not advertised: JSON is the
    // default, and a parameter named `format` costs every list tool about 100 characters. Pass it
    // per call, or set it for a session with `scryer serve --format`.
    let mut hidden: Vec<&str> = vec!["format"];
    if !BUDGET_PARAM_TOOLS.contains(&tool_name) {
        hidden.extend(["max_tokens", "no_truncate"]);
    }
    if !PROJECT_PARAM_TOOLS.contains(&tool_name) {
        hidden.push("project");
    }
    if let Some(Value::Object(props)) = schema.get_mut("properties") {
        for key in &hidden {
            props.remove(*key);
        }
    }
    if let Some(Value::Array(required)) = schema.get_mut("required") {
        required.retain(|r| r.as_str().is_none_or(|r| !hidden.contains(&r)));
    }
}

fn set_description(node: &mut Value, text: &str) {
    if let Value::Object(map) = node {
        map.insert("description".into(), Value::String(text.to_string()));
    }
}

/// Simplify one property (or array item) schema in place.
fn simplify_node(node: &mut Value, defs: &Map<String, Value>) {
    let Value::Object(map) = node else { return };

    // Inline `$ref`, keeping a description written next to the reference.
    if let Some(Value::String(reference)) = map.remove("$ref")
        && let Some(def) = reference
            .rsplit('/')
            .next()
            .and_then(|name| defs.get(name))
            .and_then(|d| d.as_object())
    {
        let description = map.remove("description");
        for (k, v) in def {
            map.entry(k.clone()).or_insert_with(|| v.clone());
        }
        if let Some(d) = description {
            map.insert("description".into(), d);
        }
    }

    // `anyOf: [X, {type: null}]` (an `Option<X>`) becomes X.
    for combinator in ["anyOf", "oneOf"] {
        if let Some(Value::Array(variants)) = map.get(combinator) {
            let non_null: Vec<Value> = variants
                .iter()
                .filter(|v| v.get("type").and_then(|t| t.as_str()) != Some("null"))
                .cloned()
                .collect();
            // An enum written as `oneOf: [{const: "a"}, {const: "b"}]` becomes `enum`.
            let consts: Option<Vec<Value>> =
                non_null.iter().map(|v| v.get("const").cloned()).collect();
            if let Some(values) = consts.filter(|c| !c.is_empty()) {
                map.remove(combinator);
                map.insert("type".into(), Value::String("string".into()));
                map.insert("enum".into(), Value::Array(values));
            } else if non_null.len() == 1 {
                map.remove(combinator);
                if let Some(Value::Object(only)) = non_null.into_iter().next() {
                    let description = map.remove("description");
                    for (k, v) in only {
                        map.entry(k).or_insert(v);
                    }
                    if let Some(d) = description {
                        map.insert("description".into(), d);
                    }
                }
                // The merged node may itself be a `$ref`.
                simplify_node(node, defs);
                return;
            }
        }
    }

    // `type: ["string", "null"]` becomes `type: "string"`.
    if let Some(Value::Array(types)) = map.get("type") {
        let first = types.iter().find(|t| t.as_str() != Some("null")).cloned();
        if let Some(t) = first {
            map.insert("type".into(), t);
        }
    }

    map.remove("format");
    map.remove("minimum");
    map.remove("title");

    if let Some(items) = map.get_mut("items") {
        simplify_node(items, defs);
    }
}

/// Schemas by tool name, as the server accepts them (every parameter, whether or not it is
/// advertised).
fn accepted_schemas() -> &'static HashMap<String, Map<String, Value>> {
    static SCHEMAS: OnceLock<HashMap<String, Map<String, Value>>> = OnceLock::new();
    SCHEMAS.get_or_init(|| {
        super::accepted_tool_definitions()
            .into_iter()
            .map(|t| (t.name.to_string(), (*t.input_schema).clone()))
            .collect()
    })
}

/// Make `arguments` match the tool's accepted schema where a client was sloppy:
/// - `null` values are dropped (an absent optional is the same thing);
/// - numeric strings become integers for `integer` parameters (`"12"` → `12`);
/// - `"true"` / `"false"` become booleans for `boolean` parameters;
/// - a bare string becomes a one-element array for `array` parameters.
pub fn coerce_arguments(tool_name: &str, arguments: &mut Value) {
    let Value::Object(args) = arguments else {
        return;
    };
    let Some(props) = accepted_schemas()
        .get(tool_name)
        .and_then(|s| s.get("properties"))
        .and_then(|p| p.as_object())
    else {
        return;
    };

    args.retain(|_, v| !v.is_null());
    for (key, value) in args.iter_mut() {
        let Some(kind) = props
            .get(key)
            .and_then(|p| p.get("type"))
            .and_then(|t| t.as_str())
        else {
            continue;
        };
        match (kind, &*value) {
            ("integer", Value::String(s)) => {
                if let Ok(n) = s.trim().parse::<u64>() {
                    *value = Value::from(n);
                }
            }
            ("boolean", Value::String(s)) => match s.trim().to_ascii_lowercase().as_str() {
                "true" => *value = Value::Bool(true),
                "false" => *value = Value::Bool(false),
                _ => {}
            },
            ("array", Value::String(s)) => *value = Value::Array(vec![Value::String(s.clone())]),
            _ => {}
        }
    }
}

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

    fn simplified(schema: Value, tool: &str) -> Value {
        let Value::Object(mut map) = schema else {
            panic!()
        };
        simplify_schema(tool, &mut map);
        Value::Object(map)
    }

    #[test]
    fn nullable_types_formats_and_metadata_are_flattened() {
        let out = simplified(
            json!({
                "$schema": "https://json-schema.org/draft/2020-12/schema",
                "title": "P",
                "type": "object",
                "properties": {
                    "a": { "type": ["string", "null"], "description": "A" },
                    "n": { "type": ["integer", "null"], "format": "uint", "minimum": 0 },
                    "tags": { "type": ["array", "null"], "items": { "type": "string" } }
                },
                "required": ["a"]
            }),
            "inspect_symbol",
        );
        assert_eq!(
            out["properties"]["a"],
            json!({ "type": "string", "description": "A" })
        );
        assert_eq!(out["properties"]["n"], json!({ "type": "integer" }));
        assert_eq!(out["properties"]["tags"]["type"], "array");
        assert!(out.get("$schema").is_none() && out.get("title").is_none());
        assert_eq!(out["required"], json!(["a"]));
    }

    #[test]
    fn optional_enum_refs_become_inline_enums() {
        let out = simplified(
            json!({
                "type": "object",
                "properties": {
                    "scope": {
                        "description": "Where to look.",
                        "anyOf": [{ "$ref": "#/$defs/Scope" }, { "type": "null" }]
                    }
                },
                "$defs": { "Scope": { "type": "string", "enum": ["project", "file"] } }
            }),
            "find_references",
        );
        assert!(out.get("$defs").is_none());
        let scope = &out["properties"]["scope"];
        assert_eq!(scope["type"], "string");
        assert_eq!(scope["enum"], json!(["project", "file"]));
        assert_eq!(scope["description"], "Where to look.");
        assert!(scope.get("anyOf").is_none() && scope.get("$ref").is_none());
    }

    #[test]
    fn documented_enums_written_as_one_of_consts_become_enum() {
        let out = simplified(
            json!({
                "type": "object",
                "properties": {
                    "dir": {
                        "anyOf": [{ "$ref": "#/$defs/Dir" }, { "type": "null" }]
                    }
                },
                "$defs": { "Dir": { "oneOf": [
                    { "type": "string", "const": "inbound", "description": "callers" },
                    { "type": "string", "const": "outbound", "description": "callees" }
                ] } }
            }),
            "trace_call_hierarchy",
        );
        let dir = &out["properties"]["dir"];
        assert_eq!(dir["type"], "string");
        assert_eq!(dir["enum"], json!(["inbound", "outbound"]));
    }

    #[test]
    fn shared_parameters_are_described_once_and_published_only_where_needed() {
        let schema = json!({
            "type": "object",
            "properties": {
                "max_tokens": { "type": ["integer", "null"], "description": "x" },
                "no_truncate": { "type": ["boolean", "null"], "description": "y" },
                "project": { "type": ["string", "null"], "description": "z" }
            }
        });
        // Accepting is independent of publishing: the accepted schema keeps everything.
        let accepted = simplified(schema, "find_references");
        assert_eq!(
            accepted["properties"]["max_tokens"]["description"],
            MAX_TOKENS_DESC
        );
        assert_eq!(
            accepted["properties"]["no_truncate"]["description"],
            NO_TRUNCATE_DESC
        );
        assert_eq!(
            accepted["properties"]["project"]["description"],
            PROJECT_DESC
        );

        let published = |tool: &str| {
            let Value::Object(mut map) = accepted.clone() else {
                panic!()
            };
            publish_schema(tool, &mut map);
            Value::Object(map)
        };
        let budget_tool = published("find_references");
        assert!(budget_tool["properties"].get("max_tokens").is_some());
        assert!(budget_tool["properties"].get("no_truncate").is_some());
        assert!(budget_tool["properties"].get("project").is_none());

        let admin_tool = published("delete_project");
        assert!(admin_tool["properties"].get("max_tokens").is_none());
        assert!(admin_tool["properties"].get("project").is_some());

        let plain = published("inspect_symbol");
        assert!(plain["properties"].as_object().unwrap().is_empty());
    }

    #[test]
    fn arguments_are_coerced_to_the_published_types() {
        let mut args = json!({
            "file_path": "src/lib.rs",
            "line": "3",
            "col": " 7 ",
            "compact": "TRUE",
            "kinds": "fn",
            "limit": null,
            "unknown": "kept"
        });
        coerce_arguments("get_file_outline", &mut args);
        assert_eq!(args["file_path"], "src/lib.rs");
        assert_eq!(args["kinds"], json!(["fn"]));
        assert_eq!(args["compact"], true);
        assert!(args.get("limit").is_none());
        assert_eq!(args["unknown"], "kept");

        let mut args = json!({ "file_path": "a.rs", "line": "3", "col": " 7 " });
        coerce_arguments("resolve_definition", &mut args);
        assert_eq!(args["line"], 3);
        assert_eq!(args["col"], 7);
    }

    #[test]
    fn unparsable_values_are_left_for_validation_to_report() {
        let mut args = json!({ "file_path": "a.rs", "line": "three" });
        coerce_arguments("get_enclosing_scope", &mut args);
        assert_eq!(args["line"], "three");
    }

    #[test]
    fn every_published_schema_is_portable() {
        for tool in super::super::all_tool_definitions() {
            let json = serde_json::to_string(&*tool.input_schema).unwrap();
            for banned in [
                "\"null\"",
                "$ref",
                "$defs",
                "$schema",
                "\"format\"",
                "\"minimum\"",
                "anyOf",
                "oneOf",
            ] {
                assert!(
                    !json.contains(banned),
                    "{} schema contains {banned}: {json}",
                    tool.name
                );
            }
        }
    }
}