Add output schema to tools

src/fastmcp/tools/tool.py

@@ -4,12 +4,13 @@ import inspect
 import json
 from collections.abc import Callable
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Annotated, Any
 
+import mcp.types
 import pydantic_core
 from mcp.types import ContentBlock, TextContent, ToolAnnotations
 from mcp.types import Tool as MCPTool
-from pydantic import Field
+from pydantic import Field, PydanticSchemaGenerationError
 
 import fastmcp
 from fastmcp.server.dependencies import get_context
@@ -20,8 +21,11 @@ from fastmcp.utilities.types import (
     Audio,
     File,
     Image,
+    NotSet,
+    NotSetT,
     find_kwarg_by_type,
     get_cached_typeadapter,
+    replace_type,
 )
 
 if TYPE_CHECKING:
@@ -37,19 +41,27 @@ def default_serializer(data: Any) -> str:
 class Tool(FastMCPComponent):
     """Internal tool registration info."""
 
-    parameters: dict[str, Any] = Field(description="JSON schema for tool parameters")
-    annotations: ToolAnnotations | None = Field(
-        default=None, description="Additional annotations about the tool"
-    )
-    serializer: Callable[[Any], str] | None = Field(
-        default=None, description="Optional custom serializer for tool results"
-    )
+    parameters: Annotated[
+        dict[str, Any], Field(description="JSON schema for tool parameters")
+    ]
+    output_schema: Annotated[
+        dict[str, Any] | None, Field(description="JSON schema for tool output")
+    ] = None
+    annotations: Annotated[
+        ToolAnnotations | None,
+        Field(description="Additional annotations about the tool"),
+    ] = None
+    serializer: Annotated[
+        Callable[[Any], str] | None,
+        Field(description="Optional custom serializer for tool results"),
+    ] = None
 
     def to_mcp_tool(self, **overrides: Any) -> MCPTool:
         kwargs = {
             "name": self.name,
             "description": self.description,
             "inputSchema": self.parameters,
+            "outputSchema": self.output_schema,
             "annotations": self.annotations,
         }
         return MCPTool(**kwargs | overrides)
@@ -62,6 +74,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,
         serializer: Callable[[Any], str] | None = None,
         enabled: bool | None = None,
     ) -> FunctionTool:
@@ -73,6 +86,7 @@ class Tool(FastMCPComponent):
             tags=tags,
             annotations=annotations,
             exclude_args=exclude_args,
+            output_schema=output_schema,
             serializer=serializer,
             enabled=enabled,
         )
@@ -121,6 +135,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,
         serializer: Callable[[Any], str] | None = None,
         enabled: bool | None = None,
     ) -> FunctionTool:
@@ -131,13 +146,17 @@ class FunctionTool(Tool):
         if name is None and parsed_fn.name == "<lambda>":
             raise ValueError("You must provide a name for lambda functions")
 
+        if isinstance(output_schema, NotSetT):
+            output_schema = parsed_fn.output_schema
+
         return cls(
             fn=parsed_fn.fn,
             name=name or parsed_fn.name,
             description=description or parsed_fn.description,
-            parameters=parsed_fn.parameters,
-            tags=tags or set(),
+            parameters=parsed_fn.input_schema,
+            output_schema=output_schema,
             annotations=annotations,
+            tags=tags or set(),
             serializer=serializer,
             enabled=enabled if enabled is not None else True,
         )
@@ -194,7 +213,8 @@ class ParsedFunction:
     fn: Callable[..., Any]
     name: str
     description: str | None
