Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 29 additions & 0 deletions docs/servers/middleware.mdx
Original file line number Diff line number Diff line change
Expand Up @@ -484,6 +484,35 @@ mcp.add_middleware(StructuredLoggingMiddleware(include_payloads=True))

The built-in versions include payload logging, structured JSON output, custom logger support, payload size limits, and operation-specific hooks for granular control.

### Schema Dereference Middleware

Some MCP clients have limited support for JSON Schema `$ref` references in tool and resource template schemas. This can lead to null parameters or validation failures when a client cannot resolve references during tool invocation.

<Info>
The schema dereference middleware flattens internal `$ref` references (those pointing to `#/$defs/...`) in listing responses, so clients receive inline schemas that are easier to consume.
</Info>

```python
from fastmcp import FastMCP
from fastmcp.server.middleware.schema_dereference import SchemaDereferenceMiddleware

mcp = FastMCP("MyServer")
mcp.add_middleware(SchemaDereferenceMiddleware())
```

What it does
- Inlines internal `$ref` to `#/$defs/...` for:
- on_list_tools: `Tool.parameters`
- on_list_resource_templates: `ResourceTemplate.parameters`
- Handles common schema shapes:
- Properties, arrays (`items`), maps (`additionalProperties`)
- Composition keywords (`allOf`, `oneOf`, `anyOf`)
- Transitive and reused definitions

<Tip>
Enable this middleware when serving tools or templates to clients that cannot resolve JSON Schema `$ref` internally. It is safe to use with clients that fully support `$ref` as well.
</Tip>

### Rate Limiting Middleware

Rate limiting is essential for protecting your server from abuse, ensuring fair resource usage, and maintaining performance under load. FastMCP includes sophisticated rate limiting middleware at `fastmcp.server.middleware.rate_limiting`.
Expand Down
201 changes: 201 additions & 0 deletions src/fastmcp/server/middleware/schema_dereference.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,201 @@
from copy import deepcopy
from typing import Any

import mcp.types as mt

from fastmcp.server.middleware.middleware import CallNext, Middleware, MiddlewareContext


def _detect_self_reference(schema: dict) -> bool:
"""
Detect if the schema contains self-referencing definitions.
Args:
schema: The JSON schema to check
Returns:
True if self-referencing is detected
"""
defs = schema.get("$defs", {})

def find_refs_in_value(value: Any, parent_def: str) -> bool:
"""Check if a value contains a reference to its parent definition."""
if isinstance(value, dict):
if "$ref" in value:
ref_path = value["$ref"]
# Check if this references the parent definition
if ref_path == f"#/$defs/{parent_def}":
return True
# Check all values in the dict
for v in value.values():
if find_refs_in_value(v, parent_def):
return True
elif isinstance(value, list):
# Check all items in the list
for item in value:
if find_refs_in_value(item, parent_def):
return True
return False

# Check each definition for self-reference
for def_name, def_content in defs.items():
if find_refs_in_value(def_content, def_name):
# Self-reference detected, return original schema
return True

return False


def dereference_json_schema(schema: dict) -> dict:
"""
Dereference a JSON schema by resolving $ref references while preserving $defs only when corner cases occur.
This function flattens schema properties by:
1. Check for self-reference - if found, return original schema with $defs
2. When encountering $refs in properties, resolve them on-demand
3. Track visited definitions globally to prevent circular expansion
4. Only preserve original $defs if corner cases are encountered:
- Self-reference detected
- Circular references between definitions
- Reference not found in $defs
Args:
schema: The JSON schema to flatten
Returns:
Schema with references resolved in properties, keeping $defs only when corner cases occur
"""
# Step 1: Check for self-reference
if _detect_self_reference(schema):
# Self-referencing detected, return original schema with $defs
return schema

# Make a deep copy to work with
result = deepcopy(schema)

# Keep original $defs for potential corner cases
defs = deepcopy(schema.get("$defs", {}))

# Track corner cases that require preserving $defs
corner_cases_detected = {
"circular_ref": False,
"ref_not_found": False,
}

# Step 2: Define resolution function that tracks visits globally and corner cases
def resolve_refs_in_value(value: Any, depth: int, visiting: set[str]) -> Any:
"""
Recursively resolve $refs in a value.
Args:
value: The value to process
depth: Current depth in resolution
visiting: Set of definitions currently being resolved (for cycle detection)
Returns:
Value with $refs resolved (or kept if corner cases occur)
"""
if isinstance(value, dict):
if "$ref" in value:
ref_path = value["$ref"]

# Only handle internal references to $defs
if isinstance(ref_path, str) and ref_path.startswith("#/$defs/"):
def_name = ref_path.split("/")[-1]

