Add output schema generation and related tests

Author: bzsurbhi

src/mcp/server/fastmcp/server.py

@@ -41,6 +41,7 @@
     GetPromptResult,
     ImageContent,
     TextContent,
+    DataContent,
     ToolAnnotations,
 )
 from mcp.types import Prompt as MCPPrompt
@@ -169,14 +170,17 @@ def _setup_handlers(self) -> None:
         self._mcp_server.get_prompt()(self.get_prompt)
         self._mcp_server.list_resource_templates()(self.list_resource_templates)
 
+
     async def list_tools(self) -> list[MCPTool]:
         """List all available tools."""
         tools = self._tool_manager.list_tools()
+        logger.info(f"list tools: {tools}")
         return [
             MCPTool(
                 name=info.name,
                 description=info.description,
                 inputSchema=info.parameters,
+                outputSchema=info.outputSchema,
                 annotations=info.annotations,
             )
             for info in tools
@@ -195,10 +199,11 @@ def get_context(self) -> Context[ServerSession, object]:
 
     async def call_tool(
         self, name: str, arguments: dict[str, Any]
-    ) -> Sequence[TextContent | ImageContent | EmbeddedResource]:
+    ) -> Sequence[TextContent | ImageContent | DataContent | EmbeddedResource]:
         """Call a tool by name with arguments."""
         context = self.get_context()
         result = await self._tool_manager.call_tool(name, arguments, context=context)
+        logger.info(f"call tool: {name} with args: {arguments} -> {result}")
         converted_result = _convert_to_content(result)
         return converted_result
 
