Skip to content
Merged
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
9 changes: 7 additions & 2 deletions src/fastmcp/tools/tool.py
Original file line number Diff line number Diff line change
Expand Up @@ -497,7 +497,10 @@ def from_function(
# Compress and handle exclude_args
prune_params = list(exclude_args) if exclude_args else None
input_schema = compress_schema(
input_schema, prune_params=prune_params, prune_titles=True
input_schema,
prune_params=prune_params,
prune_titles=True,
dereference=True,
)

output_schema = None
Expand Down Expand Up @@ -557,7 +560,9 @@ def from_function(
else:
output_schema = base_schema

output_schema = compress_schema(output_schema, prune_titles=True)
output_schema = compress_schema(
output_schema, prune_titles=True, dereference=True
)

# Resolve root-level $ref to meet MCP spec requirement for type: object
# Self-referential Pydantic models generate schemas with $ref at root
Expand Down
29 changes: 29 additions & 0 deletions src/fastmcp/utilities/json_schema.py
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,9 @@ def dereference_refs(schema: dict[str, Any]) -> dict[str, Any]:
if "$defs" in dereferenced:
dereferenced = {k: v for k, v in dereferenced.items() if k != "$defs"}

# Remove discriminator.mapping entries that referenced $defs
_strip_discriminator_mappings(dereferenced)

return dereferenced

except JsonRefError:
Expand All @@ -61,6 +64,26 @@ def dereference_refs(schema: dict[str, Any]) -> dict[str, Any]:
return resolve_root_ref(schema)


def _strip_discriminator_mappings(schema: Any, depth: int = 0) -> None:
"""Remove discriminator.mapping entries whose values are $defs references.

Pydantic emits discriminator.mapping with plain-string references like
``"#/$defs/Cat"`` that become dangling after $defs are removed by
dereference_refs(). The oneOf/anyOf variants already carry their own
const fields, so the mapping is redundant once refs are inlined.
"""
if depth > 50 or not isinstance(schema, dict):
return
if "discriminator" in schema and isinstance(schema["discriminator"], dict):
schema["discriminator"].pop("mapping", None)
for value in schema.values():
if isinstance(value, dict):
_strip_discriminator_mappings(value, depth + 1)
elif isinstance(value, list):
for item in value:
_strip_discriminator_mappings(item, depth + 1)


def _merge_ref_siblings(
original: Any,
dereferenced: Any,
Expand Down Expand Up @@ -367,6 +390,7 @@ def compress_schema(
prune_defs: bool = True,
prune_additional_properties: bool = True,
prune_titles: bool = False,
dereference: bool = False,
) -> dict[str, Any]:
"""
Remove the given parameters from the schema.
Expand All @@ -377,6 +401,7 @@ def compress_schema(
prune_defs: Whether to remove unused definitions
prune_additional_properties: Whether to remove additionalProperties: false
prune_titles: Whether to remove title fields from the schema
dereference: Whether to inline $ref references for client compatibility
"""
# Remove specific parameters if requested
for param in prune_params or []:
Expand All @@ -391,4 +416,8 @@ def compress_schema(
prune_defs=prune_defs,
)

# Inline $ref references for MCP clients that don't handle them
if dereference:
schema = dereference_refs(schema)
Comment on lines +420 to +421
Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Preserve discriminator mappings when inlining $ref schemas

Calling compress_schema(..., dereference=True) now runs dereference_refs, which removes $defs after resolving only $ref nodes; this breaks discriminated-union schemas because Pydantic emits discriminator.mapping values like #/$defs/Cat as plain strings, so those mappings become dangling once $defs is stripped. For tools whose input/output annotations use discriminated unions, clients/validators that rely on discriminator.mapping can no longer resolve variants correctly.

Useful? React with 👍 / 👎.


return schema
9 changes: 3 additions & 6 deletions tests/tools/test_tool.py
Original file line number Diff line number Diff line change
Expand Up @@ -194,18 +194,15 @@ def create_user(user: UserInput, flag: bool) -> dict:
"tags": set(),
"enabled": True,
"parameters": {
"$defs": {
"UserInput": {
"properties": {
"user": {
"properties": {
"name": {"type": "string"},
"age": {"type": "integer"},
},
"required": ["name", "age"],
"type": "object",
}
},
"properties": {
"user": {"$ref": "#/$defs/UserInput"},
},
"flag": {"type": "boolean"},
},
"required": ["user", "flag"],
Expand Down
5 changes: 3 additions & 2 deletions tests/tools/test_tool_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -75,8 +75,9 @@ def create_user(user: UserInput, flag: bool) -> dict:
assert tool is not None
assert tool.name == "create_user"
assert tool.description == "Create a new user."
assert "name" in tool.parameters["$defs"]["UserInput"]["properties"]
assert "age" in tool.parameters["$defs"]["UserInput"]["properties"]
# $refs are dereferenced so UserInput is inlined
assert "name" in tool.parameters["properties"]["user"]["properties"]
assert "age" in tool.parameters["properties"]["user"]["properties"]
assert "flag" in tool.parameters["properties"]

async def test_callable_object(self):
Expand Down
54 changes: 27 additions & 27 deletions tests/tools/test_tool_transform.py
Original file line number Diff line number Diff line change
Expand Up @@ -216,10 +216,14 @@ def tool_with_refs(a: VisibleType, b: HiddenType | None = None) -> int:
schema = new_tool.parameters
# Only 'a' should be visible
assert list(schema["properties"].keys()) == ["a"]
# $defs should only contain VisibleType, not HiddenType
defs = schema.get("$defs", {})
assert "VisibleType" in defs
assert "HiddenType" not in defs
# $refs are dereferenced so VisibleType is inlined, not in $defs
assert "$defs" not in schema
# VisibleType's structure should be inlined into the property
assert schema["properties"]["a"] == {
"properties": {"x": {"type": "integer"}},
"required": ["x"],
"type": "object",
}


async def test_forward_with_argument_mapping(add_tool):
Expand Down Expand Up @@ -427,7 +431,8 @@ def parent_tool(cool_model: CoolModel) -> int:

new_tool = Tool.from_tool(tool)

assert new_tool.parameters["$defs"] == tool.parameters["$defs"]
# $refs are dereferenced, so schemas should match (both inlined)
assert new_tool.parameters == tool.parameters


def test_transform_args_validation_unknown_arg(add_tool):
Expand Down Expand Up @@ -1567,20 +1572,20 @@ def complex_tool(
complex_tool, transform_args={"unused_param": ArgTransform(hide=True)}
)

assert "UnusedType" not in transformed_tool.parameters["$defs"]
# $refs are dereferenced, so no $defs section
assert "$defs" not in transformed_tool.parameters

assert transformed_tool.parameters == snapshot(
{
"type": "object",
"properties": {"used_param": {"$ref": "#/$defs/UsedType"}},
"required": ["used_param"],
"$defs": {
"UsedType": {
"properties": {
"used_param": {
"properties": {"value": {"type": "string"}},
"required": ["value"],
"type": "object",
}
},
"required": ["used_param"],
}
)

Expand Down Expand Up @@ -1610,15 +1615,14 @@ async def transform_function(renamed_input: InputType):
assert transformed.parameters == snapshot(
{
"type": "object",
"properties": {"renamed_input": {"$ref": "#/$defs/InputType"}},
"required": ["renamed_input"],
"$defs": {
"InputType": {
"properties": {
"renamed_input": {
"properties": {"data": {"type": "string"}},
"required": ["data"],
"type": "object",
}
},
"required": ["renamed_input"],
}
)

Expand Down Expand Up @@ -1648,46 +1652,42 @@ def base_tool(param_a: TypeA, param_b: TypeB, param_c: TypeC) -> str:
{
"type": "object",
"properties": {
"param_a": {"$ref": "#/$defs/TypeA"},
"param_b": {"$ref": "#/$defs/TypeB"},
},
"required": IsList("param_b", "param_a", check_order=False),
"$defs": {
"TypeA": {
"param_a": {
"properties": {"a": {"type": "string"}},
"required": ["a"],
"type": "object",
},
"TypeB": {
"param_b": {
"properties": {"b": {"type": "integer"}},
"required": ["b"],
"type": "object",
},
},
"required": IsList("param_b", "param_a", check_order=False),
}
)

assert "TypeA" in transform1.parameters["$defs"]
# $refs are dereferenced, so TypeA is inlined
assert "$defs" not in transform1.parameters

# Second transform: hide param_b
transform2 = Tool.from_tool(
transform1,
transform_args={"param_b": ArgTransform(hide=True, default=TypeB(b=42))},
)

assert "TypeB" not in transform2.parameters["$defs"]
assert "$defs" not in transform2.parameters

assert transform2.parameters == snapshot(
{
"type": "object",
"properties": {"param_a": {"$ref": "#/$defs/TypeA"}},
"required": ["param_a"],
"$defs": {
"TypeA": {
"properties": {
"param_a": {
"properties": {"a": {"type": "string"}},
"required": ["a"],
"type": "object",
}
},
"required": ["param_a"],
}
)