Update output schema control

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
+from typing import TYPE_CHECKING, Annotated, Any, Literal
 
 import mcp.types
 import pydantic_core
@@ -40,6 +40,27 @@ 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,
@@ -64,7 +85,11 @@ class ToolResult:
                 )
                 raise
             if not isinstance(structured_content, dict):
-                structured_content = {"result": structured_content}
+                raise ValueError(
+                    "structured_content must be a dict or None. "
+                    f"Got {type(structured_content).__name__}: {structured_content!r}. "
+                    "Tools should wrap non-dict values based on their output_schema."
+                )
         self.structured_content: dict[str, Any] | None = structured_content
 
     def to_mcp_result(
@@ -127,7 +152,7 @@ class Tool(FastMCPComponent):
         tags: set[str] | None = None,
         annotations: ToolAnnotations | None = None,
         exclude_args: list[str] | None = None,
-        output_schema: dict[str, Any] | None | NotSetT = NotSet,
+        output_schema: dict[str, Any] | None | NotSetT | Literal[False] = NotSet,
         serializer: Callable[[Any], str] | None = None,
         enabled: bool | None = None,
     ) -> FunctionTool:
@@ -166,6 +191,7 @@ class Tool(FastMCPComponent):
         description: str | None = None,
         tags: set[str] | None = None,
         annotations: ToolAnnotations | None = None,
+        output_schema: dict[str, Any] | None | Literal[False] = None,
         serializer: Callable[[Any], str] | None = None,
         enabled: bool | None = None,
     ) -> TransformedTool:
@@ -179,6 +205,7 @@ class Tool(FastMCPComponent):
             description=description,
             tags=tags,
             annotations=annotations,
+            output_schema=output_schema,
             serializer=serializer,
             enabled=enabled,
         )
@@ -196,7 +223,7 @@ class FunctionTool(Tool):
         tags: set[str] | None = None,
         annotations: ToolAnnotations | None = None,
         exclude_args: list[str] | None = None,
-        output_schema: dict[str, Any] | None | NotSetT = NotSet,
+        output_schema: dict[str, Any] | None | NotSetT | Literal[False] = NotSet,
         serializer: Callable[[Any], str] | None = None,
         enabled: bool | None = None,
     ) -> FunctionTool:
@@ -208,14 +235,10 @@ class FunctionTool(Tool):
             raise ValueError("You must provide a name for lambda functions")
 
         if isinstance(output_schema, NotSetT):
