Skip to content

Commit aa95cdf

Browse files
committed
fix(mcpserver): base64-encode bytes tool output instead of crashing on non-UTF-8
A tool annotated `-> bytes` advertises output schema {"type":"string","format":"binary"}, but returning non-UTF-8 binary data (e.g. PNG magic bytes) crashed with PydanticSerializationError. Root cause: _convert_to_content fed bytes to pydantic_core.to_json, which UTF-8-decodes them (fallback=str only applies to types pydantic cannot serialize), and the generated structured-output model serialized bytes with the default ser_json_bytes='utf8'. Both raise on non-UTF-8 bytes; UTF-8-decodable bytes also came back as a raw JSON string rather than base64. Fix: - _convert_to_content now base64-encodes bytes before the to_json branch, matching Image/Audio and the lowlevel server, so all bytes become a consistent base64 string. - Generated output models (_create_wrapped_model, _create_model_from_class) set ser_json_bytes='base64' so the structured-content path serializes bytes as base64 instead of crashing. - StrictJsonSchema.bytes_schema pins the advertised schema to format: binary, so ser_json_bytes='base64' does not leak base64url into outputSchema. Only generated output models are covered; a user-defined BaseModel/TypedDict with a bytes field is unchanged (pydantic forbids a TypeAdapter config override there). Signed-off-by: Udaya Tejas <udayatejas2004@gmail.com>
1 parent 91941ed commit aa95cdf

2 files changed

Lines changed: 99 additions & 4 deletions

File tree

‎src/mcp/server/mcpserver/utilities/func_metadata.py‎