# Check for circular reference
if def_name in visiting:
# Circular reference detected, keep the $ref
corner_cases_detected["circular_ref"] = True
return value

if def_name in defs:
# Add to visiting set
visiting.add(def_name)

# Get the definition and resolve any refs within it
resolved = resolve_refs_in_value(
deepcopy(defs[def_name]), depth + 1, visiting
)

# Remove from visiting set
visiting.remove(def_name)

# Merge resolved definition with additional properties
# Additional properties from the original object take precedence
for key, val in value.items():
if key != "$ref":
resolved[key] = val

return resolved
else:
# Definition not found, keep the $ref
corner_cases_detected["ref_not_found"] = True
return value
else:
# External ref or other type - keep as is
return value
else:
# Regular dict - process all values
return {
key: resolve_refs_in_value(val, depth, visiting)
for key, val in value.items()
}
elif isinstance(value, list):
# Process each item in the list
return [resolve_refs_in_value(item, depth, visiting) for item in value]
else:
# Primitive value - return as is
return value

# Step 3: Process main schema properties with shared visiting set
for key, value in result.items():
if key != "$defs":
# Each top-level property gets its own visiting set
# This allows the same definition to be used in different contexts
result[key] = resolve_refs_in_value(value, 0, set())

# Step 4: Conditionally preserve $defs based on corner cases
if any(corner_cases_detected.values()):
# Corner case detected, preserve original $defs
if "$defs" in schema: # Only add if original schema had $defs
result["$defs"] = defs
else:
# No corner cases, remove $defs if it exists
result.pop("$defs", None)

return result


class SchemaDereferenceMiddleware(Middleware):
"""Middleware that dereferences $ref in schemas for tools, resource templates.

Applies to list handlers so that clients like Claude Desktop receive flattened schemas
without $ref in properties, preventing null parameter values.
"""

async def on_list_tools(
self,
context: MiddlewareContext[mt.ListToolsRequest],
call_next: CallNext[mt.ListToolsRequest, list],
) -> list:
tools = await call_next(context)
flattened = []
for tool in tools:
params = getattr(tool, "parameters", None)
update: dict[str, Any] = {}
if isinstance(params, dict):
update["parameters"] = dereference_json_schema(params)
if update:
tool = tool.model_copy(update=update)
flattened.append(tool)
return flattened

async def on_list_resource_templates(
self,
context: MiddlewareContext[mt.ListResourceTemplatesRequest],
call_next: CallNext[mt.ListResourceTemplatesRequest, list],
) -> list:
templates = await call_next(context)
flattened = []
for template in templates:
params = getattr(template, "parameters", None)
if isinstance(params, dict):
template = template.model_copy(
update={"parameters": dereference_json_schema(params)}
)
flattened.append(template)
return flattened
60 changes: 60 additions & 0 deletions tests/server/middleware/test_schema_dereference.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,60 @@
import json
from enum import Enum

from fastmcp import Client, FastMCP
from fastmcp.server.middleware.schema_dereference import (
SchemaDereferenceMiddleware,
)


class ColorEnum(str, Enum):
RED = "red"
GREEN = "green"
BLUE = "blue"


class TestSchemaDereferenceMiddleware:
async def test_dereference_enum_in_tool_parameters(self):
mcp = FastMCP("SchemaDereferenceTest")
mcp.add_middleware(SchemaDereferenceMiddleware())

@mcp.tool
def choose_color(color: ColorEnum) -> str:
return color.value

async with Client(mcp) as client:
tools = await client.list_tools()

tool = next(t for t in tools if t.name == "choose_color")
schema = tool.inputSchema

# Ensure $ref was inlined and $defs removed for simple enum case
assert "$ref" not in json.dumps(schema)
assert "$defs" not in schema

assert "properties" in schema and "color" in schema["properties"]
color_schema = schema["properties"]["color"]
assert color_schema.get("enum") == ["red", "green", "blue"]
assert color_schema.get("type") == "string"

async def test_dereference_enum_in_resource_template_parameters(self):
mcp = FastMCP("SchemaDereferenceTemplateTest")
mcp.add_middleware(SchemaDereferenceMiddleware())

@mcp.resource("color://{color}")
def color_resource(color: ColorEnum) -> str:
return color.value

# Use internal list to inspect template parameters after middleware
templates = await mcp._list_resource_templates()
assert len(templates) == 1
params = templates[0].parameters

# Ensure $ref was inlined and $defs removed for simple enum case
assert "$ref" not in json.dumps(params)
assert "$defs" not in params

assert "properties" in params and "color" in params["properties"]
color_schema = params["properties"]["color"]
assert color_schema.get("enum") == ["red", "green", "blue"]
assert color_schema.get("type") == "string"
Loading