emperor-mcp 0.4.0

//! Lean tool definitions: strip JSON Schema boilerplate a model pays for on every request.
//!
//! A `tools/list` answer is read by the model on every turn. Schemas generated by `schemars`
//! (what rmcp's `#[tool]` uses) carry boilerplate that costs tokens and says nothing a model
//! needs: a `$schema` URI on every tool, `"type": ["integer", "null"]` and `"default": null`
//! on every optional argument, Rust integer formats (`"format": "uint32"`, `"minimum": 0`),
//! and `$defs` indirection for types used once.
//!
//! [`compact`] rewrites an input schema in place:
//!
//! - drops `$schema`;
//! - an optional property (not in `required`) loses its `null` alternative:
//!   `"type": ["integer", "null"]` becomes `"type": "integer"`, a `{"type": "null"}` branch of
//!   `anyOf`/`oneOf` goes (a single remaining branch is merged into the property), `null` leaves
//!   an `enum` — leaving the argument out already means "none";
//! - drops `"default": null`;
//! - drops Rust numeric formats (`int8` … `int128`, `int`, `uint8` … `uint128`, `uint`,
//!   `float`, `double`) and the `"minimum": 0` an unsigned format brings with it; the ranges of
//!   8- and 16-bit types (`"maximum": 255`) and every other bound are kept;
//! - drops an empty `title` and a `title` that only repeats the property name;
//! - inlines a `$ref` to `$defs` used once, and drops `$defs` nothing refers to;
//! - recurses into `properties`, `items`, `anyOf`/`oneOf`/`allOf` and the other subschema
//!   keywords.
//!
//! Descriptions, `enum`, `const`, `required`, real bounds, non-null defaults,
//! `additionalProperties` and every other keyword are kept. The result is still a valid
//! JSON Schema 2020-12 document. The only input it stops describing is an explicit `null` for
//! an optional argument, which a server's `Option<T>` argument (with `#[serde(default)]`)
//! treats exactly like a missing one; rmcp deserialises arguments with serde, not by
//! validating the schema, so what the server accepts does not change.
//!
//! [`compact_output`] applies the same rules except the `null` ones to an output schema, where
//! a server may well send `null`. [`compact_tool`] does both for one rmcp [`Tool`].
//!
//! A server opts in with one line where it builds its router:
//!
//! ```ignore
//! use emperor_mcp::schema::CompactSchemas;
//!
//! MyServer { tool_router: Self::tool_router().compact_schemas() }
//! ```

use std::collections::{HashMap, HashSet};
use std::sync::Arc;

use rmcp::handler::server::router::tool::ToolRouter;
use rmcp::model::{JsonObject, ListToolsResult, Tool};
use serde_json::{Map, Value};

/// Keywords that annotate a schema without changing what it accepts. On a merge, the outer
/// schema's annotation wins (a property's own description over its type's).
const ANNOTATIONS: &[&str] = &[
    "title",
    "description",
    "default",
    "examples",
    "deprecated",
    "readOnly",
    "writeOnly",
    "$comment",
];

/// Keywords whose value is one subschema.
const ONE_SUBSCHEMA: &[&str] = &[
    "items",
    "additionalProperties",
    "unevaluatedProperties",
    "unevaluatedItems",
    "contains",
    "propertyNames",
    "not",
    "if",
    "then",
    "else",
];

/// Keywords whose value is an array of subschemas.
const SUBSCHEMA_ARRAYS: &[&str] = &["anyOf", "oneOf", "allOf", "prefixItems"];

/// Keywords whose value maps names to subschemas (`properties` is handled on its own).
const SUBSCHEMA_MAPS: &[&str] = &["patternProperties", "dependentSchemas", "$defs"];

/// Compact a tool's input schema in place (see the [module docs](self) for the rules).
pub fn compact(schema: &mut JsonObject) {
    compact_with(schema, true);
}

/// Compact an output schema in place: the rules of [`compact`] except those about `null`,
/// which an output may contain.
pub fn compact_output(schema: &mut JsonObject) {
    compact_with(schema, false);
}

/// Compact one rmcp [`Tool`]: [`compact`] on its input schema, [`compact_output`] on its output
/// schema.
pub fn compact_tool(tool: &mut Tool) {
    compact(Arc::make_mut(&mut tool.input_schema));
    if let Some(output) = tool.output_schema.as_mut() {
        compact_output(Arc::make_mut(output));
    }
}

/// [`compact_tool`] on every tool of a list (e.g. `ToolRouter::list_all()`).
pub fn compact_tools(tools: &mut [Tool]) {
    tools.iter_mut().for_each(compact_tool);
}

