diff --git a/src/mcp/server/mcpserver/tools/base.py b/src/mcp/server/mcpserver/tools/base.py index 754313eb8..39d323180 100644 --- a/src/mcp/server/mcpserver/tools/base.py +++ b/src/mcp/server/mcpserver/tools/base.py @@ -9,6 +9,7 @@ from mcp.server.mcpserver.exceptions import ToolError from mcp.server.mcpserver.utilities.context_injection import find_context_parameter from mcp.server.mcpserver.utilities.func_metadata import FuncMetadata, func_metadata +from mcp.server.mcpserver.utilities.schema import dereference_local_refs from mcp.shared._callable_inspection import is_async_callable from mcp.shared.exceptions import UrlElicitationRequiredError from mcp.shared.tool_name_validation import validate_and_warn_tool_name @@ -72,7 +73,12 @@ def from_function( skip_names=[context_kwarg] if context_kwarg is not None else [], structured_output=structured_output, ) - parameters = func_arg_metadata.arg_model.model_json_schema(by_alias=True) + # Pydantic emits $ref/$defs for nested models, which LLM clients often + # can't resolve — they serialize referenced parameters as stringified + # JSON instead of structured objects. Inline local refs so tool schemas + # are self-contained and LLM-consumable. Matches behavior of + # typescript-sdk (#1563) and go-sdk. + parameters = dereference_local_refs(func_arg_metadata.arg_model.model_json_schema(by_alias=True)) return cls( fn=fn, diff --git a/src/mcp/server/mcpserver/utilities/schema.py b/src/mcp/server/mcpserver/utilities/schema.py new file mode 100644 index 000000000..89c68f106 --- /dev/null +++ b/src/mcp/server/mcpserver/utilities/schema.py @@ -0,0 +1,122 @@ +"""JSON Schema utilities for tool input schema preparation. + +LLM clients consuming `tools/list` often cannot resolve JSON Schema ``$ref`` +pointers and serialize referenced parameters as stringified JSON instead of +structured objects. This module provides :func:`dereference_local_refs` which +inlines local ``$ref`` pointers so emitted tool schemas are self-contained. + +This matches the behavior of the typescript-sdk (see +`modelcontextprotocol/typescript-sdk#1563`_) and go-sdk. + +.. _modelcontextprotocol/typescript-sdk#1563: + https://github.com/modelcontextprotocol/typescript-sdk/pull/1563 +""" + +from __future__ import annotations + +from typing import Any + + +def dereference_local_refs(schema: dict[str, Any]) -> dict[str, Any]: + """Inline local ``$ref`` pointers in a JSON Schema. + + Behavior mirrors ``dereferenceLocalRefs`` in the TypeScript SDK: + + - Caches resolved defs so diamond references (A→B→D, A→C→D) only resolve D once. + - Cycles are detected and left in place — cyclic ``$ref`` pointers are kept + along with their ``$defs`` entries so existing recursive schemas continue + to work (degraded). Non-cyclic refs in the same schema are still inlined. + - Sibling keywords alongside ``$ref`` are preserved per JSON Schema 2020-12 + (e.g. ``{"$ref": "#/$defs/X", "description": "override"}``). + - Non-local ``$ref`` (external URLs, fragments outside ``$defs``) are left as-is. + - Root self-references (``$ref: "#"``) are not handled — no library produces them. + + If the schema has no ``$defs`` (or ``definitions``) container, it is returned + unchanged. + + Args: + schema: The JSON Schema to process. Not mutated. + + Returns: + A new schema dict with local refs inlined. The ``$defs`` container is + pruned to only the cyclic entries that remain referenced. + """ + # ``$defs`` is the standard keyword since JSON Schema 2019-09. + # ``definitions`` is the legacy equivalent from drafts 04–07. + # If both exist (malformed), ``$defs`` takes precedence. + if "$defs" in schema: + defs_key = "$defs" + elif "definitions" in schema: + defs_key = "definitions" + else: + return schema + + defs: dict[str, Any] = schema[defs_key] or {} + if not defs: + return schema + + # Cache resolved defs to avoid redundant traversal on diamond references. + resolved_defs: dict[str, Any] = {} + # Def names where a cycle was detected — their $ref is left in place and + # their $defs entries must be preserved in the output. + cyclic_defs: set[str] = set() + prefix = f"#/{defs_key}/" + + def inline(node: Any, stack: set[str]) -> Any: + if node is None or isinstance(node, (str, int, float, bool)): + return node + if isinstance(node, list): + return [inline(item, stack) for item in node] + if not isinstance(node, dict): # pragma: no cover + # Defensive: valid JSON only contains None/str/int/float/bool/list/dict. + # Reachable only if a non-JSON-shaped value sneaks into a schema. + return node + + ref = node.get("$ref") + if isinstance(ref, str): + if not ref.startswith(prefix): + # External or non-local ref — leave as-is. + return node + def_name = ref[len(prefix) :] + if def_name not in defs: + # Unknown def — leave the ref untouched (pydantic shouldn't produce these). + return node + if def_name in stack: + # Cycle detected — leave $ref in place, mark def for preservation. + cyclic_defs.add(def_name) + return node + + if def_name in resolved_defs: + resolved = resolved_defs[def_name] + else: + stack.add(def_name) + resolved = inline(defs[def_name], stack) + stack.discard(def_name) + resolved_defs[def_name] = resolved + + # Siblings of $ref (JSON Schema 2020-12). + siblings = {k: v for k, v in node.items() if k != "$ref"} + if siblings and isinstance(resolved, dict): + resolved_siblings = {k: inline(v, stack) for k, v in siblings.items()} + return {**resolved, **resolved_siblings} + return resolved + + # Regular object — recurse into values, but skip the top-level $defs container. + result: dict[str, Any] = {} + for key, value in node.items(): + if node is schema and key in ("$defs", "definitions"): + continue + result[key] = inline(value, stack) + return result + + inlined = inline(schema, set()) + if not isinstance(inlined, dict): + # Shouldn't happen — a schema object always produces an object. + return schema # pragma: no cover + + # Preserve only cyclic defs in the output. + if cyclic_defs: + preserved = {name: defs[name] for name in cyclic_defs if name in defs} + inlined[defs_key] = preserved + + return inlined diff --git a/tests/server/mcpserver/test_tool_manager.py b/tests/server/mcpserver/test_tool_manager.py index e4dfd4ff9..87eb40dac 100644 --- a/tests/server/mcpserver/test_tool_manager.py +++ b/tests/server/mcpserver/test_tool_manager.py @@ -100,8 +100,10 @@ def create_user(user: UserInput, flag: bool) -> dict[str, Any]: # pragma: no co assert tool.name == "create_user" assert tool.description == "Create a new user." assert tool.is_async is False - assert "name" in tool.parameters["$defs"]["UserInput"]["properties"] - assert "age" in tool.parameters["$defs"]["UserInput"]["properties"] + # $ref is now inlined (see dereference_local_refs in utilities/schema.py). + # The UserInput definition is merged directly into properties.user. + assert "name" in tool.parameters["properties"]["user"]["properties"] + assert "age" in tool.parameters["properties"]["user"]["properties"] assert "flag" in tool.parameters["properties"] def test_add_callable_object(self): diff --git a/tests/server/mcpserver/utilities/__init__.py b/tests/server/mcpserver/utilities/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/server/mcpserver/utilities/test_schema.py b/tests/server/mcpserver/utilities/test_schema.py new file mode 100644 index 000000000..8b3f8908c --- /dev/null +++ b/tests/server/mcpserver/utilities/test_schema.py @@ -0,0 +1,172 @@ +"""Tests for mcp.server.mcpserver.utilities.schema.dereference_local_refs.""" + +from __future__ import annotations + +from typing import Any + +from mcp.server.mcpserver.utilities.schema import dereference_local_refs + + +class TestDereferenceLocalRefs: + def test_no_defs_returns_schema_unchanged(self) -> None: + schema = {"type": "object", "properties": {"x": {"type": "string"}}} + assert dereference_local_refs(schema) == schema + + def test_inlines_simple_ref(self) -> None: + schema = { + "type": "object", + "properties": {"user": {"$ref": "#/$defs/User"}}, + "$defs": {"User": {"type": "object", "properties": {"name": {"type": "string"}}}}, + } + result = dereference_local_refs(schema) + assert result["properties"]["user"] == { + "type": "object", + "properties": {"name": {"type": "string"}}, + } + # $defs pruned when fully resolved + assert "$defs" not in result + + def test_inlines_definitions_legacy_keyword(self) -> None: + schema = { + "properties": {"x": {"$ref": "#/definitions/Thing"}}, + "definitions": {"Thing": {"type": "integer"}}, + } + result = dereference_local_refs(schema) + assert result["properties"]["x"] == {"type": "integer"} + + def test_dollar_defs_wins_when_both_present(self) -> None: + schema = { + "properties": {"x": {"$ref": "#/$defs/T"}}, + "$defs": {"T": {"type": "string"}}, + "definitions": {"T": {"type": "number"}}, + } + result = dereference_local_refs(schema) + assert result["properties"]["x"] == {"type": "string"} + + def test_diamond_reference_resolved_once(self) -> None: + schema = { + "properties": { + "a": {"$ref": "#/$defs/A"}, + "c": {"$ref": "#/$defs/C"}, + }, + "$defs": { + "A": {"type": "object", "properties": {"d": {"$ref": "#/$defs/D"}}}, + "C": {"type": "object", "properties": {"d": {"$ref": "#/$defs/D"}}}, + "D": {"type": "string", "title": "the-d"}, + }, + } + result = dereference_local_refs(schema) + assert result["properties"]["a"]["properties"]["d"] == {"type": "string", "title": "the-d"} + assert result["properties"]["c"]["properties"]["d"] == {"type": "string", "title": "the-d"} + assert "$defs" not in result + + def test_cycle_leaves_ref_in_place_and_preserves_def(self) -> None: + # Node -> children[0] -> Node ... cyclic + schema = { + "type": "object", + "properties": {"root": {"$ref": "#/$defs/Node"}}, + "$defs": { + "Node": { + "type": "object", + "properties": { + "value": {"type": "string"}, + "next": {"$ref": "#/$defs/Node"}, + }, + } + }, + } + result = dereference_local_refs(schema) + # Cyclic ref left in place + assert result["properties"]["root"]["properties"]["next"] == {"$ref": "#/$defs/Node"} + # $defs entry for Node preserved so ref is resolvable + assert "Node" in result["$defs"] + + def test_sibling_keywords_preserved_via_2020_12_semantics(self) -> None: + schema = { + "properties": {"x": {"$ref": "#/$defs/Base", "description": "override"}}, + "$defs": { + "Base": { + "type": "string", + "description": "original", + "minLength": 1, + } + }, + } + result = dereference_local_refs(schema) + # Siblings override resolved, but other fields preserved + assert result["properties"]["x"] == { + "type": "string", + "description": "override", + "minLength": 1, + } + + def test_external_ref_left_as_is(self) -> None: + schema = { + "properties": {"x": {"$ref": "https://example.com/schema.json"}}, + "$defs": {"Local": {"type": "string"}}, + } + result = dereference_local_refs(schema) + assert result["properties"]["x"] == {"$ref": "https://example.com/schema.json"} + + def test_unknown_local_ref_left_as_is(self) -> None: + schema = { + "properties": {"x": {"$ref": "#/$defs/DoesNotExist"}}, + "$defs": {"Other": {"type": "string"}}, + } + result = dereference_local_refs(schema) + assert result["properties"]["x"] == {"$ref": "#/$defs/DoesNotExist"} + + def test_nested_arrays_and_objects_are_traversed(self) -> None: + schema = { + "type": "object", + "properties": { + "list": { + "type": "array", + "items": {"$ref": "#/$defs/Item"}, + } + }, + "$defs": {"Item": {"type": "integer"}}, + } + result = dereference_local_refs(schema) + assert result["properties"]["list"]["items"] == {"type": "integer"} + + def test_original_schema_not_mutated(self) -> None: + schema = { + "properties": {"x": {"$ref": "#/$defs/A"}}, + "$defs": {"A": {"type": "string"}}, + } + original_defs = dict(schema["$defs"]) + _ = dereference_local_refs(schema) + # Original still has $defs intact + assert schema["$defs"] == original_defs + assert schema["properties"]["x"] == {"$ref": "#/$defs/A"} + + def test_empty_defs_returns_schema_unchanged(self) -> None: + """`$defs: {}` (empty container) is a no-op — returns input as-is.""" + schema = {"type": "object", "$defs": {}} + result = dereference_local_refs(schema) + assert result is schema # same object — no copy made on the empty path + + def test_null_defs_returns_schema_unchanged(self) -> None: + """`$defs: null` falls through the same empty-defs path.""" + schema: dict[str, Any] = {"type": "object", "$defs": None} + result = dereference_local_refs(schema) + assert result is schema + + def test_inlines_through_array_of_objects(self) -> None: + """Refs nested inside arrays of dict items are recursed properly. + + Covers the `if isinstance(node, list)` branch of the inner inline(). + """ + schema = { + "anyOf": [ + {"$ref": "#/$defs/A"}, + {"$ref": "#/$defs/B"}, + ], + "$defs": { + "A": {"type": "string"}, + "B": {"type": "integer"}, + }, + } + result = dereference_local_refs(schema) + assert result["anyOf"] == [{"type": "string"}, {"type": "integer"}]