feat(tools): inline local $ref in tool inputSchema (#2384)

src/mcp/client/stdio.py

@@ -174,7 +174,12 @@ async def stdin_writer():
                             errors=server.encoding_error_handler,
                         )
                     )
-        except anyio.ClosedResourceError:  # pragma: no cover
+        except (
+            anyio.BrokenResourceError,
+            anyio.ClosedResourceError,
+            BrokenPipeError,
+            ConnectionResetError,
+        ):  # pragma: no cover
             await anyio.lowlevel.checkpoint()
 
     async with anyio.create_task_group() as tg, process:

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,

src/mcp/server/mcpserver/utilities/schema.py

@@ -0,0 +1,133 @@
+"""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 TypeAlias, cast
+
+JSONPrimitive: TypeAlias = None | str | int | float | bool
+JSONValue: TypeAlias = JSONPrimitive | list["JSONValue"] | dict[str, "JSONValue"]
+JSONObject: TypeAlias = dict[str, JSONValue]
+
+
+def dereference_local_refs(schema: JSONObject) -> JSONObject:
+    """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
+
+    raw_defs = schema[defs_key]
+    if raw_defs is None:
+        return schema
+    if not isinstance(raw_defs, dict):
+        return schema
+
+    defs: JSONObject = raw_defs
+    if not defs:
+        return schema
+
+    # Cache resolved defs to avoid redundant traversal on diamond references.
+    resolved_defs: dict[str, JSONValue] = {}
+    # 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: JSONValue, stack: set[str]) -> JSONValue:
+        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: JSONObject = {k: v for k, v in node.items() if k != "$ref"}
+            if siblings and isinstance(resolved, dict):
+                resolved_schema = cast(JSONObject, resolved)
+                resolved_siblings: JSONObject = {key: inline(value, stack) for key, value in siblings.items()}
+                return {**resolved_schema, **resolved_siblings}
+            return resolved
+
+        # Regular object — recurse into values, but skip the top-level $defs container.
+        result: JSONObject = {}
+        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: JSONObject = {name: defs[name] for name in cyclic_defs if name in defs}
+        inlined[defs_key] = preserved
+
+    return inlined

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):

tests/server/mcpserver/utilities/test_schema.py

@@ -0,0 +1,178 @@
+"""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_non_object_defs_returns_schema_unchanged(self) -> None:
+        """Malformed non-object defs are ignored without copying the schema."""
+        schema: dict[str, Any] = {"type": "object", "$defs": ["not", "an", "object"]}
+        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: dict[str, Any] = {
+            "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"}]