ceos/fastmcp
1
0
Fork 0
mirror of https://github.com/PrefectHQ/fastmcp.git synced 2026-08-09 15:19:10 +02:00
Code Issues Projects Releases Packages Wiki Activity

Merge pull request #995 from jlowin/fix-union-type-schema-generation

Browse source 
Fix output schema generation edge case
This commit is contained in:
Jeremiah Lowin 2025-06-30 14:18:30 -04:00 • committed by GitHub
commit d7969b159e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 101 additions and 60 deletions
Download patch file Download diff file Expand all files Collapse all files

74
src/fastmcp/tools/tool.py
View file

@ -3,7 +3,7 @@ from __future__ import annotations
import inspect
from collections.abc import Callable
from dataclasses import dataclass
from typing import TYPE_CHECKING, Annotated, Any, Literal
from typing import TYPE_CHECKING, Annotated, Any, Generic, Literal, TypeVar
import mcp.types
import pydantic_core
@ -31,6 +31,15 @@ if TYPE_CHECKING:
logger = get_logger(__name__)
T = TypeVar("T")
@dataclass
class _WrappedResult(Generic[T]):
"""Generic wrapper for non-object return types."""
result: T
class _UnserializableType:
pass
@ -40,27 +49,6 @@ def default_serializer(data: Any) -> str:
return pydantic_core.to_json(data, fallback=str, indent=2).decode()
def _wrap_schema_if_needed(schema: dict[str, Any] | None) -> dict[str, Any] | None:
"""Wrap non-object schemas with result property for structured output.
This wrapping allows primitive types (int, str, etc.) to be returned as
structured content by placing them under a "result" key.
Args:
schema: The JSON schema to potentially wrap
Returns:
Wrapped schema if needed, or original schema if already an object type
"""
if schema and schema.get("type") != "object":
return {
"type": "object",
"properties": {"result": schema},
"x-fastmcp-wrap-result": True,
}
return schema
class ToolResult:
def __init__(
self,
@ -246,7 +234,7 @@ class FunctionTool(Tool):
raise ValueError("You must provide a name for lambda functions")
if isinstance(output_schema, NotSetT):
output_schema = _wrap_schema_if_needed(parsed_fn.output_schema)
output_schema = parsed_fn.output_schema
elif output_schema is False:
output_schema = None
# Note: explicit schemas (dict) are used as-is without auto-wrapping
@ -329,8 +317,8 @@ class ParsedFunction:
cls,
fn: Callable[..., Any],
exclude_args: list[str] | None = None,
ignore_response_types: list[type] | None = None,
validate: bool = True,
wrap_non_object_output_schema: bool = True,
) -> ParsedFunction:
from fastmcp.server.context import Context
@ -389,7 +377,7 @@ class ParsedFunction:
# or are MCP content types that explicitly don't form structured
# content. By replacing them with an explicitly unserializable type,
# we ensure that no output schema is automatically generated.
output_type = replace_type(
clean_output_type = replace_type(
output_type,
{
t: _UnserializableType
@ -408,8 +396,25 @@ class ParsedFunction:
)
try:
output_type_adapter = get_cached_typeadapter(output_type)
output_schema = output_type_adapter.json_schema()
type_adapter = get_cached_typeadapter(clean_output_type)
base_schema = type_adapter.json_schema()
# Generate schema for wrapped type if it's non-object
# because MCP requires that output schemas are objects
if (
wrap_non_object_output_schema
and base_schema.get("type") != "object"
):
# Use the wrapped result schema directly
wrapped_type = _WrappedResult[clean_output_type]
wrapped_adapter = get_cached_typeadapter(wrapped_type)
output_schema = wrapped_adapter.json_schema()
output_schema["x-fastmcp-wrap-result"] = True
else:
output_schema = base_schema
output_schema = compress_schema(output_schema)
except PydanticSchemaGenerationError as e:
if "_UnserializableType" not in str(e):
logger.debug(f"Unable to generate schema for type {output_type!r}")
@ -422,21 +427,6 @@ class ParsedFunction:
output_schema=output_schema or None,
)
try:
output_type_adapter = get_cached_typeadapter(output_type)
output_schema = output_type_adapter.json_schema()
except PydanticSchemaGenerationError as e:
if "_UnserializableType" not in str(e):
logger.debug(f"Unable to generate schema for type {output_type!r}")
return cls(
fn=fn,
name=fn_name,
description=fn_doc,
input_schema=input_schema,
output_schema=output_schema or None,
)
def _convert_to_content(
result: Any,

4
src/fastmcp/tools/tool_transform.py
View file

@ -9,7 +9,7 @@ from typing import Any, Literal
from mcp.types import ToolAnnotations
from pydantic import ConfigDict
from fastmcp.tools.tool import ParsedFunction, Tool, ToolResult, _wrap_schema_if_needed
from fastmcp.tools.tool import ParsedFunction, Tool, ToolResult
from fastmcp.utilities.logging import get_logger
from fastmcp.utilities.types import NotSet, NotSetT, get_cached_typeadapter
@ -430,7 +430,7 @@ class TransformedTool(Tool):
# Smart fallback: try custom function, then parent, then None
if transform_fn is not None:
parsed_fn = ParsedFunction.from_function(transform_fn, validate=False)
final_output_schema = _wrap_schema_if_needed(parsed_fn.output_schema)
final_output_schema = parsed_fn.output_schema
if final_output_schema is None:
# Check if function returns ToolResult - if so, don't fall back to parent
import inspect

17
tests/server/test_server_interactions.py
View file

@ -27,6 +27,7 @@ from fastmcp.prompts.prompt import Prompt, PromptMessage
from fastmcp.resources import FileResource, ResourceTemplate
from fastmcp.resources.resource import FunctionResource
from fastmcp.tools.tool import Tool, ToolResult
from fastmcp.utilities.json_schema import compress_schema
from fastmcp.utilities.types import Audio, File, Image
@ -894,7 +895,9 @@ class TestToolOutputSchema:
# this line will fail until MCP adds output schemas!!
assert tools[0].outputSchema == {
"type": "object",
"properties": {"result": type_schema},
"properties": {"result": {**type_schema, "title": "Result"}},
"required": ["result"],
"title": "_WrappedResult",
"x-fastmcp-wrap-result": True,
}
@ -912,7 +915,7 @@ class TestToolOutputSchema:
async with Client(mcp) as client:
tools = await client.list_tools()
type_schema = TypeAdapter(annotation).json_schema()
type_schema = compress_schema(TypeAdapter(annotation).json_schema())
assert len(tools) == 1
assert tools[0].outputSchema == type_schema
@ -1020,7 +1023,9 @@ class TestToolOutputSchema:
tool = next(t for t in tools if t.name == "primitive_tool")
expected_schema = {
"type": "object",
"properties": {"result": {"type": "string"}},
"properties": {"result": {"type": "string", "title": "Result"}},
"required": ["result"],
"title": "_WrappedResult",
"x-fastmcp-wrap-result": True,
}
assert tool.outputSchema == expected_schema
@ -1045,7 +1050,9 @@ class TestToolOutputSchema:
expected_inner_schema = TypeAdapter(list[dict[str, int]]).json_schema()
expected_schema = {
"type": "object",
"properties": {"result": expected_inner_schema},
"properties": {"result": {**expected_inner_schema, "title": "Result"}},
"required": ["result"],
"title": "_WrappedResult",
"x-fastmcp-wrap-result": True,
}
assert tool.outputSchema == expected_schema
@ -1074,7 +1081,7 @@ class TestToolOutputSchema:
# List tools and verify schema is object type (not wrapped)
tools = await client.list_tools()
tool = next(t for t in tools if t.name == "dataclass_tool")
expected_schema = TypeAdapter(User).json_schema()
expected_schema = compress_schema(TypeAdapter(User).json_schema())
assert tool.outputSchema == expected_schema
assert (
tool.outputSchema and "x-fastmcp-wrap-result" not in tool.outputSchema

58
tests/tools/test_tool.py
View file

@ -14,6 +14,7 @@ from pydantic import AnyUrl, BaseModel, Field, TypeAdapter
from typing_extensions import TypedDict
from fastmcp.tools.tool import Tool, _convert_to_content
from fastmcp.utilities.json_schema import compress_schema
from fastmcp.utilities.types import Audio, File, Image
@ -35,7 +36,9 @@ class TestToolFromFunction:
# With primitive wrapping, int return type becomes object with result property
expected_schema = {
"type": "object",
"properties": {"result": {"type": "integer"}},
"properties": {"result": {"type": "integer", "title": "Result"}},
"required": ["result"],
"title": "_WrappedResult",
"x-fastmcp-wrap-result": True,
}
assert tool.output_schema == expected_schema
@ -296,7 +299,9 @@ class TestToolFromFunctionOutputSchema:
# Non-object types get wrapped
expected_schema = {
"type": "object",
"properties": {"result": base_schema},
"properties": {"result": {**base_schema, "title": "Result"}},
"required": ["result"],
"title": "_WrappedResult",
"x-fastmcp-wrap-result": True,
}
assert tool.output_schema == expected_schema
@ -321,7 +326,9 @@ class TestToolFromFunctionOutputSchema:
expected_schema = {
"type": "object",
"properties": {"result": base_schema},
"properties": {"result": {**base_schema, "title": "Result"}},
"required": ["result"],
"title": "_WrappedResult",
"x-fastmcp-wrap-result": True,
}
assert tool.output_schema == expected_schema
@ -369,7 +376,8 @@ class TestToolFromFunctionOutputSchema:
return Person(name="John", age=30)
tool = Tool.from_function(func)
assert tool.output_schema == TypeAdapter(Person).json_schema()
expected_schema = compress_schema(TypeAdapter(Person).json_schema())
assert tool.output_schema == expected_schema
async def test_base_model_return_annotation(self):
class Person(BaseModel):
@ -380,7 +388,8 @@ class TestToolFromFunctionOutputSchema:
return Person(name="John", age=30)
tool = Tool.from_function(func)
assert tool.output_schema == TypeAdapter(Person).json_schema()
expected_schema = compress_schema(TypeAdapter(Person).json_schema())
assert tool.output_schema == expected_schema
async def test_typeddict_return_annotation(self):
class Person(TypedDict):
@ -391,7 +400,8 @@ class TestToolFromFunctionOutputSchema:
return Person(name="John", age=30)
tool = Tool.from_function(func)
assert tool.output_schema == TypeAdapter(Person).json_schema()
expected_schema = compress_schema(TypeAdapter(Person).json_schema())
assert tool.output_schema == expected_schema
async def test_unserializable_return_annotation(self):
class Unserializable:
@ -568,7 +578,9 @@ class TestToolFromFunctionOutputSchema:
tool = Tool.from_function(func)
expected_schema = {
"type": "object",
"properties": {"result": {"type": "integer"}},
"properties": {"result": {"type": "integer", "title": "Result"}},
"required": ["result"],
"title": "_WrappedResult",
"x-fastmcp-wrap-result": True,
}
assert tool.output_schema == expected_schema
@ -638,7 +650,9 @@ class TestToolFromFunctionOutputSchema:
tool = Tool.from_function(func)
expected_schema = {
"type": "object",
"properties": {"result": {"type": "string"}},
"properties": {"result": {"type": "string", "title": "Result"}},
"required": ["result"],
"title": "_WrappedResult",
"x-fastmcp-wrap-result": True,
}
assert tool.output_schema == expected_schema
@ -1220,13 +1234,39 @@ class TestAutomaticStructuredContent:
async with Client(mcp) as client:
result = await client.call_tool("get_profile", {"user_id": "456"})
# Client should deserialize back to a dataclass (type name will match)
# Client should deserialize back to a dataclass (type name preserved with new compression)
assert result.data.__class__.__name__ == "UserProfile"
assert result.data.name == "Bob"
assert result.data.age == 25
assert result.data.verified is True
class TestUnionReturnTypes:
"""Tests for tools with union return types."""
async def test_dataclass_union_string_works(self):
"""Test that union of dataclass and string works correctly."""
@dataclass
class Data:
value: int
def get_data(return_error: bool) -> Data | str:
if return_error:
return "error occurred"
return Data(value=42)
tool = Tool.from_function(get_data)
# Test returning dataclass
result1 = await tool.run({"return_error": False})
assert result1.structured_content == {"result": {"value": 42}}
# Test returning string
result2 = await tool.run({"return_error": True})
assert result2.structured_content == {"result": "error occurred"}
class TestToolTitle:
"""Tests for tool title functionality."""

8
tests/tools/test_tool_transform.py
View file

@ -1063,7 +1063,9 @@ class TestTransformToolOutputSchema:
# Should inherit parent's wrapped string schema
expected_schema = {
"type": "object",
"properties": {"result": {"type": "string"}},
"properties": {"result": {"type": "string", "title": "Result"}},
"required": ["result"],
"title": "_WrappedResult",
"x-fastmcp-wrap-result": True,
}
assert new_tool.output_schema == expected_schema
@ -1121,7 +1123,9 @@ class TestTransformToolOutputSchema:
# Should infer string schema from custom function and wrap it
expected_schema = {
"type": "object",
"properties": {"result": {"type": "string"}},
"properties": {"result": {"type": "string", "title": "Result"}},
"required": ["result"],
"title": "_WrappedResult",
"x-fastmcp-wrap-result": True,
}
assert new_tool.output_schema == expected_schema