-            output_schema = parsed_fn.output_schema
-
-            if output_schema and output_schema.get("type") != "object":
-                output_schema = {
-                    "type": "object",
-                    "properties": {"result": output_schema},
-                    "x-fastmcp-wrap-result": True,
-                }
+            output_schema = _wrap_schema_if_needed(parsed_fn.output_schema)
+        elif output_schema is False:
+            output_schema = None
+        # Note: explicit schemas (dict) are used as-is without auto-wrapping
 
         return cls(
             fn=parsed_fn.fn,
@@ -249,6 +272,7 @@ class FunctionTool(Tool):
 
         unstructured_result = _convert_to_content(result, serializer=self.serializer)
 
+        # Handle structured content based on output schema
         if self.output_schema is not None:
             if self.output_schema.get("x-fastmcp-wrap-result"):
                 structured_output = {"result": result}

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
+from fastmcp.tools.tool import ParsedFunction, Tool, ToolResult, _wrap_schema_if_needed
 from fastmcp.utilities.logging import get_logger
 from fastmcp.utilities.types import NotSet, NotSetT, get_cached_typeadapter
 
@@ -52,7 +52,7 @@ async def forward(**kwargs) -> ToolResult:
     return await tool.forwarding_fn(**kwargs)
 
 
-async def forward_raw(**kwargs) -> Any:
+async def forward_raw(**kwargs) -> ToolResult:
     """Forward directly to parent tool without transformation.
 
     This function bypasses all argument transformation and validation, calling the parent
@@ -66,7 +66,7 @@ async def forward_raw(**kwargs) -> Any:
         **kwargs: Arguments to pass directly to the parent tool (using original names).
 
     Returns:
-        The result from the parent tool execution.
+        The ToolResult from the parent tool execution.
 
     Raises:
         RuntimeError: If called outside a transformed tool context.
@@ -268,15 +268,31 @@ class TransformedTool(Tool):
         token = _current_tool.set(self)
         try:
             result = await self.fn(**arguments)
-            
+
             # If transform function returns ToolResult, use it directly
             if isinstance(result, ToolResult):
                 return result
-            
-            # Otherwise convert to content and create basic ToolResult
+
+            # Otherwise convert to content and create ToolResult with proper structured content
             from fastmcp.tools.tool import _convert_to_content
-            unstructured_result = _convert_to_content(result, serializer=self.serializer)
-            return ToolResult(content=unstructured_result)
+
+            unstructured_result = _convert_to_content(
+                result, serializer=self.serializer
+            )
+
+            # Handle structured content based on output schema
+            if self.output_schema is not None:
+                if self.output_schema.get("x-fastmcp-wrap-result"):
+                    structured_output = {"result": result}
+                else:
+                    structured_output = result
+            else:
+                structured_output = None
+
+            return ToolResult(
+                content=unstructured_result,
+                structured_content=structured_output,
+            )
         finally:
             _current_tool.reset(token)
 
@@ -290,6 +306,7 @@ class TransformedTool(Tool):
         transform_fn: Callable[..., Any] | None = None,
         transform_args: dict[str, ArgTransform] | None = None,
         annotations: ToolAnnotations | None = None,
+        output_schema: dict[str, Any] | None | Literal[False] = None,
         serializer: Callable[[Any], str] | None = None,
         enabled: bool | None = None,
     ) -> TransformedTool:
@@ -352,13 +369,30 @@ class TransformedTool(Tool):
         # Always create the forwarding transform
         schema, forwarding_fn = cls._create_forwarding_transform(tool, transform_args)
 
+        # Handle output schema with smart fallback
+        if output_schema is False:
+            final_output_schema = None
+        elif output_schema is not None:
+            # Explicit schema provided - use as-is
+            final_output_schema = output_schema
+        else:
+            # 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)
+                if final_output_schema is None:
+                    final_output_schema = tool.output_schema
+            else:
+                final_output_schema = tool.output_schema
+
         if transform_fn is None:
             # User wants pure transformation - use forwarding_fn as the main function
             final_fn = forwarding_fn
             final_schema = schema
         else:
             # User provided custom function - merge schemas
-            parsed_fn = ParsedFunction.from_function(transform_fn, validate=False)
+            if "parsed_fn" not in locals():
+                parsed_fn = ParsedFunction.from_function(transform_fn, validate=False)
             final_fn = transform_fn
 
             has_kwargs = cls._function_has_kwargs(transform_fn)
@@ -426,6 +460,7 @@ class TransformedTool(Tool):
             name=name or tool.name,
             description=final_description,
             parameters=final_schema,
+            output_schema=final_output_schema,
             tags=tags or tool.tags,
             annotations=annotations or tool.annotations,
             serializer=serializer or tool.serializer,

tests/tools/test_tool.py

@@ -31,12 +31,10 @@ class TestToolFromFunction:
         assert len(tool.parameters["properties"]) == 2
         assert tool.parameters["properties"]["a"]["type"] == "integer"
         assert tool.parameters["properties"]["b"]["type"] == "integer"
-        # With primitive wrapping, int return type becomes object with value property
+        # With primitive wrapping, int return type becomes object with result property
         expected_schema = {
             "type": "object",
-            "properties": {"value": {"title": "Value", "type": "integer"}},
-            "required": ["value"],
-            "title": "Result",
+            "properties": {"result": {"type": "integer"}},
             "x-fastmcp-wrap-result": True,
         }
         assert tool.output_schema == expected_schema
@@ -251,7 +249,7 @@ class TestToolFromFunction:
         assert isinstance(result.content[0], TextContent)
         assert result.content[0].text == "Custom serializer: 15"
         # Structured output should have the raw value
-        assert result.structured_content == {"value": 15}
+        assert result.structured_content == {"result": 15}
 
 
 class TestToolFromFunctionOutputSchema:
@@ -287,27 +285,20 @@ class TestToolFromFunctionOutputSchema:
 
         base_schema = TypeAdapter(annotation).json_schema()
 
-        # Only pure primitives (just type + optional title) get wrapped
-        primitive_types = {"string", "number", "integer", "boolean", "null"}
+        # Non-object types get wrapped
         schema_type = base_schema.get("type")
-        is_pure_primitive = (
-            schema_type in primitive_types
-            and len(base_schema) <= 2  # Only 'type' and optionally 'title'
-            and all(key in {"type", "title"} for key in base_schema.keys())
-        )
+        is_object_type = schema_type == "object"
 
-        if is_pure_primitive:
-            # Pure primitives get wrapped
+        if not is_object_type:
+            # Non-object types get wrapped
             expected_schema = {
                 "type": "object",
-                "properties": {"value": base_schema | {"title": "Value"}},
-                "required": ["value"],
-                "title": "Result",
+                "properties": {"result": base_schema},
                 "x-fastmcp-wrap-result": True,
             }
             assert tool.output_schema == expected_schema
         else:
-            # Complex types (objects, unions, constrained types) remain unwrapped
+            # Object types remain unwrapped
             assert tool.output_schema == base_schema
 
     @pytest.mark.parametrize(
@@ -326,8 +317,17 @@ class TestToolFromFunctionOutputSchema:
         tool = Tool.from_function(func)
         base_schema = TypeAdapter(annotation).json_schema()
 
-        # Complex types with constraints are not wrapped - they remain as-is
-        assert tool.output_schema == base_schema
+        # Special case for Any type - it generates an empty schema and doesn't get wrapped
+        if annotation is Any:
+            assert tool.output_schema == base_schema  # Should be {}
+        else:
+            # All other non-object types get wrapped, including complex constrained types
+            expected_schema = {
+                "type": "object",
+                "properties": {"result": base_schema},
+                "x-fastmcp-wrap-result": True,
+            }
+            assert tool.output_schema == expected_schema
 
     @pytest.mark.parametrize(
         "annotation, expected",

tests/tools/test_tool_transform.py

@@ -255,7 +255,7 @@ async def test_forward_raw_without_argument_mapping(add_tool):
 async def test_custom_fn_with_kwargs_and_no_transform_args(add_tool):
     async def custom_fn(extra: int, **kwargs) -> int:
         sum = await forward(**kwargs)
-        return int(sum[0].text) + extra  # type: ignore[attr-defined]
+        return int(sum.content[0].text) + extra  # type: ignore[attr-defined]
 
     new_tool = Tool.from_tool(add_tool, transform_fn=custom_fn)
     result = await new_tool.run(arguments={"extra": 1, "old_x": 2, "old_y": 3})
@@ -360,7 +360,7 @@ async def test_fn_with_kwargs_dropped_args_not_in_kwargs(add_tool):
     )  # drop 'old_y'
     result = await new_tool.run(arguments={"new_x": 8})
     # 8 + 10 (default value of b in parent)
-    assert result[0].text == "18"  # type: ignore[attr-defined]
+    assert result.content[0].text == "18"  # type: ignore[attr-defined]
 
 
 async def test_forward_outside_context_raises_error():
@@ -480,18 +480,18 @@ async def test_tool_transform_chaining(add_tool):
     tool2 = Tool.from_tool(tool1, transform_args={"x": ArgTransform(name="final_x")})
 
     result = await tool2.run(arguments={"final_x": 5})
-    assert result[0].text == "15"  # type: ignore[attr-defined]
+    assert result.content[0].text == "15"  # type: ignore[attr-defined]
 
     # Transform tool1 with custom function that handles all parameters
     async def custom(final_x: int, **kwargs) -> str:
         result = await forward(final_x=final_x, **kwargs)
-        return f"custom {result[0].text}"  # Extract text from content
+        return f"custom {result.content[0].text}"  # Extract text from content
 
     tool3 = Tool.from_tool(
         tool1, transform_fn=custom, transform_args={"x": ArgTransform(name="final_x")}
     )
     result = await tool3.run(arguments={"final_x": 3, "old_y": 5})
-    assert result[0].text == "custom 8"  # type: ignore[attr-defined]
+    assert result.content[0].text == "custom 8"  # type: ignore[attr-defined]
 
 
 class MyModel(BaseModel):
@@ -619,7 +619,7 @@ async def test_arg_transform_precedence_over_function_with_kwargs():
     # Function signature has different types/defaults than ArgTransform
     async def custom_fn(x: str = "function_default", **kwargs) -> str:
         result = await forward(x=x, **kwargs)
-        return f"custom: {result}"
+        return f"custom: {result.content[0].text}"
 
     tool = Tool.from_tool(
         base,
@@ -646,7 +646,7 @@ async def test_arg_transform_precedence_over_function_with_kwargs():
     # Test it works at runtime
     result = await tool.run(arguments={"y": "test"})
     # Should use ArgTransform default of 42
-    assert "42: test" in result[0].text  # type: ignore[attr-defined]
+    assert "42: test" in result.content[0].text  # type: ignore[attr-defined]
 
 
 def test_arg_transform_combined_attributes():
@@ -691,7 +691,7 @@ async def test_arg_transform_type_precedence_runtime():
         # Convert string back to int for the original function
         result = await forward_raw(x=int(x), y=y)
         # Extract the text from the result
-        result_text = result[0].text
+        result_text = result.content[0].text
         return f"String input '{x}' converted to result: {result_text}"
 
     tool = Tool.from_tool(
@@ -703,8 +703,8 @@ async def test_arg_transform_type_precedence_runtime():
 
     # Test it works with string input
     result = await tool.run(arguments={"x": "5", "y": 3})
-    assert "String input '5'" in result[0].text  # type: ignore[attr-defined]
-    assert "result: 8" in result[0].text  # type: ignore[attr-defined]
+    assert "String input '5'" in result.content[0].text  # type: ignore[attr-defined]
+    assert "result: 8" in result.content[0].text  # type: ignore[attr-defined]
 
 
 class TestProxy:
@@ -739,7 +739,7 @@ class TestProxy:
         async with Client(proxy_server) as client:
             # The tool should be registered with its transformed name
             result = await client.call_tool("add_transformed", {"new_x": 1, "old_y": 2})
-            assert result[0].text == "3"  # type: ignore[attr-defined]
+            assert result.content[0].text == "3"  # type: ignore[attr-defined]
 
 
 async def test_arg_transform_default_factory():
@@ -762,7 +762,7 @@ async def test_arg_transform_default_factory():
 
     # Should work without providing timestamp (gets value from factory)
     result = await new_tool.run(arguments={"x": 42})
-    assert result[0].text == "42_12345.0"  # type: ignore[attr-defined]
+    assert result.content[0].text == "42_12345.0"  # type: ignore[attr-defined]
 
 
 async def test_arg_transform_default_factory_called_each_time():
@@ -790,11 +790,11 @@ async def test_arg_transform_default_factory_called_each_time():
 
     # First call
     result1 = await new_tool.run(arguments={"x": 1})
-    assert result1[0].text == "1_1"  # type: ignore[attr-defined]
+    assert result1.content[0].text == "1_1"  # type: ignore[attr-defined]
 
     # Second call should get a different value
     result2 = await new_tool.run(arguments={"x": 2})
-    assert result2[0].text == "2_2"  # type: ignore[attr-defined]
+    assert result2.content[0].text == "2_2"  # type: ignore[attr-defined]
 
 
 async def test_arg_transform_hidden_with_default_factory():
@@ -819,7 +819,7 @@ async def test_arg_transform_hidden_with_default_factory():
 
     # Should pass hidden request_id with factory value
     result = await new_tool.run(arguments={"x": 42})
-    assert result[0].text == "42_req_123"  # type: ignore[attr-defined]
+    assert result.content[0].text == "42_req_123"  # type: ignore[attr-defined]
 
 
 async def test_arg_transform_default_and_factory_raises_error():
@@ -856,7 +856,7 @@ async def test_arg_transform_required_true():
 
     # Should work when parameter is provided
     result = await new_tool.run(arguments={"optional_param": 100})
-    assert result[0].text == "value: 100"  # type: ignore
+    assert result.content[0].text == "value: 100"  # type: ignore
 
     # Should fail when parameter is not provided
     with pytest.raises(TypeError, match="Missing required argument"):
@@ -903,7 +903,7 @@ async def test_arg_transform_required_with_rename():
 
     # Should work with new name
     result = await new_tool.run(arguments={"new_param": 200})
-    assert result[0].text == "value: 200"  # type: ignore
+    assert result.content[0].text == "value: 200"  # type: ignore
 
 
 async def test_arg_transform_required_true_with_default_raises_error():
@@ -945,7 +945,7 @@ async def test_arg_transform_required_no_change():
 
     # Should work as expected
     result = await new_tool.run(arguments={"req": 1})
-    assert result[0].text == "values: 1, 42"  # type: ignore
+    assert result.content[0].text == "values: 1, 42"  # type: ignore
 
 
 async def test_arg_transform_hide_and_required_raises_error():
@@ -977,7 +977,7 @@ class TestEnableDisable:
             assert {tool.name for tool in tools} == {"new_add"}
 
             result = await client.call_tool("new_add", {"x": 1, "y": 2})
-            assert result[0].text == "3"  # type: ignore[attr-defined]
+            assert result.content[0].text == "3"  # type: ignore[attr-defined]
 
             with pytest.raises(ToolError):
                 await client.call_tool("add", {"x": 1, "y": 2})