diff --git a/src/mcp/server/mcpserver/utilities/func_metadata.py b/src/mcp/server/mcpserver/utilities/func_metadata.py index 4a76106371..3ea9bf7ce5 100644 --- a/src/mcp/server/mcpserver/utilities/func_metadata.py +++ b/src/mcp/server/mcpserver/utilities/func_metadata.py @@ -478,7 +478,7 @@ def _create_wrapped_model(func_name: str, annotation: Any) -> type[BaseModel]: """ model_name = f"{func_name}Output" - return create_model(model_name, result=annotation) + return create_model(model_name, result=(annotation, ...)) def _create_dict_model(func_name: str, dict_annotation: Any) -> type[BaseModel]: diff --git a/tests/server/mcpserver/test_func_metadata.py b/tests/server/mcpserver/test_func_metadata.py index c57d1ee9f0..ec73e6ee99 100644 --- a/tests/server/mcpserver/test_func_metadata.py +++ b/tests/server/mcpserver/test_func_metadata.py @@ -677,6 +677,9 @@ def func_dict_str_int() -> dict[str, int]: # pragma: no cover def func_union() -> str | int: # pragma: no cover return "hello" + def func_pipe_union_containers() -> dict[str, int] | list[str] | str: # pragma: no cover + return {"a": 1} + def func_optional() -> str | None: # pragma: no cover return None @@ -706,6 +709,24 @@ def func_optional() -> str | None: # pragma: no cover "title": "func_unionOutput", } + # Test PEP 604 union with containers + meta = func_metadata(func_pipe_union_containers) + assert meta.output_schema == { + "type": "object", + "properties": { + "result": { + "title": "Result", + "anyOf": [ + {"additionalProperties": {"type": "integer"}, "type": "object"}, + {"items": {"type": "string"}, "type": "array"}, + {"type": "string"}, + ], + } + }, + "required": ["result"], + "title": "func_pipe_union_containersOutput", + } + # Test Optional meta = func_metadata(func_optional) assert meta.output_schema == { diff --git a/tests/server/mcpserver/test_server.py b/tests/server/mcpserver/test_server.py index 3457ec944a..a79de81f17 100644 --- a/tests/server/mcpserver/test_server.py +++ b/tests/server/mcpserver/test_server.py @@ -1,5 +1,6 @@ import base64 from pathlib import Path +from types import TracebackType from typing import Any from unittest.mock import AsyncMock, MagicMock, patch @@ -65,6 +66,41 @@ async def test_create_server(self): assert len(mcp.icons) == 1 assert mcp.icons[0].src == "https://example.com/icon.png" + def test_run_stdio_transport(self): + mcp = MCPServer("test") + + with patch("mcp.server.mcpserver.server.anyio.run") as run: + mcp.run("stdio") + + run.assert_called_once_with(mcp.run_stdio_async) + + async def test_run_stdio_async_uses_stdio_transport(self): + class StdioServer: + async def __aenter__(self): + return "read", "write" + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + tb: TracebackType | None, + ) -> None: + return None + + mcp = MCPServer("test") + run = AsyncMock() + + with ( + patch("mcp.server.mcpserver.server.stdio_server", return_value=StdioServer()), + patch.object(mcp._lowlevel_server, "run", run), + ): + await mcp.run_stdio_async() + + run.assert_awaited_once() + await_args = run.await_args + assert await_args is not None + assert await_args.args[:2] == ("read", "write") + def test_dependencies(self): """Dependencies list is read by `mcp install` / `mcp dev` CLI commands.""" mcp = MCPServer("test", dependencies=["pandas", "numpy"])