-    parameters: dict[str, Any]
+    input_schema: dict[str, Any]
+    output_schema: dict[str, Any] | None
 
     @classmethod
     def from_function(
@@ -240,9 +260,6 @@ class ParsedFunction:
         if isinstance(fn, staticmethod):
             fn = fn.__func__
 
-        type_adapter = get_cached_typeadapter(fn)
-        schema = type_adapter.json_schema()
-
         prune_params: list[str] = []
         context_kwarg = find_kwarg_by_type(fn, kwarg_type=Context)
         if context_kwarg:
@@ -250,12 +267,33 @@ class ParsedFunction:
         if exclude_args:
             prune_params.extend(exclude_args)
 
-        schema = compress_schema(schema, prune_params=prune_params)
+        input_type_adapter = get_cached_typeadapter(fn)
+        input_schema = input_type_adapter.json_schema()
+        input_schema = compress_schema(input_schema, prune_params=prune_params)
+
+        output_schema = None
+        output_type = inspect.signature(fn).return_annotation
+        if output_type is not inspect._empty:
+            try:
+                replaced_output_type = replace_type(
+                    output_type,
+                    {
+                        Image: mcp.types.ImageContent,
+                        Audio: mcp.types.AudioContent,
+                        File: mcp.types.EmbeddedResource,
+                    },
+                )
+                output_type_adapter = get_cached_typeadapter(replaced_output_type)
+                output_schema = output_type_adapter.json_schema()
+            except PydanticSchemaGenerationError:
+                logger.debug(f"Unable to generate schema for type {output_type!r}")
+
         return cls(
             fn=fn,
             name=fn_name,
             description=fn_doc,
-            parameters=schema,
+            input_schema=input_schema,
+            output_schema=output_schema,
         )
 
 

src/fastmcp/tools/tool_transform.py

@@ -4,7 +4,6 @@ import inspect
 from collections.abc import Callable
 from contextvars import ContextVar
 from dataclasses import dataclass
-from types import EllipsisType
 from typing import Any, Literal
 
 from mcp.types import ContentBlock, ToolAnnotations
@@ -12,12 +11,10 @@ from pydantic import ConfigDict
 
 from fastmcp.tools.tool import ParsedFunction, Tool
 from fastmcp.utilities.logging import get_logger
-from fastmcp.utilities.types import get_cached_typeadapter
+from fastmcp.utilities.types import NotSet, NotSetT, get_cached_typeadapter
 
 logger = get_logger(__name__)
 
-NotSet = ...
-
 
 # Context variable to store current transformed tool
 _current_tool: ContextVar[TransformedTool | None] = ContextVar(
@@ -131,14 +128,14 @@ class ArgTransform:
         ArgTransform(name="new_name", description="New desc", default=None, type=int)
     """
 
-    name: str | EllipsisType = NotSet
-    description: str | EllipsisType = NotSet
-    default: Any | EllipsisType = NotSet
-    default_factory: Callable[[], Any] | EllipsisType = NotSet
-    type: Any | EllipsisType = NotSet
+    name: str | NotSetT = NotSet
+    description: str | NotSetT = NotSet
+    default: Any | NotSetT = NotSet
+    default_factory: Callable[[], Any] | NotSetT = NotSet
+    type: Any | NotSetT = NotSet
     hide: bool = False
-    required: Literal[True] | EllipsisType = NotSet
-    examples: Any | EllipsisType = NotSet
+    required: Literal[True] | NotSetT = NotSet
+    examples: Any | NotSetT = NotSet
 
     def __post_init__(self):
         """Validate that only one of default or default_factory is provided."""
@@ -334,7 +331,7 @@ class TransformedTool(Tool):
             has_kwargs = cls._function_has_kwargs(transform_fn)
 
             # Validate function parameters against transformed schema
-            fn_params = set(parsed_fn.parameters.get("properties", {}).keys())
+            fn_params = set(parsed_fn.input_schema.get("properties", {}).keys())
             transformed_params = set(schema.get("properties", {}).keys())
 
             if not has_kwargs:
@@ -351,7 +348,7 @@ class TransformedTool(Tool):
                 # ArgTransform takes precedence over function signature
                 # Start with function schema as base, then override with transformed schema
                 final_schema = cls._merge_schema_with_precedence(
-                    parsed_fn.parameters, schema
+                    parsed_fn.input_schema, schema
                 )
             else:
                 # With **kwargs, function can access all transformed params
@@ -360,7 +357,7 @@ class TransformedTool(Tool):
 
                 # Start with function schema as base, then override with transformed schema
                 final_schema = cls._merge_schema_with_precedence(
-                    parsed_fn.parameters, schema
+                    parsed_fn.input_schema, schema
                 )
 
         # Additional validation: check for naming conflicts after transformation

src/fastmcp/utilities/types.py

@@ -6,21 +6,19 @@ import mimetypes
 from collections.abc import Callable
 from functools import lru_cache
 from pathlib import Path
-from types import UnionType
-from typing import Annotated, TypeVar, Union, get_args, get_origin
-
-from mcp.types import (
-    Annotations,
-    AudioContent,
-    BlobResourceContents,
-    EmbeddedResource,
-    ImageContent,
-    TextResourceContents,
-)
+from types import EllipsisType, UnionType
+from typing import Annotated, TypeAlias, TypeVar, Union, get_args, get_origin
+
+import mcp.types
+from mcp.types import Annotations
 from pydantic import AnyUrl, BaseModel, ConfigDict, TypeAdapter, UrlConstraints
 
 T = TypeVar("T")
 
+# sentinel values for optional arguments
+NotSet = ...
+NotSetT: TypeAlias = EllipsisType
+
 
 class FastMCPBaseModel(BaseModel):
     """Base model for FastMCP models."""
@@ -129,7 +127,7 @@ class Image:
         self,
         mime_type: str | None = None,
         annotations: Annotations | None = None,
-    ) -> ImageContent:
+    ) -> mcp.types.ImageContent:
         """Convert to MCP ImageContent."""
         if self.path:
             with open(self.path, "rb") as f:
@@ -139,7 +137,7 @@ class Image:
         else:
             raise ValueError("No image data available")
 
-        return ImageContent(
+        return mcp.types.ImageContent(
             type="image",
             data=data,
             mimeType=mime_type or self._mime_type,
@@ -188,7 +186,7 @@ class Audio:
         self,
         mime_type: str | None = None,
         annotations: Annotations | None = None,
-    ) -> AudioContent:
+    ) -> mcp.types.AudioContent:
         if self.path:
             with open(self.path, "rb") as f:
                 data = base64.b64encode(f.read()).decode()
@@ -197,7 +195,7 @@ class Audio:
         else:
             raise ValueError("No audio data available")
 
-        return AudioContent(
+        return mcp.types.AudioContent(
             type="audio",
             data=data,
             mimeType=mime_type or self._mime_type,
@@ -248,7 +246,7 @@ class File:
         self,
         mime_type: str | None = None,
         annotations: Annotations | None = None,
-    ) -> EmbeddedResource:
+    ) -> mcp.types.EmbeddedResource:
         if self.path:
             with open(self.path, "rb") as f:
                 raw_data = f.read()
@@ -271,21 +269,57 @@ class File:
                 text = raw_data.decode("utf-8")
             except UnicodeDecodeError:
                 text = raw_data.decode("latin-1")
-            resource = TextResourceContents(
+            resource = mcp.types.TextResourceContents(
                 text=text,
                 mimeType=mime,
                 uri=uri,
             )
         else:
             data = base64.b64encode(raw_data).decode()
-            resource = BlobResourceContents(
+            resource = mcp.types.BlobResourceContents(
                 blob=data,
                 mimeType=mime,
                 uri=uri,
             )
 
-        return EmbeddedResource(
+        return mcp.types.EmbeddedResource(
             type="resource",
             resource=resource,
             annotations=annotations or self.annotations,
         )
+
+
+def replace_type(type_, type_map: dict[type, type]):
+    """
+    Given a (possibly generic, nested, or otherwise complex) type, replaces all
+    instances of old_type with new_type.
+
+    This is useful for transforming types when creating tools.
+
+    Args:
+        type_: The type to replace instances of old_type with new_type.
+        old_type: The type to replace.
+        new_type: The type to replace old_type with.
+
+    Examples:
+        >>> replace_type(list[int | bool], {int: str})
+        list[str | bool]
+
+        >>> replace_type(list[list[int]], {int: str})
+        list[list[str]]
+
+    """
+    if type_ in type_map:
+        return type_map[type_]
+
+    origin = get_origin(type_)
+    if not origin:
+        return type_
+
+    args = get_args(type_)
+    new_args = tuple(replace_type(arg, type_map) for arg in args)
+
+    if origin is UnionType:
+        return Union[new_args]  # type: ignore # noqa: UP007
+    else:
+        return origin[new_args]

tests/server/test_server_interactions.py

@@ -17,7 +17,7 @@ from mcp.types import (
     TextContent,
     TextResourceContents,
 )
-from pydantic import AnyUrl, Field
+from pydantic import AnyUrl, Field, TypeAdapter
 
 from fastmcp import Client, Context, FastMCP
 from fastmcp.client.transports import FastMCPTransport
@@ -826,6 +826,22 @@ class TestToolParameters:
             assert result[0].text == "0:16:40"  # type: ignore[attr-defined]
 
 
+class TestToolOutputSchema:
+    @pytest.mark.parametrize("annotation", [str, int, float, bool, list, dict, AnyUrl])
+    async def test_output_schema(self, annotation):
+        mcp = FastMCP()
+
+        @mcp.tool
+        def f() -> annotation:  # type: ignore
+            return "hello"
+
+        async with Client(mcp) as client:
+            tools = await client.list_tools()
+            assert len(tools) == 1
+            # this line will fail until MCP adds output schemas!!
+            assert tools[0].outputSchema == TypeAdapter(annotation).json_schema()  # type: ignore
+
+
 class TestToolContextInjection:
     """Test context injection in tools."""
 

tests/tools/test_tool.py

@@ -1,4 +1,6 @@
 import json
+from dataclasses import dataclass
+from typing import Annotated, Any, TypedDict
 
 import pytest
 from mcp.types import (
@@ -8,7 +10,7 @@ from mcp.types import (
     TextContent,
     TextResourceContents,
 )
-from pydantic import AnyUrl, BaseModel
+from pydantic import AnyUrl, BaseModel, Field, TypeAdapter
 
 from fastmcp import FastMCP
 from fastmcp.client import Client
@@ -33,6 +35,7 @@ class TestToolFromFunction:
         assert len(tool.parameters["properties"]) == 2
         assert tool.parameters["properties"]["a"]["type"] == "integer"
         assert tool.parameters["properties"]["b"]["type"] == "integer"
+        assert tool.output_schema == {"type": "integer"}
 
     async def test_async_function(self):
         """Test registering and running an async function."""
@@ -244,6 +247,129 @@ class TestToolFromFunction:
         assert result[0].text == "Custom serializer: 15"
 
 
+class TestToolFromFunctionOutputSchema:
+    async def test_no_return_annotation(self):
+        def func():
+            pass
+
+        tool = Tool.from_function(func)
+        assert tool.output_schema is None
+
+    @pytest.mark.parametrize(
+        "annotation",
+        [
+            None,
+            int,
+            float,
+            bool,
+            str,
+            int | float,
+            list[int],
+            list[int | float],
+            dict[str, int | None],
+            tuple[int, str],
+            set[int],
+            list[tuple[int, str]],
+        ],
+    )
+    async def test_simple_return_annotation(self, annotation):
+        def func() -> annotation:  # type: ignore
+            return 1
+
+        tool = Tool.from_function(func)
+        assert tool.output_schema == TypeAdapter(annotation).json_schema()
+
+    @pytest.mark.parametrize(
+        "annotation",
+        [
+            Any,
+            AnyUrl,
+            Annotated[int, Field(ge=1)],
+            Annotated[int, Field(ge=1)],
+        ],
+    )
+    async def test_complex_return_annotation(self, annotation):
+        def func() -> annotation:  # type: ignore
+            return 1
+
+        tool = Tool.from_function(func)
+        assert tool.output_schema == TypeAdapter(annotation).json_schema()
+
+    @pytest.mark.parametrize(
+        "annotation, expected",
+        [
+            (Image, ImageContent),
+            (Audio, AudioContent),
+            (File, EmbeddedResource),
+            (Image | int, ImageContent | int),
+            (Image | Audio, ImageContent | AudioContent),
+            (list[Image | Audio], list[ImageContent | AudioContent]),
+        ],
+    )
+    async def test_converted_return_annotation(self, annotation, expected):
+        def func() -> annotation:  # type: ignore
+            return 1
+
+        tool = Tool.from_function(func)
+        assert tool.output_schema == TypeAdapter(expected).json_schema()
+
+    async def test_dataclass_return_annotation(self):
+        @dataclass
+        class Person:
+            name: str
+            age: int
+
+        def func() -> Person:
+            return Person(name="John", age=30)
+
+        tool = Tool.from_function(func)
+        assert tool.output_schema == TypeAdapter(Person).json_schema()
+
+    async def test_base_model_return_annotation(self):
+        class Person(BaseModel):
+            name: str
+            age: int
+
+        def func() -> Person:
+            return Person(name="John", age=30)
+
+        tool = Tool.from_function(func)
+        assert tool.output_schema == TypeAdapter(Person).json_schema()
+
+    async def test_typeddict_return_annotation(self):
+        class Person(TypedDict):
+            name: str
+            age: int
+
+        def func() -> Person:
+            return Person(name="John", age=30)
+
+        tool = Tool.from_function(func)
+        assert tool.output_schema == TypeAdapter(Person).json_schema()
+
+    async def test_unserializable_return_annotation(self):
+        class Unserializable:
+            def __init__(self, data: Any):
+                self.data = data
+
+        def func() -> Unserializable:
+            return Unserializable(data="test")
+
+        tool = Tool.from_function(func)
+        assert tool.output_schema is None
+
+    async def test_mixed_unserializable_return_annotation(self):
+        class Unserializable:
+            def __init__(self, data: Any):
+                self.data = data
+
+        def func() -> Unserializable | int:
+            return Unserializable(data="test")
+
+        tool = Tool.from_function(func)
+        assert tool.output_schema is None
+
+
 class TestLegacyToolJsonParsing:
     """Tests for Tool's JSON pre-parsing functionality."""
 

tests/utilities/test_types.py

@@ -12,6 +12,7 @@ from fastmcp.utilities.types import (
     find_kwarg_by_type,
     is_class_member_of_type,
     issubclass_safe,
+    replace_type,
 )
 
 
@@ -536,3 +537,29 @@ class TestFindKwargByType:
             pass
 
         assert find_kwarg_by_type(func, str) == "c"
+
+
+class TestReplaceType:
+    @pytest.mark.parametrize(
+        "input,type_map,expected",
+        [
+            (int, {}, int),
+            (int, {int: str}, str),
+            (int, {int: int}, int),
+            (int, {int: float, bool: str}, float),
+            (bool, {int: float, bool: str}, str),
+            (int, {int: list[int]}, list[int]),
+            (list[int], {int: str}, list[str]),
+            (list[int], {int: list[str]}, list[list[str]]),
+            (
+                list[int],
+                {int: float, list[int]: bool},
+                bool,
+            ),  # list[int] will match before int
+            (list[int | bool], {int: str}, list[str | bool]),
+            (list[list[int]], {int: str}, list[list[str]]),
+        ],
+    )
+    def test_replace_type(self, input, type_map, expected):
+        """Test replacing a type with another type."""
+        assert replace_type(input, type_map) == expected