Fix schema wrapping for non-object union types in output schemas

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,34 +31,22 @@ if TYPE_CHECKING:
 
 logger = get_logger(__name__)
 
-
-class _UnserializableType:
-    pass
+T = TypeVar("T")
 
 
-def default_serializer(data: Any) -> str:
-    return pydantic_core.to_json(data, fallback=str, indent=2).decode()
+@dataclass
+class _WrappedResult(Generic[T]):
+    """Generic wrapper for non-object return types."""
 
+    result: T
 
-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.
+class _UnserializableType:
+    pass
 
-    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
+def default_serializer(data: Any) -> str:
+    return pydantic_core.to_json(data, fallback=str, indent=2).decode()
 
 
 class ToolResult:
@@ -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,

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

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

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."""
 

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