Python: fix(core): rebuild nested Content of image generation and MCP tool results in Content.from_dict

python/packages/core/agent_framework/_types.py

@@ -24,7 +24,20 @@
 from copy import deepcopy
 from datetime import datetime
 from inspect import isawaitable, ismethod
-from typing import TYPE_CHECKING, Any, ClassVar, Final, Generic, Literal, NamedTuple, NewType, TypeAlias, cast, overload
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    ClassVar,
+    Final,
+    Generic,
+    Literal,
+    NamedTuple,
+    NewType,
+    TypeAlias,
+    cast,
+    get_args,
+    overload,
+)
 
 from typing_extensions import Required, TypedDict
 
@@ -442,6 +455,8 @@ def _restore_compaction_annotation_in_additional_properties(
     "code_interpreter_tool_result",
 })
 
+_CONTENT_TYPES: Final[frozenset[str]] = frozenset(get_args(ContentType))
+
 
 class TextSpanRegion(TypedDict, total=False):
     """TypedDict representation of a text span region annotation."""
@@ -1906,6 +1921,13 @@ def from_dict(cls: type[ContentT], data: Mapping[str, Any]) -> ContentT:
             remaining["outputs"] = [cls.from_dict(item) if isinstance(item, dict) else item for item in output_items]  # type: ignore[reportUnknownVariableType]
         if (content_items := remaining.get("items")) and isinstance(content_items, list):
             remaining["items"] = [cls.from_dict(item) if isinstance(item, dict) else item for item in content_items]  # type: ignore[reportUnknownVariableType]
+        # A single nested Content (e.g., outputs in image_generation_tool_result) and Content items in a tool output
+        # (e.g., output in mcp_server_tool_result). These fields can also hold plain values such as tool JSON.
+        # A list in outputs was converted above and passes through unchanged.
+        if (single_output := remaining.get("outputs")) is not None:
+            remaining["outputs"] = cls._from_serialized_value(single_output)
+        if (tool_output := remaining.get("output")) and isinstance(tool_output, list):
+            remaining["output"] = [cls._from_serialized_value(item) for item in cast(list[Any], tool_output)]
 
         return cls(
             type=content_type,
@@ -1915,6 +1937,30 @@ def from_dict(cls: type[ContentT], data: Mapping[str, Any]) -> ContentT:
             **remaining,
         )
 
+    @classmethod
+    def _from_serialized_value(cls: type[ContentT], value: Any) -> Any:
+        """Rebuild a Content from its serialized mapping, or return the value unchanged.
+
+        A mapping is rebuilt only when it has a known content type and the rebuilt Content serializes back to the
+        same mapping, with or without ``None`` values. Any other value, such as tool JSON that only shares a
+        content type name, is kept as it is. A mapping identical to a serialized Content (including its
+        ``additional_properties``) cannot be told apart from one and is rebuilt.
+        """
+        if not isinstance(value, dict):
+            return value
+        mapping = cast(dict[str, Any], value)
+        content_type = mapping.get("type")
+        if not isinstance(content_type, str) or content_type not in _CONTENT_TYPES:
+            return mapping
+        try:
+            content = cls.from_dict(mapping)
+            if content.to_dict() == mapping or content.to_dict(exclude_none=False) == mapping:
+                return content
+        except Exception:
+            # The mapping is not a serialized Content (for example tool JSON with other fields), so keep it.
+            return mapping
+        return mapping
+
     def __add__(self, other: Content) -> Content:
         """Concatenate or merge two Content instances."""
         if not isinstance(other, Content):

python/packages/core/tests/core/test_types.py

@@ -4705,6 +4705,77 @@ def test_function_result_items_roundtrip_via_dict():
     assert restored.items[1].type == "data"
 
 
+@pytest.mark.parametrize("serialization", ["dict", "dict_with_none", "json"])
+def test_tool_result_nested_content_roundtrips(serialization: str):
+    """Test that nested Content in image generation and MCP server tool results survives a round trip."""
+    image = Content.from_uri("data:image/png;base64,iVBORw0KGgo=")
+    output = [
+        Content.from_text("42", annotations=[{"type": "citation", "url": "https://example.com/source"}]),
+        Content.from_uri("https://example.com/report.pdf", media_type="application/pdf"),
+        Content.from_hosted_file("file_1"),
+        Content.from_function_result(call_id="call_1", result=[Content.from_text("done")]),
+    ]
+    message = Message(
+        "assistant",
+        [
+            Content.from_image_generation_tool_result(image_id="ig_1", outputs=image),
+            Content.from_mcp_server_tool_result("mcp_1", output=output),
+        ],
+    )
+
+    if serialization == "dict":
+        restored = Message.from_dict(message.to_dict())
+    elif serialization == "dict_with_none":
+        restored = Message.from_dict(message.to_dict(exclude_none=False))
+    else:
+        restored = Message.from_json(message.to_json())
+    restored_image, restored_mcp = restored.contents
+
+    assert isinstance(restored_image.outputs, Content)
+    assert restored_image.outputs == image
+    assert restored_mcp.output is not None
+    assert all(isinstance(item, Content) for item in restored_mcp.output)
+    assert restored_mcp.output == output
+
+
+@pytest.mark.parametrize(
+    "tool_json",
+    [
+        {"type": "text", "value": 1},
+        {"type": "text", "text": "x"},
+        {"type": "data", "data": "AAAA", "media_type": "image/png"},
+        {"type": "text", "text": "x", "annotations": {"audience": ["user"]}},
+        {"type": "text", "function_call": {"a": 1}},
+    ],
+    ids=["unknown_field", "text_shape", "data_with_str_data", "mcp_annotations", "invalid_nested_content"],
+)
+def test_tool_result_json_with_content_type_name_roundtrips(tool_json: dict[str, Any]):
+    """Test that tool JSON whose type is a content type name is kept as it is in a round trip."""
+    mcp_result = Content.from_mcp_server_tool_result("mcp_1", output=[Content.from_text("42"), tool_json])
+    image_result = Content.from_image_generation_tool_result(image_id="ig_1", outputs=tool_json)
+
+    restored_mcp = Content.from_dict(mcp_result.to_dict())
+    restored_image = Content.from_dict(image_result.to_dict())
+
+    assert restored_mcp.output == [Content.from_text("42"), tool_json]
+    assert isinstance(restored_mcp.output[1], dict)
+    assert restored_image.outputs == tool_json
+    assert restored_mcp.to_dict()["output"][1] == tool_json
+    assert restored_image.to_dict()["outputs"] == tool_json
+
+
+def test_tool_result_plain_outputs_roundtrip_via_dict():
+    """Test that output values that are not serialized Content stay unchanged in a round trip."""
+    output = ["plain", {"type": "object", "value": 1}, {"type": ["x"]}, {"answer": 42}]
+    mcp_result = Content.from_mcp_server_tool_result("mcp_1", output=output)
+    image_result = Content.from_image_generation_tool_result(
+        image_id="ig_1", outputs={"url": "https://example.com/a.png"}
+    )
+
+    assert Content.from_dict(mcp_result.to_dict()).output == output
+    assert Content.from_dict(image_result.to_dict()).outputs == {"url": "https://example.com/a.png"}
+
+
 def test_from_function_result_with_non_content_list():
     """Test Content.from_function_result with a list of non-Content objects falls back to str."""
     result = Content.from_function_result(call_id="test-789", result=["hello", "world"])