Skip to content

Commit a2d08e8

Browse files
committed
fix: wrap structured union outputs explicitly
1 parent 161834d commit a2d08e8

2 files changed

Lines changed: 21 additions & 1 deletion

File tree

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -478,7 +478,7 @@ def _create_wrapped_model(func_name: str, annotation: Any) -> type[BaseModel]:
478478
"""
479479
model_name = f"{func_name}Output"
480480

481-
return create_model(model_name, result=annotation)
481+
return create_model(model_name, result=(annotation, ...))
482482

483483

484484
def _create_dict_model(func_name: str, dict_annotation: Any) -> type[BaseModel]:

tests/server/mcpserver/test_func_metadata.py

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -680,6 +680,9 @@ def func_union() -> str | int: # pragma: no cover
680680
def func_optional() -> str | None: # pragma: no cover
681681
return None
682682

683+
def func_builtin_union() -> dict | list | str: # pragma: no cover
684+
return {"hello": "world"}
685+
683686
# Test list
684687
meta = func_metadata(func_list_str)
685688
assert meta.output_schema == {
@@ -715,6 +718,23 @@ def func_optional() -> str | None: # pragma: no cover
715718
"title": "func_optionalOutput",
716719
}
717720

721+
meta = func_metadata(func_builtin_union)
722+
assert meta.output_schema == {
723+
"type": "object",
724+
"properties": {
725+
"result": {
726+
"title": "Result",
727+
"anyOf": [
728+
{"type": "object", "additionalProperties": True},
729+
{"type": "array", "items": {}},
730+
{"type": "string"},
731+
],
732+
}
733+
},
734+
"required": ["result"],
735+
"title": "func_builtin_unionOutput",
736+
}
737+
718738

719739
def test_structured_output_dataclass():
720740
"""Test structured output with dataclass return types"""

0 commit comments

Comments
 (0)