From 7e0a237945ca44683c3473836fc1937f2ef4e634 Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Mon, 30 Jun 2025 14:12:40 -0400 Subject: [PATCH] Fix schema wrapping for non-object union types in output schemas --- src/fastmcp/tools/tool.py | 74 ++++++++++-------------- src/fastmcp/tools/tool_transform.py | 4 +- tests/server/test_server_interactions.py | 17 ++++-- tests/tools/test_tool.py | 58 ++++++++++++++++--- tests/tools/test_tool_transform.py | 8 ++- 5 files changed, 101 insertions(+), 60 deletions(-) diff --git a/src/fastmcp/tools/tool.py b/src/fastmcp/tools/tool.py index 03cab7d3c..6cca87919 100644 --- a/src/fastmcp/tools/tool.py +++ b/src/fastmcp/tools/tool.py @@ -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, diff --git a/src/fastmcp/tools/tool_transform.py b/src/fastmcp/tools/tool_transform.py index 61d9ecc9e..95c4a877e 100644 --- a/src/fastmcp/tools/tool_transform.py +++ b/src/fastmcp/tools/tool_transform.py @@ -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 diff --git a/tests/server/test_server_interactions.py b/tests/server/test_server_interactions.py index f37434194..8d905919e 100644 --- a/tests/server/test_server_interactions.py +++ b/tests/server/test_server_interactions.py @@ -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 diff --git a/tests/tools/test_tool.py b/tests/tools/test_tool.py index a2278d6d1..119d52f38 100644 --- a/tests/tools/test_tool.py +++ b/tests/tools/test_tool.py @@ -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.""" diff --git a/tests/tools/test_tool_transform.py b/tests/tools/test_tool_transform.py index 5e976b381..f0dedf4e0 100644 --- a/tests/tools/test_tool_transform.py +++ b/tests/tools/test_tool_transform.py @@ -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