/// [`compact_tool`] on every tool of a `tools/list` result.
pub fn compact_list(result: &mut ListToolsResult) {
    compact_tools(&mut result.tools);
}

/// [`compact_tool`] on every route of an rmcp tool router, so `tools/list` answers with the
/// compact definitions. Disabled routes are compacted too, ready for when they are enabled.
pub fn compact_router<S>(router: &mut ToolRouter<S>) {
    for route in router.map.values_mut() {
        compact_tool(&mut route.attr);
    }
}

/// One-line opt-in: `Self::tool_router().compact_schemas()`.
pub trait CompactSchemas {
    /// Return `self` with every tool definition compacted ([`compact_tool`]).
    #[must_use]
    fn compact_schemas(self) -> Self;
}

impl<S> CompactSchemas for ToolRouter<S> {
    fn compact_schemas(mut self) -> Self {
        compact_router(&mut self);
        self
    }
}

impl CompactSchemas for Vec<Tool> {
    fn compact_schemas(mut self) -> Self {
        compact_tools(&mut self);
        self
    }
}

impl CompactSchemas for ListToolsResult {
    fn compact_schemas(mut self) -> Self {
        compact_list(&mut self);
        self
    }
}

fn compact_with(schema: &mut JsonObject, strip_null: bool) {
    schema.remove("$schema");
    inline_defs(schema);
    walk(schema, strip_null);
}

/// Compact one schema object and everything below it.
fn walk(node: &mut Map<String, Value>, strip_null: bool) {
    if node.get("default").is_some_and(Value::is_null) {
        node.remove("default");
    }
    if node
        .get("title")
        .and_then(Value::as_str)
        .is_some_and(|t| t.trim().is_empty())
    {
        node.remove("title");
    }
    drop_numeric_format(node);

    let required: HashSet<String> = node
        .get("required")
        .and_then(Value::as_array)
        .map(|r| {
            r.iter()
                .filter_map(Value::as_str)
                .map(String::from)
                .collect()
        })
        .unwrap_or_default();
    if let Some(Value::Object(props)) = node.get_mut("properties") {
        for (name, prop) in props.iter_mut() {
            let Value::Object(prop) = prop else { continue };
            if prop
                .get("title")
                .and_then(Value::as_str)
                .is_some_and(|t| same_name(t, name))
            {
                prop.remove("title");
            }
            if strip_null && !required.contains(name) {
                drop_null(prop);
            }
            walk(prop, strip_null);
        }
    }
    for key in ONE_SUBSCHEMA {
        if let Some(Value::Object(sub)) = node.get_mut(*key) {
            walk(sub, strip_null);
        }
    }
    for key in SUBSCHEMA_ARRAYS {
        if let Some(Value::Array(subs)) = node.get_mut(*key) {
            for sub in subs.iter_mut().filter_map(Value::as_object_mut) {
                walk(sub, strip_null);
            }
        }
    }
    for key in SUBSCHEMA_MAPS {
        if let Some(Value::Object(subs)) = node.get_mut(*key) {
            for sub in subs.values_mut().filter_map(Value::as_object_mut) {
                walk(sub, strip_null);
            }
        }
    }
}

/// A title that only repeats the property name (`"max_hits"` / `"Max Hits"` / `"maxHits"`).
fn same_name(title: &str, name: &str) -> bool {
    let norm = |s: &str| -> String {
        s.chars()
            .filter(char::is_ascii_alphanumeric)
            .map(|c| c.to_ascii_lowercase())
            .collect()
    };
    norm(title) == norm(name)
}

/// Drop a Rust numeric `format` and the `minimum: 0` an unsigned one implies.
fn drop_numeric_format(node: &mut Map<String, Value>) {
    let Some(format) = node.get("format").and_then(Value::as_str) else {
        return;
    };
    let unsigned = match format {
        "uint" | "uint8" | "uint16" | "uint32" | "uint64" | "uint128" => true,
        "int" | "int8" | "int16" | "int32" | "int64" | "int128" | "float" | "double" => false,
        _ => return,
    };
    node.remove("format");
    if unsigned && node.get("minimum").and_then(Value::as_f64) == Some(0.0) {
        node.remove("minimum");
    }
}

/// A schema that accepts only `null` (annotations aside).
fn is_null_only(schema: &Value) -> bool {
    let Some(obj) = schema.as_object() else {
        return false;
    };
    let mut validation = obj
        .iter()
        .filter(|(k, _)| !ANNOTATIONS.contains(&k.as_str()));
    let Some((key, value)) = validation.next() else {
        return false;
    };
    validation.next().is_none()
        && match key.as_str() {
            "type" => value == "null",
            "const" => value.is_null(),
            "enum" => value
                .as_array()
                .is_some_and(|v| v.len() == 1 && v[0].is_null()),
            _ => false,
        }
}

