diff --git a/src/mcp/server/mcpserver/utilities/func_metadata.py b/src/mcp/server/mcpserver/utilities/func_metadata.py index 0ffac07c4e..4b6ddbb80d 100644 --- a/src/mcp/server/mcpserver/utilities/func_metadata.py +++ b/src/mcp/server/mcpserver/utilities/func_metadata.py @@ -1,3 +1,4 @@ +import base64 import functools import inspect import json @@ -22,7 +23,8 @@ create_model, ) from pydantic.fields import FieldInfo -from pydantic.json_schema import GenerateJsonSchema, JsonSchemaWarningKind +from pydantic.json_schema import GenerateJsonSchema, JsonSchemaValue, JsonSchemaWarningKind +from pydantic_core import core_schema from typing_extensions import NotRequired, ReadOnly, TypedDict, deprecated, get_type_hints, is_typeddict from typing_inspection.introspection import ( UNKNOWN, @@ -73,6 +75,11 @@ def emit_warning(self, kind: JsonSchemaWarningKind, detail: str) -> None: # Raise an exception instead of emitting a warning raise ValueError(f"JSON schema warning: {kind} - {detail}") + def bytes_schema(self, schema: core_schema.BytesSchema) -> JsonSchemaValue: + # `bytes` output is serialized to base64 (see `_create_wrapped_model`), but the advertised + # schema keeps the plain `format: binary` regardless of the serializer's `ser_json_bytes`. + return {"type": "string", "format": "binary"} + _LOCAL_DEFS_PREFIX = "#/$defs/" @@ -577,7 +584,11 @@ def _create_model_from_class(cls: type[Any], type_hints: dict[str, Any]) -> type else: model_fields[field_name] = (field_type, default) - return create_model(cls.__name__, __config__=ConfigDict(from_attributes=True), **model_fields) + return create_model( + cls.__name__, + __config__=ConfigDict(from_attributes=True, ser_json_bytes="base64"), + **model_fields, + ) def _pydantic_readable_typeddict(output_model: type[Any]) -> type[Any]: @@ -618,7 +629,14 @@ def _create_wrapped_model(func_name: str, annotation: Any) -> type[BaseModel]: """ model_name = f"{func_name}Output" - return create_model(model_name, result=annotation) + # Serialize any `bytes` field (e.g. a `-> bytes` return) as base64 for JSON so structured output + # does not crash on non-UTF-8 data. Only the serializer is changed, so the advertised schema keeps + # its `format: binary`. + return create_model( + model_name, + __config__=ConfigDict(ser_json_bytes="base64"), + result=annotation, + ) def _convert_to_content(result: Any) -> list[ContentBlock]: @@ -649,6 +667,11 @@ def _convert_to_content(result: Any) -> list[ContentBlock]: ) ) + # `bytes` advertise `format: binary`; base64-encode them (matching `Image`/`Audio` and the lowlevel + # server) rather than letting `to_json` UTF-8-decode them and raise on non-UTF-8 data. + if isinstance(result, bytes): + result = base64.b64encode(result).decode() + if not isinstance(result, str): result = pydantic_core.to_json(result, fallback=str, indent=2).decode() diff --git a/tests/server/mcpserver/test_func_metadata.py b/tests/server/mcpserver/test_func_metadata.py index dba0637ded..cf9a1562d5 100644 --- a/tests/server/mcpserver/test_func_metadata.py +++ b/tests/server/mcpserver/test_func_metadata.py @@ -3,6 +3,7 @@ # pyright: reportMissingParameterType=false # pyright: reportUnknownArgumentType=false # pyright: reportUnknownLambdaType=false +import base64 from collections.abc import Callable from dataclasses import dataclass from typing import TYPE_CHECKING, Annotated, Any, Final, NamedTuple, TypedDict @@ -17,7 +18,12 @@ from mcp import MCPDeprecationWarning from mcp.server.mcpserver import Audio, Image from mcp.server.mcpserver.exceptions import InvalidSignature -from mcp.server.mcpserver.utilities.func_metadata import ArgModelBase, FuncMetadata, func_metadata +from mcp.server.mcpserver.utilities.func_metadata import ( + ArgModelBase, + FuncMetadata, + _convert_to_content, + func_metadata, +) if TYPE_CHECKING: from decimal import Decimal @@ -724,6 +730,72 @@ def func_bytes() -> bytes: # pragma: no cover } +def test_convert_to_content_base64_encodes_non_utf8_bytes(): + """Non-UTF-8 `bytes` (e.g. PNG magic bytes) base64-encode instead of crashing in `to_json`.""" + png = b"\x89PNG\r\n\x1a\n" + content = _convert_to_content(png) + assert len(content) == 1 + assert isinstance(content[0], TextContent) + assert content[0].text == base64.b64encode(png).decode() + + +def test_convert_to_content_base64_encodes_utf8_bytes(): + """UTF-8-decodable `bytes` are base64-encoded too, consistent with the advertised `format: binary`.""" + content = _convert_to_content(b"hello") + assert len(content) == 1 + assert isinstance(content[0], TextContent) + assert content[0].text == base64.b64encode(b"hello").decode() + + +def test_structured_output_bytes_return_is_base64(): + """A `-> bytes` tool returning non-UTF-8 bytes yields base64 in both unstructured and structured output.""" + png = b"\x89PNG\r\n\x1a\n" + expected = base64.b64encode(png).decode() + + def func_bytes() -> bytes: + return png + + meta = func_metadata(func_bytes) + # Schema stays `format: binary` (regression guard). + assert meta.output_schema == { + "type": "object", + "properties": {"result": {"title": "Result", "type": "string", "format": "binary"}}, + "required": ["result"], + "title": "func_bytesOutput", + } + + result = meta.convert_result(png) + assert isinstance(result, CallToolResult) + assert not result.is_error + assert result.structured_content == {"result": expected} + assert isinstance(result.content[0], TextContent) + assert result.content[0].text == expected + + +def test_structured_output_bytes_field_round_trips(): + """A `bytes` field in an output model (ordinary class) serializes to base64 and round-trips back.""" + png = b"\x89PNG\r\n\x1a\n" + + class Thumb: + data: bytes + + def __init__(self, data: bytes) -> None: + self.data = data + + def make_thumb() -> Thumb: + return Thumb(png) + + meta = func_metadata(make_thumb) + result = meta.convert_result(Thumb(png)) + assert isinstance(result, CallToolResult) + assert not result.is_error + assert result.structured_content == {"data": base64.b64encode(png).decode()} + + # The base64 structured content decodes back to the original bytes. + assert result.structured_content is not None + assert base64.b64decode(result.structured_content["data"]) == png + + def test_structured_output_generic_types(): """Test structured output with generic types (list, dict, Union, etc.)"""