diff --git a/src/mcp/server/fastmcp/server.py b/src/mcp/server/fastmcp/server.py index bc2b105e..1b2561c6 100644 --- a/src/mcp/server/fastmcp/server.py +++ b/src/mcp/server/fastmcp/server.py @@ -177,6 +177,7 @@ async def list_tools(self) -> list[MCPTool]: name=info.name, description=info.description, inputSchema=info.parameters, + outputSchema=info.outputSchema, annotations=info.annotations, ) for info in tools @@ -260,7 +261,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 +308,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 @@ -552,6 +559,7 @@ def _convert_to_content( if result is None: return [] + # Handle existing content types if isinstance(result, TextContent | ImageContent | EmbeddedResource): return [result] diff --git a/src/mcp/server/fastmcp/tools/base.py b/src/mcp/server/fastmcp/tools/base.py index 21eb1841..4401ccd1 100644 --- a/src/mcp/server/fastmcp/tools/base.py +++ b/src/mcp/server/fastmcp/tools/base.py @@ -15,7 +15,6 @@ from mcp.server.session import ServerSessionT from mcp.shared.context import LifespanContextT - class Tool(BaseModel): """Internal tool registration info.""" @@ -23,6 +22,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 +72,8 @@ def from_function( ) parameters = func_arg_metadata.arg_model.model_json_schema() + output_schema = getattr(func_arg_metadata, "outputSchema", None) + return cls( fn=fn, name=func_name, @@ -78,6 +82,7 @@ def from_function( fn_metadata=func_arg_metadata, is_async=is_async, context_kwarg=context_kwarg, + outputSchema=output_schema, annotations=annotations, ) @@ -87,6 +92,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, diff --git a/src/mcp/server/fastmcp/tools/tool_manager.py b/src/mcp/server/fastmcp/tools/tool_manager.py index cfdaeb35..13825ab9 100644 --- a/src/mcp/server/fastmcp/tools/tool_manager.py +++ b/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: diff --git a/src/mcp/server/fastmcp/utilities/func_metadata.py b/src/mcp/server/fastmcp/utilities/func_metadata.py index 37439132..78a47afe 100644 --- a/src/mcp/server/fastmcp/utilities/func_metadata.py +++ b/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 @@ -172,7 +173,61 @@ 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) + # 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 +249,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 +269,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 diff --git a/src/mcp/server/lowlevel/server.py b/src/mcp/server/lowlevel/server.py index dbaff305..45a639f3 100644 --- a/src/mcp/server/lowlevel/server.py +++ b/src/mcp/server/lowlevel/server.py @@ -398,7 +398,9 @@ def decorator( ..., Awaitable[ Iterable[ - types.TextContent | types.ImageContent | types.EmbeddedResource + types.TextContent + | types.ImageContent + | types.EmbeddedResource ] ], ], diff --git a/src/mcp/types.py b/src/mcp/types.py index 6ab7fba5..dfe41179 100644 --- a/src/mcp/types.py +++ b/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 diff --git a/tests/server/fastmcp/test_func_metadata.py b/tests/server/fastmcp/test_func_metadata.py index b1828ffe..0e9c8825 100644 --- a/tests/server/fastmcp/test_func_metadata.py +++ b/tests/server/fastmcp/test_func_metadata.py @@ -414,3 +414,144 @@ def func_with_str_and_int(a: str, b: int): result = meta.pre_parse_json({"a": "123", "b": 123}) assert result["a"] == "123" assert result["b"] == 123 + + +def test_output_schema_generation(): + """Test automatic generation of output schemas from return type annotations.""" + + # Test with simple return types + def fn_returns_str() -> str: + """Function that returns a string.""" + return "hello" + + meta = func_metadata(fn_returns_str) + assert meta.outputSchema is not None + assert meta.outputSchema["type"] == "string" + + def fn_returns_int() -> int: + """Function that returns an integer.""" + return 42 + + meta = func_metadata(fn_returns_int) + assert meta.outputSchema is not None + assert meta.outputSchema["type"] == "integer" + + def fn_returns_bool() -> bool: + """Function that returns a boolean.""" + return True + + meta = func_metadata(fn_returns_bool) + assert meta.outputSchema is not None + assert meta.outputSchema["type"] == "boolean" + + # Test with container types + def fn_returns_list_of_str() -> list[str]: + """Function that returns a list of strings.""" + return ["hello", "world"] + + meta = func_metadata(fn_returns_list_of_str) + assert meta.outputSchema is not None + assert meta.outputSchema["type"] == "array" + assert meta.outputSchema["items"]["type"] == "string" + + def fn_returns_dict_str_int() -> dict[str, int]: + """Function that returns a dictionary mapping strings to integers.""" + return {"a": 1, "b": 2} + + meta = func_metadata(fn_returns_dict_str_int) + assert meta.outputSchema is not None + assert meta.outputSchema["type"] == "object" + # Check that it's a dict with string keys and integer values + assert "additionalProperties" in meta.outputSchema + assert meta.outputSchema["additionalProperties"]["type"] == "integer" + + # Test with optional types + def fn_returns_optional_str() -> str | None: + """Function that returns an optional string.""" + return "hello" + + meta = func_metadata(fn_returns_optional_str) + assert meta.outputSchema is not None + assert "anyOf" in meta.outputSchema or "oneOf" in meta.outputSchema + # The schema should allow either string or null + + # Test with union types + def fn_returns_union() -> str | int: + """Function that returns either a string or an integer.""" + return "hello" + + meta = func_metadata(fn_returns_union) + assert meta.outputSchema is not None + assert "anyOf" in meta.outputSchema or "oneOf" in meta.outputSchema + # The schema should allow either string or integer + + # Test with Pydantic models + class UserProfile(BaseModel): + name: str + email: str + age: int + is_active: bool = True + + def fn_returns_pydantic_model() -> UserProfile: + """Function that returns a Pydantic model.""" + return UserProfile(name="John", email="john@example.com", age=30) + + meta = func_metadata(fn_returns_pydantic_model) + assert meta.outputSchema is not None + assert meta.outputSchema["type"] == "object" + assert "properties" in meta.outputSchema + assert "name" in meta.outputSchema["properties"] + assert meta.outputSchema["properties"]["name"]["type"] == "string" + assert "email" in meta.outputSchema["properties"] + assert meta.outputSchema["properties"]["email"]["type"] == "string" + assert "age" in meta.outputSchema["properties"] + assert meta.outputSchema["properties"]["age"]["type"] == "integer" + assert "is_active" in meta.outputSchema["properties"] + assert meta.outputSchema["properties"]["is_active"]["type"] == "boolean" + assert "required" in meta.outputSchema + assert "name" in meta.outputSchema["required"] + assert "email" in meta.outputSchema["required"] + assert "age" in meta.outputSchema["required"] + assert "is_active" not in meta.outputSchema["required"] # It has a default value + + # Test with nested Pydantic models + class Address(BaseModel): + street: str + city: str + zip_code: str + + class Person(BaseModel): + name: str + age: int + address: Address + + def fn_returns_nested_model() -> Person: + """Function that returns a nested Pydantic model.""" + return Person( + name="John", + age=30, + address=Address(street="123 Main St", city="Anytown", zip_code="12345"), + ) + + meta = func_metadata(fn_returns_nested_model) + assert meta.outputSchema is not None + assert meta.outputSchema["type"] == "object" + assert "properties" in meta.outputSchema + assert "address" in meta.outputSchema["properties"] + + # Test with a function that has no return type annotation + def fn_no_return_type(): + """Function with no return type annotation.""" + return "hello" + + meta = func_metadata(fn_no_return_type) + assert meta.outputSchema is None + + # Test with a function that returns None + def fn_returns_none() -> None: + """Function that returns None.""" + return None + + meta = func_metadata(fn_returns_none) + assert meta.outputSchema is not None + assert meta.outputSchema["type"] == "null"