/// Remove the `null` alternative of an optional property.
fn drop_null(prop: &mut Map<String, Value>) {
    if let Some(Value::Array(types)) = prop.get_mut("type")
        && types.len() > 1
        && types.iter().any(|t| t == "null")
    {
        types.retain(|t| t != "null");
        if types.len() == 1 {
            let only = types.remove(0);
            prop.insert("type".into(), only);
        }
    }
    if let Some(Value::Array(values)) = prop.get_mut("enum")
        && values.len() > 1
        && values.iter().any(Value::is_null)
    {
        values.retain(|v| !v.is_null());
    }
    for key in ["anyOf", "oneOf"] {
        let Some(Value::Array(branches)) = prop.get_mut(key) else {
            continue;
        };
        if branches.len() < 2 || !branches.iter().any(is_null_only) {
            continue;
        }
        branches.retain(|b| !is_null_only(b));
        if branches.len() == 1 {
            let Some(Value::Object(only)) = branches.pop() else {
                continue;
            };
            prop.remove(key);
            if !merge_into(prop, &only) {
                // The property and its one branch clash: keep the branch as a one-item list.
                prop.insert(key.into(), Value::Array(vec![Value::Object(only)]));
            }
        }
    }
}

/// Merge `inner` into `outer`: the outer annotation wins, a validation keyword is copied over.
/// Refuses (and leaves `outer` unchanged) when both set one validation keyword differently.
fn merge_into(outer: &mut Map<String, Value>, inner: &Map<String, Value>) -> bool {
    let clash = inner
        .iter()
        .any(|(k, v)| !ANNOTATIONS.contains(&k.as_str()) && outer.get(k).is_some_and(|o| o != v));
    if clash {
        return false;
    }
    for (k, v) in inner {
        outer.entry(k.clone()).or_insert_with(|| v.clone());
    }
    true
}

/// Inline every `$defs` entry referenced once (and itself free of `$ref`s, so recursion stays
/// as it is), then drop `$defs` entries nothing refers to.
fn inline_defs(root: &mut JsonObject) {
    let mut keep: HashSet<String> = HashSet::new();
    loop {
        let Some(Value::Object(defs)) = root.get("$defs") else {
            return;
        };
        let counts = ref_counts(root);
        let candidate = defs
            .iter()
            .find(|(name, body)| {
                !keep.contains(*name) && counts.get(*name) == Some(&1) && !contains_any_ref(body)
            })
            .map(|(name, _)| name.clone());
        let Some(name) = candidate else { break };
        let Some(Value::Object(defs)) = root.get_mut("$defs") else {
            return;
        };
        let Some(Value::Object(mut body)) = defs.remove(&name) else {
            keep.insert(name);
            continue;
        };
        if body
            .get("title")
            .and_then(Value::as_str)
            .is_some_and(|t| same_name(t, &name))
        {
            body.remove("title");
        }
        let target = def_pointer(&name);
        if !replace_ref(root, &target, &body) {
            if let Some(Value::Object(defs)) = root.get_mut("$defs") {
                defs.insert(name.clone(), Value::Object(body));
            }
            keep.insert(name);
        }
    }
    // Drop unused definitions until none is left unused.
    loop {
        let counts = ref_counts(root);
        let Some(Value::Object(defs)) = root.get_mut("$defs") else {
            return;
        };
        let before = defs.len();
        defs.retain(|name, _| counts.contains_key(name));
        if defs.is_empty() {
            root.remove("$defs");
            return;
        }
        if defs.len() == before {
            return;
        }
    }
}

const DEFS_PREFIX: &str = "#/$defs/";

/// The `$ref` value pointing at a definition (JSON Pointer escaping applied).
fn def_pointer(name: &str) -> String {
    format!(
        "{DEFS_PREFIX}{}",
        name.replace('~', "~0").replace('/', "~1")
    )
}

/// The definition name a `$ref` value points at, if it is a local `$defs` reference.
fn def_name(reference: &str) -> Option<String> {
    let rest = reference.strip_prefix(DEFS_PREFIX)?;
    (!rest.contains('/')).then(|| rest.replace("~1", "/").replace("~0", "~"))
}