Lines changed: 26 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import base64
12
import functools
23
import inspect
34
import json
@@ -22,7 +23,8 @@
2223
create_model,
2324
)
2425
from pydantic.fields import FieldInfo
25-
from pydantic.json_schema import GenerateJsonSchema, JsonSchemaWarningKind
26+
from pydantic.json_schema import GenerateJsonSchema, JsonSchemaValue, JsonSchemaWarningKind
27+
from pydantic_core import core_schema
2628
from typing_extensions import NotRequired, ReadOnly, TypedDict, deprecated, get_type_hints, is_typeddict
2729
from typing_inspection.introspection import (
2830
UNKNOWN,
@@ -73,6 +75,11 @@ def emit_warning(self, kind: JsonSchemaWarningKind, detail: str) -> None:
7375
# Raise an exception instead of emitting a warning
7476
raise ValueError(f"JSON schema warning: {kind} - {detail}")
7577

78+
def bytes_schema(self, schema: core_schema.BytesSchema) -> JsonSchemaValue:
79+
# `bytes` output is serialized to base64 (see `_create_wrapped_model`), but the advertised
80+
# schema keeps the plain `format: binary` regardless of the serializer's `ser_json_bytes`.
81+
return {"type": "string", "format": "binary"}
82+
7683

7784
_LOCAL_DEFS_PREFIX = "#/$defs/"
7885

@@ -577,7 +584,11 @@ def _create_model_from_class(cls: type[Any], type_hints: dict[str, Any]) -> type
577584
else:
578585
model_fields[field_name] = (field_type, default)
579586

580-
return create_model(cls.__name__, __config__=ConfigDict(from_attributes=True), **model_fields)
587+
return create_model(
588+
cls.__name__,
589+
__config__=ConfigDict(from_attributes=True, ser_json_bytes="base64"),
590+
**model_fields,
591+
)
581592

582593

583594
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]:
618629
"""
619630
model_name = f"{func_name}Output"
620631

621-
return create_model(model_name, result=annotation)
632+
# Serialize any `bytes` field (e.g. a `-> bytes` return) as base64 for JSON so structured output
633+
# does not crash on non-UTF-8 data. Only the serializer is changed, so the advertised schema keeps
634+
# its `format: binary`.
635+
return create_model(
636+
model_name,
637+
__config__=ConfigDict(ser_json_bytes="base64"),
638+
result=annotation,
639+
)
622640

623641

624642
def _convert_to_content(result: Any) -> list[ContentBlock]:
@@ -649,6 +667,11 @@ def _convert_to_content(result: Any) -> list[ContentBlock]:
649667
)
650668
)
651669

670+
# `bytes` advertise `format: binary`; base64-encode them (matching `Image`/`Audio` and the lowlevel
671+
# server) rather than letting `to_json` UTF-8-decode them and raise on non-UTF-8 data.
672+
if isinstance(result, bytes):
673+
result = base64.b64encode(result).decode()
674+
652675
if not isinstance(result, str):
653676
result = pydantic_core.to_json(result, fallback=str, indent=2).decode()
654677

‎tests/server/mcpserver/test_func_metadata.py‎

Lines changed: 73 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
# pyright: reportMissingParameterType=false
44
# pyright: reportUnknownArgumentType=false
55
# pyright: reportUnknownLambdaType=false
6+
import base64
67
from collections.abc import Callable
78
from dataclasses import dataclass
89
from typing import TYPE_CHECKING, Annotated, Any, Final, NamedTuple, TypedDict
@@ -17,7 +18,12 @@
1718
from mcp import MCPDeprecationWarning
1819
from mcp.server.mcpserver import Audio, Image
1920
from mcp.server.mcpserver.exceptions import InvalidSignature
20-
from mcp.server.mcpserver.utilities.func_metadata import ArgModelBase, FuncMetadata, func_metadata
21+
from mcp.server.mcpserver.utilities.func_metadata import (
22+
ArgModelBase,
23+
FuncMetadata,
24+
_convert_to_content,
25+
func_metadata,
26+
)
2127

2228
if TYPE_CHECKING:
2329
from decimal import Decimal
@@ -724,6 +730,72 @@ def func_bytes() -> bytes: # pragma: no cover
724730
}
725731

726732

733+
def test_convert_to_content_base64_encodes_non_utf8_bytes():
734+
"""Non-UTF-8 `bytes` (e.g. PNG magic bytes) base64-encode instead of crashing in `to_json`."""
735+
png = b"\x89PNG\r\n\x1a\n"
736+
content = _convert_to_content(png)
737+
assert len(content) == 1
738+
assert isinstance(content[0], TextContent)
739+
assert content[0].text == base64.b64encode(png).decode()
740+
741+
742+
def test_convert_to_content_base64_encodes_utf8_bytes():
743+
"""UTF-8-decodable `bytes` are base64-encoded too, consistent with the advertised `format: binary`."""
744+
content = _convert_to_content(b"hello")
745+
assert len(content) == 1
746+
assert isinstance(content[0], TextContent)
747+
assert content[0].text == base64.b64encode(b"hello").decode()
748+
749+
750+
def test_structured_output_bytes_return_is_base64():
751+
"""A `-> bytes` tool returning non-UTF-8 bytes yields base64 in both unstructured and structured output."""
752+
png = b"\x89PNG\r\n\x1a\n"
753+
expected = base64.b64encode(png).decode()
754+
755+
def func_bytes() -> bytes:
756+
return png
757+
758+
meta = func_metadata(func_bytes)
759+
# Schema stays `format: binary` (regression guard).
760+
assert meta.output_schema == {
761+
"type": "object",
762+
"properties": {"result": {"title": "Result", "type": "string", "format": "binary"}},
763+
"required": ["result"],
764+
"title": "func_bytesOutput",
765+
}
766+
767+
result = meta.convert_result(png)
768+
assert isinstance(result, CallToolResult)
769+
assert not result.is_error
770+
assert result.structured_content == {"result": expected}
771+
assert isinstance(result.content[0], TextContent)
772+
assert result.content[0].text == expected
773+
774+
775+
def test_structured_output_bytes_field_round_trips():
776+
"""A `bytes` field in an output model (ordinary class) serializes to base64 and round-trips back."""
777+
png = b"\x89PNG\r\n\x1a\n"
778+
779+
class Thumb:
780+
data: bytes
781+
782+
def __init__(self, data: bytes) -> None:
783+
self.data = data
784+
785+
def make_thumb() -> Thumb:
786+
return Thumb(png)
787+
788+
meta = func_metadata(make_thumb)
789+
result = meta.convert_result(Thumb(png))
790+
assert isinstance(result, CallToolResult)
791+
assert not result.is_error
792+
assert result.structured_content == {"data": base64.b64encode(png).decode()}
793+
794+
# The base64 structured content decodes back to the original bytes.
795+
assert result.structured_content is not None
796+
assert base64.b64decode(result.structured_content["data"]) == png
797+
798+
727799
def test_structured_output_generic_types():
728800
"""Test structured output with generic types (list, dict, Union, etc.)"""
729801

0 commit comments

Comments
 (0)