mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-24 06:24:18 +02:00
Backport: Dereference $ref in tool schemas for MCP client compatibility (#2861)
This commit is contained in:
Jeremiah Lowin
GitHub
parent
559b778135
commit
bc2f601e52
6 changed files with 255 additions and 5 deletions
1
.loq_cache
Normal file
1
.loq_cache
Normal file
File diff suppressed because one or more lines are too long
|
|
@ -173,6 +173,10 @@ def my_tool() -> None:
|
|||
|
||||
By default, FastMCP converts Python functions into MCP tools by inspecting the function's signature and type annotations. This allows you to use standard Python type annotations for your tools. In general, the framework strives to "just work": idiomatic Python behaviors like parameter defaults and type annotations are automatically translated into MCP schemas. However, there are a number of ways to customize the behavior of your tools.
|
||||
|
||||
<Note>
|
||||
FastMCP automatically dereferences `$ref` entries in tool schemas to ensure compatibility with MCP clients that don't fully support JSON Schema references (e.g., VS Code Copilot, Claude Desktop). This means complex Pydantic models with shared types are inlined in the schema rather than using `$defs` references.
|
||||
</Note>
|
||||
|
||||
### Type Annotations
|
||||
|
||||
MCP tools have typed arguments, and FastMCP uses type annotations to determine those types. Therefore, you should use standard Python type annotations for tool arguments:
|
||||
|
|
|
|||
|
|
@ -20,6 +20,7 @@ dependencies = [
|
|||
"uvicorn>=0.35",
|
||||
"websockets>=15.0.1",
|
||||
"jsonschema-path>=0.3.4",
|
||||
"jsonref>=1.1.0",
|
||||
]
|
||||
|
||||
requires-python = ">=3.10"
|
||||
|
|
|
|||
|
|
@ -3,6 +3,130 @@ from __future__ import annotations
|
|||
from collections import defaultdict
|
||||
from typing import Any
|
||||
|
||||
from jsonref import JsonRefError, replace_refs
|
||||
|
||||
|
||||
def dereference_refs(schema: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Resolve all $ref references in a JSON schema by inlining definitions.
|
||||
|
||||
This function resolves $ref references that point to $defs, replacing them
|
||||
with the actual definition content while preserving sibling keywords (like
|
||||
description, default, examples) that Pydantic places alongside $ref.
|
||||
|
||||
This is necessary because some MCP clients (e.g., VS Code Copilot) don't
|
||||
properly handle $ref in tool input schemas.
|
||||
|
||||
For self-referencing/circular schemas where full dereferencing is not possible,
|
||||
this function falls back to resolving only the root-level $ref while preserving
|
||||
$defs for nested references.
|
||||
|
||||
Args:
|
||||
schema: JSON schema dict that may contain $ref references
|
||||
|
||||
Returns:
|
||||
A new schema dict with $ref resolved where possible and $defs removed
|
||||
when no longer needed
|
||||
|
||||
Example:
|
||||
>>> schema = {
|
||||
... "$defs": {"Category": {"enum": ["a", "b"], "type": "string"}},
|
||||
... "properties": {"cat": {"$ref": "#/$defs/Category", "default": "a"}}
|
||||
... }
|
||||
>>> resolved = dereference_refs(schema)
|
||||
>>> # Result: {"properties": {"cat": {"enum": ["a", "b"], "type": "string", "default": "a"}}}
|
||||
"""
|
||||
try:
|
||||
# Use jsonref to resolve all $ref references
|
||||
# proxies=False returns plain dicts (not proxy objects)
|
||||
# lazy_load=False resolves immediately
|
||||
dereferenced = replace_refs(schema, proxies=False, lazy_load=False)
|
||||
|
||||
# Merge sibling keywords that were lost during dereferencing
|
||||
# Pydantic puts description, default, examples as siblings to $ref
|
||||
defs = schema.get("$defs", {})
|
||||
merged = _merge_ref_siblings(schema, dereferenced, defs)
|
||||
# Type assertion: top-level schema is always a dict
|
||||
assert isinstance(merged, dict)
|
||||
dereferenced = merged
|
||||
|
||||
# Remove $defs since all references have been resolved
|
||||
if "$defs" in dereferenced:
|
||||
dereferenced = {k: v for k, v in dereferenced.items() if k != "$defs"}
|
||||
|
||||
return dereferenced
|
||||
|
||||
except JsonRefError:
|
||||
# Self-referencing/circular schemas can't be fully dereferenced
|
||||
# Fall back to resolving only root-level $ref (for MCP spec compliance)
|
||||
return resolve_root_ref(schema)
|
||||
|
||||
|
||||
def _merge_ref_siblings(
|
||||
original: Any,
|
||||
dereferenced: Any,
|
||||
defs: dict[str, Any],
|
||||
visited: set[str] | None = None,
|
||||
) -> Any:
|
||||
"""Merge sibling keywords from original $ref nodes into dereferenced schema.
|
||||
|
||||
When jsonref resolves $ref, it replaces the entire node with the referenced
|
||||
definition, losing any sibling keywords like description, default, or examples.
|
||||
This function walks both trees in parallel and merges those siblings back.
|
||||
|
||||
Args:
|
||||
original: The original schema with $ref and potential siblings
|
||||
dereferenced: The schema after jsonref processing
|
||||
defs: The $defs from the original schema, for looking up referenced definitions
|
||||
visited: Set of definition names already being processed (prevents cycles)
|
||||
|
||||
Returns:
|
||||
The dereferenced schema with sibling keywords restored
|
||||
"""
|
||||
if visited is None:
|
||||
visited = set()
|
||||
|
||||
if isinstance(original, dict) and isinstance(dereferenced, dict):
|
||||
# Check if original had a $ref
|
||||
if "$ref" in original:
|
||||
ref = original["$ref"]
|
||||
siblings = {k: v for k, v in original.items() if k not in ("$ref", "$defs")}
|
||||
|
||||
# Look up the referenced definition to process its nested siblings
|
||||
if isinstance(ref, str) and ref.startswith("#/$defs/"):
|
||||
def_name = ref.split("/")[-1]
|
||||
# Prevent infinite recursion on circular references
|
||||
if def_name in defs and def_name not in visited:
|
||||
# Recursively process the definition's content for nested siblings
|
||||
dereferenced = _merge_ref_siblings(
|
||||
defs[def_name], dereferenced, defs, visited | {def_name}
|
||||
)
|
||||
|
||||
if siblings:
|
||||
# Merge local siblings, which take precedence
|
||||
merged = dict(dereferenced)
|
||||
merged.update(siblings)
|
||||
return merged
|
||||
return dereferenced
|
||||
|
||||
# Recurse into nested structures
|
||||
result = {}
|
||||
for key, value in dereferenced.items():
|
||||
if key in original:
|
||||
result[key] = _merge_ref_siblings(original[key], value, defs, visited)
|
||||
else:
|
||||
result[key] = value
|
||||
return result
|
||||
|
||||
elif isinstance(original, list) and isinstance(dereferenced, list):
|
||||
# Process list items in parallel
|
||||
min_len = min(len(original), len(dereferenced))
|
||||
return [
|
||||
_merge_ref_siblings(o, d, defs, visited)
|
||||
for o, d in zip(original[:min_len], dereferenced[:min_len], strict=False)
|
||||
] + dereferenced[min_len:]
|
||||
|
||||
return dereferenced
|
||||
|
||||
|
||||
def resolve_root_ref(schema: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Resolve $ref at root level to meet MCP spec requirements.
|
||||
|
|
@ -43,7 +167,7 @@ def resolve_root_ref(schema: dict[str, Any]) -> dict[str, Any]:
|
|||
return schema
|
||||
|
||||
|
||||
def _prune_param(schema: dict, param: str) -> dict:
|
||||
def _prune_param(schema: dict[str, Any], param: str) -> dict[str, Any]:
|
||||
"""Return a new schema with *param* removed from `properties`, `required`,
|
||||
and (if no longer referenced) `$defs`.
|
||||
"""
|
||||
|
|
@ -65,11 +189,11 @@ def _prune_param(schema: dict, param: str) -> dict:
|
|||
|
||||
|
||||
def _single_pass_optimize(
|
||||
schema: dict,
|
||||
schema: dict[str, Any],
|
||||
prune_titles: bool = False,
|
||||
prune_additional_properties: bool = False,
|
||||
prune_defs: bool = True,
|
||||
) -> dict:
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Optimize JSON schemas in a single traversal for better performance.
|
||||
|
||||
|
|
@ -238,12 +362,12 @@ def _single_pass_optimize(
|
|||
|
||||
|
||||
def compress_schema(
|
||||
schema: dict,
|
||||
schema: dict[str, Any],
|
||||
prune_params: list[str] | None = None,
|
||||
prune_defs: bool = True,
|
||||
prune_additional_properties: bool = True,
|
||||
prune_titles: bool = False,
|
||||
) -> dict:
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Remove the given parameters from the schema.
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
from fastmcp.utilities.json_schema import (
|
||||
_prune_param,
|
||||
compress_schema,
|
||||
dereference_refs,
|
||||
resolve_root_ref,
|
||||
)
|
||||
|
||||
|
|
@ -508,6 +509,114 @@ class TestCompressSchema:
|
|||
assert "title" not in compressed["properties"]["normal_field"]
|
||||
|
||||
|
||||
class TestDereferenceRefs:
|
||||
"""Tests for the dereference_refs function."""
|
||||
|
||||
def test_falls_back_for_circular_refs(self):
|
||||
"""Test that circular references fall back to resolve_root_ref."""
|
||||
schema = {
|
||||
"$defs": {
|
||||
"Node": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"children": {
|
||||
"type": "array",
|
||||
"items": {"$ref": "#/$defs/Node"},
|
||||
}
|
||||
},
|
||||
}
|
||||
},
|
||||
"$ref": "#/$defs/Node",
|
||||
}
|
||||
result = dereference_refs(schema)
|
||||
|
||||
# Should fall back to resolve_root_ref behavior
|
||||
# Root should be resolved but nested refs preserved
|
||||
assert result.get("type") == "object"
|
||||
assert "$defs" in result # $defs preserved for circular refs
|
||||
|
||||
def test_preserves_sibling_keywords(self):
|
||||
"""Test that sibling keywords (default, description) are preserved.
|
||||
|
||||
Pydantic places description, default, examples as siblings to $ref.
|
||||
These should not be lost during dereferencing.
|
||||
"""
|
||||
schema = {
|
||||
"$defs": {
|
||||
"Status": {"type": "string", "enum": ["active", "inactive"]},
|
||||
},
|
||||
"properties": {
|
||||
"status": {
|
||||
"$ref": "#/$defs/Status",
|
||||
"default": "active",
|
||||
"description": "The user status",
|
||||
},
|
||||
},
|
||||
"type": "object",
|
||||
}
|
||||
result = dereference_refs(schema)
|
||||
|
||||
# $ref should be inlined with siblings preserved
|
||||
status = result["properties"]["status"]
|
||||
assert status["type"] == "string"
|
||||
assert status["enum"] == ["active", "inactive"]
|
||||
assert status["default"] == "active"
|
||||
assert status["description"] == "The user status"
|
||||
# $defs should be removed
|
||||
assert "$defs" not in result
|
||||
|
||||
def test_preserves_siblings_in_lists(self):
|
||||
"""Test that siblings are preserved for $refs inside lists (allOf, anyOf, etc)."""
|
||||
schema = {
|
||||
"$defs": {
|
||||
"StringType": {"type": "string"},
|
||||
"IntType": {"type": "integer"},
|
||||
},
|
||||
"properties": {
|
||||
"field": {
|
||||
"anyOf": [
|
||||
{"$ref": "#/$defs/StringType", "description": "As string"},
|
||||
{"$ref": "#/$defs/IntType", "description": "As integer"},
|
||||
]
|
||||
},
|
||||
},
|
||||
}
|
||||
result = dereference_refs(schema)
|
||||
|
||||
# Both items in anyOf should have their siblings preserved
|
||||
any_of = result["properties"]["field"]["anyOf"]
|
||||
assert any_of[0]["type"] == "string"
|
||||
assert any_of[0]["description"] == "As string"
|
||||
assert any_of[1]["type"] == "integer"
|
||||
assert any_of[1]["description"] == "As integer"
|
||||
assert "$defs" not in result
|
||||
|
||||
def test_preserves_nested_siblings(self):
|
||||
"""Test that siblings on nested $refs are preserved."""
|
||||
schema = {
|
||||
"$defs": {
|
||||
"Address": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"country": {"$ref": "#/$defs/Country", "default": "US"},
|
||||
},
|
||||
},
|
||||
"Country": {"type": "string", "enum": ["US", "UK", "CA"]},
|
||||
},
|
||||
"properties": {
|
||||
"home_address": {"$ref": "#/$defs/Address"},
|
||||
},
|
||||
}
|
||||
result = dereference_refs(schema)
|
||||
|
||||
# The nested $ref's sibling (default) should be preserved
|
||||
country = result["properties"]["home_address"]["properties"]["country"]
|
||||
assert country["type"] == "string"
|
||||
assert country["enum"] == ["US", "UK", "CA"]
|
||||
assert country["default"] == "US"
|
||||
assert "$defs" not in result
|
||||
|
||||
|
||||
class TestResolveRootRef:
|
||||
"""Tests for the resolve_root_ref function.
|
||||
|
||||
|
|
|
|||
11
uv.lock
generated
11
uv.lock
generated
|
|
@ -689,6 +689,7 @@ dependencies = [
|
|||
{ name = "cyclopts" },
|
||||
{ name = "exceptiongroup" },
|
||||
{ name = "httpx" },
|
||||
{ name = "jsonref" },
|
||||
{ name = "jsonschema-path" },
|
||||
{ name = "mcp" },
|
||||
{ name = "openapi-pydantic" },
|
||||
|
|
@ -745,6 +746,7 @@ requires-dist = [
|
|||
{ name = "cyclopts", specifier = ">=4.0.0" },
|
||||
{ name = "exceptiongroup", specifier = ">=1.2.2" },
|
||||
{ name = "httpx", specifier = ">=0.28.1" },
|
||||
{ name = "jsonref", specifier = ">=1.1.0" },
|
||||
{ name = "jsonschema-path", specifier = ">=0.3.4" },
|
||||
{ name = "mcp", specifier = ">=1.24.0,<2.0" },
|
||||
{ name = "openai", marker = "extra == 'openai'", specifier = ">=1.102.0" },
|
||||
|
|
@ -1100,6 +1102,15 @@ wheels = [
|
|||
{ url = "https://files.pythonhosted.org/packages/2f/9c/6753e6522b8d0ef07d3a3d239426669e984fb0eba15a315cdbc1253904e4/jiter-0.12.0-graalpy312-graalpy250_312_native-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c24e864cb30ab82311c6425655b0cdab0a98c5d973b065c66a3f020740c2324c", size = 346110, upload-time = "2025-11-09T20:49:21.817Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "jsonref"
|
||||
version = "1.1.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/aa/0d/c1f3277e90ccdb50d33ed5ba1ec5b3f0a242ed8c1b1a85d3afeb68464dca/jsonref-1.1.0.tar.gz", hash = "sha256:32fe8e1d85af0fdefbebce950af85590b22b60f9e95443176adbde4e1ecea552", size = 8814, upload-time = "2023-01-16T16:10:04.455Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/0c/ec/e1db9922bceb168197a558a2b8c03a7963f1afe93517ddd3cf99f202f996/jsonref-1.1.0-py3-none-any.whl", hash = "sha256:590dc7773df6c21cbf948b5dac07a72a251db28b0238ceecce0a2abfa8ec30a9", size = 9425, upload-time = "2023-01-16T16:10:02.255Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "jsonschema"
|
||||
version = "4.25.1"
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue