Skip to content

Commit c21eea3

Browse files
committed
feat(tools): inline local $ref pointers in tool inputSchema (#2384)
Pydantic's model_json_schema() emits $ref/$defs for nested models, but LLM clients consuming tools/list often cannot resolve $ref and serialize referenced parameters as stringified JSON instead of structured objects. Inline local $ref pointers in tool schemas so they're self-contained and LLM-consumable, matching behavior in typescript-sdk (#1563) and go-sdk. Behavior: - Caches resolved defs (diamond references resolve each def once) - Cycles left in place with $defs entries preserved (degraded but correct) - Sibling keywords alongside $ref preserved per JSON Schema 2020-12 - External $ref and unknown local refs untouched - Input schema not mutated Changes: - Add mcp/server/mcpserver/utilities/schema.py with dereference_local_refs() - Apply to tool inputSchema at Tool.from_function() - 11 unit tests covering simple refs, diamond, cycles, siblings, legacy 'definitions' keyword, external refs, unknown refs, nested arrays/objects, input immutability Closes #2384 Signed-off-by: Mukunda Katta <mukunda.vjcs6@gmail.com>
1 parent 3d7b311 commit c21eea3

File tree

4 files changed

+272
-1
lines changed

4 files changed

+272
-1
lines changed

src/mcp/server/mcpserver/tools/base.py

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
from mcp.server.mcpserver.exceptions import ToolError
1010
from mcp.server.mcpserver.utilities.context_injection import find_context_parameter
1111
from mcp.server.mcpserver.utilities.func_metadata import FuncMetadata, func_metadata
12+
from mcp.server.mcpserver.utilities.schema import dereference_local_refs
1213
from mcp.shared._callable_inspection import is_async_callable
1314
from mcp.shared.exceptions import UrlElicitationRequiredError
1415
from mcp.shared.tool_name_validation import validate_and_warn_tool_name
@@ -72,7 +73,14 @@ def from_function(
7273
skip_names=[context_kwarg] if context_kwarg is not None else [],
7374
structured_output=structured_output,
7475
)
75-
parameters = func_arg_metadata.arg_model.model_json_schema(by_alias=True)
76+
# Pydantic emits $ref/$defs for nested models, which LLM clients often
77+
# can't resolve — they serialize referenced parameters as stringified
78+
# JSON instead of structured objects. Inline local refs so tool schemas
79+
# are self-contained and LLM-consumable. Matches behavior of
80+
# typescript-sdk (#1563) and go-sdk.
81+
parameters = dereference_local_refs(
82+
func_arg_metadata.arg_model.model_json_schema(by_alias=True)
83+
)
7684

7785
return cls(
7886
fn=fn,
Lines changed: 121 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,121 @@
1+
"""JSON Schema utilities for tool input schema preparation.
2+
3+
LLM clients consuming `tools/list` often cannot resolve JSON Schema ``$ref``
4+
pointers and serialize referenced parameters as stringified JSON instead of
5+
structured objects. This module provides :func:`dereference_local_refs` which
6+
inlines local ``$ref`` pointers so emitted tool schemas are self-contained.
7+
8+
This matches the behavior of the typescript-sdk (see
9+
`modelcontextprotocol/typescript-sdk#1563`_) and go-sdk.
10+
11+
.. _modelcontextprotocol/typescript-sdk#1563:
12+
https://github.com/modelcontextprotocol/typescript-sdk/pull/1563
13+
"""
14+
15+
from __future__ import annotations
16+
17+
from copy import copy
18+
from typing import Any
19+
20+
21+
def dereference_local_refs(schema: dict[str, Any]) -> dict[str, Any]:
22+
"""Inline local ``$ref`` pointers in a JSON Schema.
23+
24+
Behavior mirrors ``dereferenceLocalRefs`` in the TypeScript SDK:
25+
26+
- Caches resolved defs so diamond references (A→B→D, A→C→D) only resolve D once.
27+
- Cycles are detected and left in place — cyclic ``$ref`` pointers are kept
28+
along with their ``$defs`` entries so existing recursive schemas continue
29+
to work (degraded). Non-cyclic refs in the same schema are still inlined.
30+
- Sibling keywords alongside ``$ref`` are preserved per JSON Schema 2020-12
31+
(e.g. ``{"$ref": "#/$defs/X", "description": "override"}``).
32+
- Non-local ``$ref`` (external URLs, fragments outside ``$defs``) are left as-is.
33+
- Root self-references (``$ref: "#"``) are not handled — no library produces them.
34+
35+
If the schema has no ``$defs`` (or ``definitions``) container, it is returned
36+
unchanged.
37+
38+
Args:
39+
schema: The JSON Schema to process. Not mutated.
40+
41+
Returns:
42+
A new schema dict with local refs inlined. The ``$defs`` container is
43+
pruned to only the cyclic entries that remain referenced.
44+
"""
45+
# ``$defs`` is the standard keyword since JSON Schema 2019-09.
46+
# ``definitions`` is the legacy equivalent from drafts 04–07.
47+
# If both exist (malformed), ``$defs`` takes precedence.
48+
if "$defs" in schema:
49+
defs_key = "$defs"
50+
elif "definitions" in schema:
51+
defs_key = "definitions"
52+
else:
53+
return schema
54+
55+
defs: dict[str, Any] = schema[defs_key] or {}
56+
if not defs:
57+
return schema
58+
59+
# Cache resolved defs to avoid redundant traversal on diamond references.
60+
resolved_defs: dict[str, Any] = {}
61+
# Def names where a cycle was detected — their $ref is left in place and
62+
# their $defs entries must be preserved in the output.
63+
cyclic_defs: set[str] = set()
64+
prefix = f"#/{defs_key}/"
65+
66+
def inline(node: Any, stack: set[str]) -> Any:
67+
if node is None or isinstance(node, (str, int, float, bool)):
68+
return node
69+
if isinstance(node, list):
70+
return [inline(item, stack) for item in node]
71+
if not isinstance(node, dict):
72+
return node
73+
74+
ref = node.get("$ref")
75+
if isinstance(ref, str):
76+
if not ref.startswith(prefix):
77+
# External or non-local ref — leave as-is.
78+
return node
79+
def_name = ref[len(prefix) :]
80+
if def_name not in defs:
81+
# Unknown def — leave the ref untouched (pydantic shouldn't produce these).
82+
return node
83+
if def_name in stack:
84+
# Cycle detected — leave $ref in place, mark def for preservation.
85+
cyclic_defs.add(def_name)
86+
return node
87+
88+
if def_name in resolved_defs:
89+
resolved = resolved_defs[def_name]
90+
else:
91+
stack.add(def_name)
92+
resolved = inline(defs[def_name], stack)
93+
stack.discard(def_name)
94+
resolved_defs[def_name] = resolved
95+
96+
# Siblings of $ref (JSON Schema 2020-12).
97+
siblings = {k: v for k, v in node.items() if k != "$ref"}
98+
if siblings and isinstance(resolved, dict):
99+
resolved_siblings = {k: inline(v, stack) for k, v in siblings.items()}
100+
return {**resolved, **resolved_siblings}
101+
return resolved
102+
103+
# Regular object — recurse into values, but skip the top-level $defs container.
104+
result: dict[str, Any] = {}
105+
for key, value in node.items():
106+
if node is schema and key in ("$defs", "definitions"):
107+
continue
108+
result[key] = inline(value, stack)
109+
return result
110+
111+
inlined = inline(schema, set())
112+
if not isinstance(inlined, dict):
113+
# Shouldn't happen — a schema object always produces an object.
114+
return schema # pragma: no cover
115+
116+
# Preserve only cyclic defs in the output.
117+
if cyclic_defs:
118+
preserved = {name: defs[name] for name in cyclic_defs if name in defs}
119+
inlined[defs_key] = preserved
120+
121+
return inlined

tests/server/mcpserver/utilities/__init__.py

Whitespace-only changes.
Lines changed: 142 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,142 @@
1+
"""Tests for mcp.server.mcpserver.utilities.schema.dereference_local_refs."""
2+
3+
from __future__ import annotations
4+
5+
import pytest
6+
7+
from mcp.server.mcpserver.utilities.schema import dereference_local_refs
8+
9+
10+
class TestDereferenceLocalRefs:
11+
def test_no_defs_returns_schema_unchanged(self) -> None:
12+
schema = {"type": "object", "properties": {"x": {"type": "string"}}}
13+
assert dereference_local_refs(schema) == schema
14+
15+
def test_inlines_simple_ref(self) -> None:
16+
schema = {
17+
"type": "object",
18+
"properties": {"user": {"$ref": "#/$defs/User"}},
19+
"$defs": {"User": {"type": "object", "properties": {"name": {"type": "string"}}}},
20+
}
21+
result = dereference_local_refs(schema)
22+
assert result["properties"]["user"] == {
23+
"type": "object",
24+
"properties": {"name": {"type": "string"}},
25+
}
26+
# $defs pruned when fully resolved
27+
assert "$defs" not in result
28+
29+
def test_inlines_definitions_legacy_keyword(self) -> None:
30+
schema = {
31+
"properties": {"x": {"$ref": "#/definitions/Thing"}},
32+
"definitions": {"Thing": {"type": "integer"}},
33+
}
34+
result = dereference_local_refs(schema)
35+
assert result["properties"]["x"] == {"type": "integer"}
36+
37+
def test_dollar_defs_wins_when_both_present(self) -> None:
38+
schema = {
39+
"properties": {"x": {"$ref": "#/$defs/T"}},
40+
"$defs": {"T": {"type": "string"}},
41+
"definitions": {"T": {"type": "number"}},
42+
}
43+
result = dereference_local_refs(schema)
44+
assert result["properties"]["x"] == {"type": "string"}
45+
46+
def test_diamond_reference_resolved_once(self) -> None:
47+
schema = {
48+
"properties": {
49+
"a": {"$ref": "#/$defs/A"},
50+
"c": {"$ref": "#/$defs/C"},
51+
},
52+
"$defs": {
53+
"A": {"type": "object", "properties": {"d": {"$ref": "#/$defs/D"}}},
54+
"C": {"type": "object", "properties": {"d": {"$ref": "#/$defs/D"}}},
55+
"D": {"type": "string", "title": "the-d"},
56+
},
57+
}
58+
result = dereference_local_refs(schema)
59+
assert result["properties"]["a"]["properties"]["d"] == {"type": "string", "title": "the-d"}
60+
assert result["properties"]["c"]["properties"]["d"] == {"type": "string", "title": "the-d"}
61+
assert "$defs" not in result
62+
63+
def test_cycle_leaves_ref_in_place_and_preserves_def(self) -> None:
64+
# Node -> children[0] -> Node ... cyclic
65+
schema = {
66+
"type": "object",
67+
"properties": {"root": {"$ref": "#/$defs/Node"}},
68+
"$defs": {
69+
"Node": {
70+
"type": "object",
71+
"properties": {
72+
"value": {"type": "string"},
73+
"next": {"$ref": "#/$defs/Node"},
74+
},
75+
}
76+
},
77+
}
78+
result = dereference_local_refs(schema)
79+
# Cyclic ref left in place
80+
assert result["properties"]["root"]["properties"]["next"] == {"$ref": "#/$defs/Node"}
81+
# $defs entry for Node preserved so ref is resolvable
82+
assert "Node" in result["$defs"]
83+
84+
def test_sibling_keywords_preserved_via_2020_12_semantics(self) -> None:
85+
schema = {
86+
"properties": {"x": {"$ref": "#/$defs/Base", "description": "override"}},
87+
"$defs": {
88+
"Base": {
89+
"type": "string",
90+
"description": "original",
91+
"minLength": 1,
92+
}
93+
},
94+
}
95+
result = dereference_local_refs(schema)
96+
# Siblings override resolved, but other fields preserved
97+
assert result["properties"]["x"] == {
98+
"type": "string",
99+
"description": "override",
100+
"minLength": 1,
101+
}
102+
103+
def test_external_ref_left_as_is(self) -> None:
104+
schema = {
105+
"properties": {"x": {"$ref": "https://example.com/schema.json"}},
106+
"$defs": {"Local": {"type": "string"}},
107+
}
108+
result = dereference_local_refs(schema)
109+
assert result["properties"]["x"] == {"$ref": "https://example.com/schema.json"}
110+
111+
def test_unknown_local_ref_left_as_is(self) -> None:
112+
schema = {
113+
"properties": {"x": {"$ref": "#/$defs/DoesNotExist"}},
114+
"$defs": {"Other": {"type": "string"}},
115+
}
116+
result = dereference_local_refs(schema)
117+
assert result["properties"]["x"] == {"$ref": "#/$defs/DoesNotExist"}
118+
119+
def test_nested_arrays_and_objects_are_traversed(self) -> None:
120+
schema = {
121+
"type": "object",
122+
"properties": {
123+
"list": {
124+
"type": "array",
125+
"items": {"$ref": "#/$defs/Item"},
126+
}
127+
},
128+
"$defs": {"Item": {"type": "integer"}},
129+
}
130+
result = dereference_local_refs(schema)
131+
assert result["properties"]["list"]["items"] == {"type": "integer"}
132+
133+
def test_original_schema_not_mutated(self) -> None:
134+
schema = {
135+
"properties": {"x": {"$ref": "#/$defs/A"}},
136+
"$defs": {"A": {"type": "string"}},
137+
}
138+
original_defs = dict(schema["$defs"])
139+
_ = dereference_local_refs(schema)
140+
# Original still has $defs intact
141+
assert schema["$defs"] == original_defs
142+
assert schema["properties"]["x"] == {"$ref": "#/$defs/A"}

0 commit comments

Comments
 (0)