/// How often each definition is referenced anywhere in the document (definitions included).
fn ref_counts(root: &JsonObject) -> HashMap<String, usize> {
    let mut counts = HashMap::new();
    for value in root.values() {
        count_refs(value, &mut counts);
    }
    if let Some(name) = root.get("$ref").and_then(Value::as_str).and_then(def_name) {
        *counts.entry(name).or_default() += 1;
    }
    counts
}

fn count_refs(value: &Value, counts: &mut HashMap<String, usize>) {
    match value {
        Value::Object(obj) => {
            if let Some(name) = obj.get("$ref").and_then(Value::as_str).and_then(def_name) {
                *counts.entry(name).or_default() += 1;
            }
            obj.values().for_each(|v| count_refs(v, counts));
        }
        Value::Array(items) => items.iter().for_each(|v| count_refs(v, counts)),
        _ => {}
    }
}

/// Any `$ref` (local or not) in a schema; such a definition is left where it is.
fn contains_any_ref(value: &Value) -> bool {
    match value {
        Value::Object(obj) => obj.contains_key("$ref") || obj.values().any(contains_any_ref),
        Value::Array(items) => items.iter().any(contains_any_ref),
        _ => false,
    }
}

/// Replace the one schema `{"$ref": target, ...}` below `root` by the definition merged with
/// the reference's own keywords. `false` when the reference was not found or could not merge.
fn replace_ref(root: &mut JsonObject, target: &str, body: &Map<String, Value>) -> bool {
    fn visit(value: &mut Value, target: &str, body: &Map<String, Value>) -> Option<bool> {
        match value {
            Value::Object(obj) => {
                if obj.get("$ref").and_then(Value::as_str) == Some(target) {
                    let mut siblings = obj.clone();
                    siblings.remove("$ref");
                    let mut inlined = body.clone();
                    // The reference's own annotations (its description) win over the type's.
                    for (k, v) in &siblings {
                        if ANNOTATIONS.contains(&k.as_str()) {
                            inlined.insert(k.clone(), v.clone());
                        }
                    }
                    if !merge_into(&mut inlined, &siblings) {
                        return Some(false);
                    }
                    *obj = inlined;
                    return Some(true);
                }
                obj.iter_mut()
                    .filter(|(k, _)| k.as_str() != "$defs")
                    .find_map(|(_, v)| visit(v, target, body))
            }
            Value::Array(items) => items.iter_mut().find_map(|v| visit(v, target, body)),
            _ => None,
        }
    }
    root.iter_mut()
        .filter(|(k, _)| k.as_str() != "$defs")
        .find_map(|(_, v)| visit(v, target, body))
        .unwrap_or(false)
}

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

    fn compacted(schema: Value) -> Value {
        let Value::Object(mut obj) = schema else {
            panic!("not an object")
        };
        compact(&mut obj);
        Value::Object(obj)
    }

    #[test]
    fn optional_integer_loses_boilerplate() {
        let before = json!({
            "$schema": "https://json-schema.org/draft/2020-12/schema",
            "type": "object",
            "properties": {
                "path": { "type": "string", "description": "File to read." },
                "end_line": { "default": null, "description": "Last line.",
                    "format": "uint", "minimum": 0, "type": ["integer", "null"] }
            },
            "required": ["path"]
        });
        assert_eq!(
            compacted(before),
            json!({
                "type": "object",
                "properties": {
                    "path": { "type": "string", "description": "File to read." },
                    "end_line": { "description": "Last line.", "type": "integer" }
                },
                "required": ["path"]
            })
        );
    }

    #[test]
    fn required_nullable_and_real_bounds_stay() {
        let before = json!({
            "type": "object",
            "properties": {
                "a": { "type": ["string", "null"] },
                "n": { "type": "integer", "format": "uint32", "minimum": 1, "maximum": 300 },
                "b": { "type": "integer", "format": "uint8", "minimum": 0, "maximum": 255 },
                "s": { "type": "integer", "format": "int16", "minimum": -32768, "maximum": 32767 },
                "f": { "type": "boolean", "default": false },
                "e": { "type": "string", "format": "email" }
            },
            "required": ["a"],
            "additionalProperties": false
        });
        assert_eq!(
            compacted(before),
            json!({
                "type": "object",
                "properties": {
                    "a": { "type": ["string", "null"] },
                    "n": { "type": "integer", "minimum": 1, "maximum": 300 },
                    "b": { "type": "integer", "maximum": 255 },
                    "s": { "type": "integer", "minimum": -32768, "maximum": 32767 },
                    "f": { "type": "boolean", "default": false },
                    "e": { "type": "string", "format": "email" }
                },
                "required": ["a"],
                "additionalProperties": false
            })
        );
    }

    #[test]
    fn titles_that_repeat_the_name_go() {
        let before = json!({ "type": "object", "properties": {
            "max_hits": { "title": "Max Hits", "type": "integer" },
            "q": { "title": "", "type": "string" },
            "r": { "title": "Search pattern", "type": "string" }
        }});
        assert_eq!(
            compacted(before),
            json!({ "type": "object", "properties": {
                "max_hits": { "type": "integer" },
                "q": { "type": "string" },
                "r": { "title": "Search pattern", "type": "string" }
            }})
        );
    }

    #[test]
    fn single_use_definition_is_inlined_and_reference_annotations_win() {
        let before = json!({
            "$defs": { "Mode": { "description": "How.", "type": "string",
                "enum": ["a", "b"] } },
            "type": "object",
            "properties": { "mode": { "$ref": "#/$defs/Mode", "default": "a",
                "description": "Pick one." } }
        });
        assert_eq!(
            compacted(before),
            json!({ "type": "object", "properties": { "mode": {
                "description": "Pick one.", "type": "string", "enum": ["a", "b"],
                "default": "a" } } })
        );
    }

    #[test]
    fn optional_reference_drops_its_null_branch() {
        let before = json!({
            "$defs": { "Ref": { "description": "A project.",
                "anyOf": [ { "type": "integer", "format": "int64" }, { "type": "string" } ] } },
            "type": "object",
            "properties": { "project": { "description": "Only this project.",
                "anyOf": [ { "$ref": "#/$defs/Ref" }, { "type": "null" } ] } }
        });
        assert_eq!(
            compacted(before),
            json!({ "type": "object", "properties": { "project": {
                "description": "Only this project.",
                "anyOf": [ { "type": "integer" }, { "type": "string" } ] } } })
        );
    }

    #[test]
    fn shared_and_recursive_definitions_stay() {
        let before = json!({
            "$defs": {
                "P": { "type": "string" },
                "Tree": { "type": "object", "properties": {
                    "kids": { "type": "array", "items": { "$ref": "#/$defs/Tree" } } } },
                "Unused": { "type": "string" }
            },
            "type": "object",
            "properties": {
                "a": { "$ref": "#/$defs/P" },
                "b": { "$ref": "#/$defs/P" },
                "t": { "$ref": "#/$defs/Tree" }
            }
        });
        let after = compacted(before);
        let defs = after["$defs"].as_object().unwrap();
        assert!(
            defs.contains_key("P") && defs.contains_key("Tree"),
            "{after}"
        );
        assert!(!defs.contains_key("Unused"), "{after}");
        assert_eq!(after["properties"]["a"], json!({ "$ref": "#/$defs/P" }));
    }

    #[test]
    fn clashing_reference_is_left_alone() {
        let before = json!({
            "$defs": { "S": { "type": "string" } },
            "type": "object",
            "properties": { "x": { "$ref": "#/$defs/S", "type": "integer" } }
        });
        let after = compacted(before.clone());
        assert_eq!(after["properties"]["x"]["$ref"], "#/$defs/S");
        assert!(after["$defs"]["S"].is_object());
    }

    #[test]
    fn nested_objects_items_and_enums() {
        let before = json!({ "type": "object", "properties": {
            "inputs": { "type": "array", "items": { "type": "object",
                "properties": {
                    "pages": { "default": null, "type": ["string", "null"] },
                    "path": { "type": "string" } },
                "required": ["path"] } },
            "color": { "type": ["string", "null"], "enum": ["red", "blue", null] }
        }});
        assert_eq!(
            compacted(before),
            json!({ "type": "object", "properties": {
                "inputs": { "type": "array", "items": { "type": "object",
                    "properties": { "pages": { "type": "string" },
                        "path": { "type": "string" } },
                    "required": ["path"] } },
                "color": { "type": "string", "enum": ["red", "blue"] }
            }})
        );
    }

    #[test]
    fn output_schema_keeps_null() {
        let mut out = json!({ "$schema": "x", "type": "object", "properties": {
            "n": { "type": ["integer", "null"], "format": "uint", "minimum": 0,
                "default": null } } })
        .as_object()
        .unwrap()
        .clone();
        compact_output(&mut out);
        assert_eq!(
            Value::Object(out),
            json!({ "type": "object", "properties": { "n": { "type": ["integer", "null"] } } })
        );
    }

    #[test]
    fn compaction_is_idempotent() {
        let before = json!({ "$schema": "x", "type": "object", "properties": {
            "a": { "type": ["integer", "null"], "format": "uint", "minimum": 0 } } });
        let once = compacted(before);
        assert_eq!(compacted(once.clone()), once);
    }
}