@@ -260,7 +265,10 @@ def add_tool(
             annotations: Optional ToolAnnotations providing additional tool information
         """
         self._tool_manager.add_tool(
-            fn, name=name, description=description, annotations=annotations
+            fn,
+            name=name,
+            description=description,
+            annotations=annotations,
         )
 
     def tool(
@@ -304,7 +312,10 @@ async def async_tool(x: int, context: Context) -> str:
 
         def decorator(fn: AnyFunction) -> AnyFunction:
             self.add_tool(
-                fn, name=name, description=description, annotations=annotations
+                fn,
+                name=name,
+                description=description,
+                annotations=annotations,
             )
             return fn
 
@@ -547,12 +558,13 @@ async def get_prompt(
 
 def _convert_to_content(
     result: Any,
-) -> Sequence[TextContent | ImageContent | EmbeddedResource]:
+) -> Sequence[TextContent | ImageContent | EmbeddedResource | DataContent]:
     """Convert a result to a sequence of content objects."""
     if result is None:
         return []
 
-    if isinstance(result, TextContent | ImageContent | EmbeddedResource):
+    # Handle existing content types
+    if isinstance(result, TextContent | ImageContent | EmbeddedResource | DataContent):
         return [result]
 
     if isinstance(result, Image):
@@ -561,9 +573,21 @@ def _convert_to_content(
     if isinstance(result, list | tuple):
         return list(chain.from_iterable(_convert_to_content(item) for item in result))  # type: ignore[reportUnknownVariableType]
 
+    # For non-string objects, convert to DataContent
     if not isinstance(result, str):
-        result = pydantic_core.to_json(result, fallback=str, indent=2).decode()
+        # Try to convert to a JSON-serializable structure
+        try:
+            # Get the data as a dict/list structure
+            data = pydantic_core.to_jsonable_python(result)
+            # Create DataContent with the data
+            return [DataContent(type="data", data=data)]
+        except Exception as e:
+            logger.warning(f"Failed to convert result to DataContent: {e}")
+            # Fall back to string representation
+            result_str = pydantic_core.to_json(result, fallback=str, indent=2).decode()
+            return [TextContent(type="text", text=result_str)]
 
+    # For strings, use TextContent
     return [TextContent(type="text", text=result)]
 
 

src/mcp/server/fastmcp/tools/base.py

@@ -3,6 +3,7 @@
 import inspect
 from collections.abc import Callable
 from typing import TYPE_CHECKING, Any, get_origin
+from mcp.server.fastmcp.utilities.logging import get_logger
 
 from pydantic import BaseModel, Field
 
@@ -15,6 +16,7 @@
     from mcp.server.session import ServerSessionT
     from mcp.shared.context import LifespanContextT
 
+logger = get_logger(__name__)
 
 class Tool(BaseModel):
     """Internal tool registration info."""
@@ -23,6 +25,9 @@ class Tool(BaseModel):
     name: str = Field(description="Name of the tool")
     description: str = Field(description="Description of what the tool does")
     parameters: dict[str, Any] = Field(description="JSON schema for tool parameters")
+    outputSchema: dict[str, Any] | None = Field(
+        None, description="Optional JSON schema for tool output"
+    )
     fn_metadata: FuncMetadata = Field(
         description="Metadata about the function including a pydantic model for tool"
         " arguments"
@@ -70,6 +75,10 @@ def from_function(
         )
         parameters = func_arg_metadata.arg_model.model_json_schema()
 
+        output_schema = getattr(func_arg_metadata, "outputSchema", None)
+
+        logger.info(f"output schema: {output_schema}")
+
         return cls(
             fn=fn,
             name=func_name,
@@ -78,6 +87,7 @@ def from_function(
             fn_metadata=func_arg_metadata,
             is_async=is_async,
             context_kwarg=context_kwarg,
+            outputSchema=output_schema,
             annotations=annotations,
         )
 
@@ -87,6 +97,7 @@ async def run(
         context: Context[ServerSessionT, LifespanContextT] | None = None,
     ) -> Any:
         """Run the tool with arguments."""
+
         try:
             return await self.fn_metadata.call_fn_with_arg_validation(
                 self.fn,

src/mcp/server/fastmcp/tools/tool_manager.py

@@ -40,7 +40,10 @@ def add_tool(
     ) -> Tool:
         """Add a tool to the server."""
         tool = Tool.from_function(
-            fn, name=name, description=description, annotations=annotations
+            fn,
+            name=name,
+            description=description,
+            annotations=annotations,
         )
         existing = self._tools.get(tool.name)
         if existing:

src/mcp/server/fastmcp/utilities/func_metadata.py

@@ -38,6 +38,7 @@ def model_dump_one_level(self) -> dict[str, Any]:
 
 class FuncMetadata(BaseModel):
     arg_model: Annotated[type[ArgModelBase], WithJsonSchema(None)]
+    outputSchema: dict[str, Any] | None = None
     # We can add things in the future like
     #  - Maybe some args are excluded from attempting to parse from JSON
     #  - Maybe some args are special (like context) for dependency injection
@@ -60,6 +61,11 @@ async def call_fn_with_arg_validation(
 
         arguments_parsed_dict |= arguments_to_pass_directly or {}
 
+        logger.info(
+            "Calling function with arguments: %s",
+            arguments_parsed_dict,
+        )
+        logger.info(f"Function is async: ${fn}")
         if fn_is_async:
             if isinstance(fn, Awaitable):
                 return await fn
@@ -172,7 +178,62 @@ def func_metadata(
         **dynamic_pydantic_model_params,
         __base__=ArgModelBase,
     )
-    resp = FuncMetadata(arg_model=arguments_model)
+
+    # Generate output schema from return type annotation
+    output_schema: dict[str, Any] | None = None
+    return_annotation = sig.return_annotation
+
+    if return_annotation is not inspect.Signature.empty:
+        try:
+            # Handle forward references
+            return_type = _get_typed_annotation(return_annotation, globalns)
+            logger.info(f"return_type: {return_type}")
+            # Special case for None
+            if return_type is type(None):  # noqa: E721
+                output_schema = {"type": "null"}
+            else:
+                # Create a temporary model to get the schema
+                class OutputModel(BaseModel):
+                    result: return_type  # type: ignore
+
+                    model_config = ConfigDict(
+                        arbitrary_types_allowed=True,
+                    )
+
+                # Extract the schema for the return type
+                full_schema = OutputModel.model_json_schema()
+
+                # If the return type is a complex type, use its schema definition
+                if "$defs" in full_schema and "result" in full_schema.get(
+                    "properties", {}
+                ):
+                    prop = full_schema["properties"]["result"]
+                    if isinstance(prop, dict) and "$ref" in prop:
+                        if isinstance(prop["$ref"], str):
+                            ref_name = prop["$ref"].split("/")[-1]
+                        else:
+                            raise TypeError("Expected to be a string")
+                        if ref_name in full_schema.get("$defs", {}):
+                            ref_schema = full_schema["$defs"][ref_name]
+                            output_schema = {
+                                "type": "object",
+                                "properties": ref_schema.get("properties", {}),
+                                "required": ref_schema.get("required", []),
+                            }
+                            # Optionally include title if present
+                            if "title" in ref_schema:
+                                output_schema["title"] = ref_schema["title"]
+                    else:
+                        output_schema = prop
+                else:
+                    # For simple types
+                    output_schema = full_schema["properties"]["result"]
+
+        except Exception as e:
+            # If we can't generate a schema, log the error but continue
+            logger.warning(f"Failed to generate output schema for {func.__name__}: {e}")
+
+    resp = FuncMetadata(arg_model=arguments_model, outputSchema=output_schema)
     return resp
 
 
@@ -194,6 +255,10 @@ def try_eval_type(
         if status is False:
             raise InvalidSignature(f"Unable to evaluate type annotation {annotation}")
 
+    # If the annotation is already a valid type, return it directly
+    if isinstance(annotation, type):
+        return annotation
+
     return annotation
 
 
@@ -210,5 +275,8 @@ def _get_typed_signature(call: Callable[..., Any]) -> inspect.Signature:
         )
         for param in signature.parameters.values()
     ]
-    typed_signature = inspect.Signature(typed_params)
+    typed_signature = inspect.Signature(
+        typed_params,
+        return_annotation=_get_typed_annotation(signature.return_annotation, globalns),
+    )
     return typed_signature

src/mcp/server/lowlevel/server.py

@@ -398,7 +398,7 @@ def decorator(
                 ...,
                 Awaitable[
                     Iterable[
-                        types.TextContent | types.ImageContent | types.EmbeddedResource
+                        types.TextContent | types.ImageContent | types.DataContent | types.EmbeddedResource
                     ]
                 ],
             ],

src/mcp/types.py

@@ -646,6 +646,27 @@ class ImageContent(BaseModel):
     model_config = ConfigDict(extra="allow")
 
 
+class DataContent(BaseModel):
+    """Structured JSON content for a message or tool result."""
+
+    type: Literal["data"]
+    data: dict[str, Any]
+    """
+    The structured JSON data. This is a JSON serializable object.
+    """
+
+    schema_definition: dict[str, Any] | str | None = None
+    """
+    An optional schema describing the structure of the data.
+    - Can be a string (schema reference URI),
+    - A dictionary (full schema definition),
+    - Or omitted if no schema is provided.
+    """
+
+    annotations: Annotations | None = None
+    model_config = ConfigDict(extra="allow")
+
+
 class SamplingMessage(BaseModel):
     """Describes a message issued to or received from an LLM API."""
 
@@ -762,6 +783,8 @@ class Tool(BaseModel):
     """A human-readable description of the tool."""
     inputSchema: dict[str, Any]
     """A JSON Schema object defining the expected parameters for the tool."""
+    outputSchema: dict[str, Any] | None = None
+    """A JSON Schema object defining the expected structure of the tool's output."""
     annotations: ToolAnnotations | None = None
     """Optional additional tool information."""
     model_config = ConfigDict(extra="allow")
@@ -791,7 +814,7 @@ class CallToolRequest(Request[CallToolRequestParams, Literal["tools/call"]]):
 class CallToolResult(Result):
     """The server's response to a tool call."""
 
-    content: list[TextContent | ImageContent | EmbeddedResource]
+    content: list[TextContent | ImageContent | DataContent | EmbeddedResource]
     isError: bool = False