fix: validate MCP image payload blocks

Signed-off-by: Nabin Mulepati <nmulepati@nvidia.com>

Author: nabinchha

architecture/mcp.md

@@ -34,9 +34,10 @@ Scoped to one `ToolConfig`. Provides the interface that `ModelFacade` uses:
 - **`process_completion_response`** — extracts tool calls from a completion, executes them in parallel via `MCPIOService`, returns `ChatMessage` list with results
 - **`refuse_completion_response`** — handles tool-call turn limits (prevents infinite tool loops)
 
-Tool result messages may contain either text or ordered multimodal content blocks. MCP image results and explicit
-base64 image payloads are preserved as canonical `image_url` data URI blocks and translated by provider adapters at
-the API boundary. Models need VLM-capable provider support to interpret those image results semantically.
+Tool result messages may contain either text or ordered multimodal content blocks. MCP image results, and generic
+base64 payloads with `image/*` MIME metadata or image data URI prefixes, are preserved as canonical `image_url` data
+URI blocks and translated by provider adapters at the API boundary. Models need VLM-capable provider support to
+interpret those image results semantically.
 
 ### MCPRegistry
 

packages/data-designer-engine/src/data_designer/engine/mcp/io.py

@@ -552,12 +552,20 @@ def _has_base64_image_payload(item: Any) -> bool:
 
 def _coerce_image_url_block(block: dict[str, Any]) -> dict[str, Any]:
     image_url = block.get("image_url")
-    if not isinstance(image_url, dict):
-        return block
+    if isinstance(image_url, str):
+        image_url = {"url": image_url}
+    elif not isinstance(image_url, dict):
+        raise MCPToolError("MCP image_url block must contain an image_url dict or string.")
 
     url = image_url.get("url")
-    if not isinstance(url, str) or url.startswith(("data:image/", "http://", "https://")):
-        return block
+    if not isinstance(url, str) or not url:
+        raise MCPToolError("MCP image_url block must contain a non-empty string URL.")
+    if url.startswith(("http://", "https://")):
+        return {"type": "image_url", "image_url": image_url}
+    if url.startswith("data:"):
+        _extract_mime_type_from_data_uri(url)
+        _coerce_base64_image_data(url)
+        return {"type": "image_url", "image_url": image_url}
 
     return _build_image_url_block({"base64": url})
 
@@ -594,7 +602,9 @@ def _coerce_image_mime_type(data: str, mime_type: Any) -> str:
 
 def _coerce_base64_image_data(data: str) -> str:
     try:
-        return extract_base64_from_data_uri(data)
+        base64_data = extract_base64_from_data_uri(data)
+        decode_base64_image(base64_data)
+        return base64_data
     except ValueError as exc:
         raise MCPToolError("MCP image content has invalid base64 data.") from exc
 

packages/data-designer-engine/tests/engine/mcp/test_mcp_io.py

@@ -227,10 +227,10 @@ def test_coerce_content_base64_payload_dict_with_media_type() -> None:
     """Explicit media_type payloads do not need image format detection."""
 
     class FakeResult:
-        content = [{"base64": "abc123", "media_type": "image/webp"}]
+        content = [{"base64": "YWJjMTIz", "media_type": "image/webp"}]
 
     assert mcp_io._coerce_tool_result_content(FakeResult()) == [
-        {"type": "image_url", "image_url": {"url": "data:image/webp;base64,abc123"}}
+        {"type": "image_url", "image_url": {"url": "data:image/webp;base64,YWJjMTIz"}}
     ]
 
 
@@ -278,6 +278,17 @@ class FakeResult:
     assert mcp_io._coerce_tool_result_content(FakeResult()) == [{"type": "text", "text": "before"}, image_block]
 
 
+def test_coerce_content_normalizes_image_url_string_shorthand() -> None:
+    """Common shorthand image_url strings are normalized to canonical blocks."""
+
+    class FakeResult:
+        content = [{"type": "image_url", "image_url": "data:image/png;base64,iVBORw0KGgo="}]
+
+    assert mcp_io._coerce_tool_result_content(FakeResult()) == [
+        {"type": "image_url", "image_url": {"url": "data:image/png;base64,iVBORw0KGgo="}}
+    ]
+
+
 def test_coerce_content_normalizes_image_url_raw_base64() -> None:
     """Non-canonical image_url blocks with raw base64 are normalized to data URIs."""
 
@@ -289,6 +300,35 @@ class FakeResult:
     ]
 
 
+@pytest.mark.parametrize(
+    "image_url",
+    [
+        pytest.param(None, id="missing"),
+        pytest.param(123, id="non-dict"),
+        pytest.param({}, id="missing-url"),
+        pytest.param({"url": 123}, id="non-string-url"),
+    ],
+)
+def test_coerce_content_rejects_malformed_image_url_blocks(image_url: object) -> None:
+    """MCP coercion enforces canonical image_url block shape."""
+
+    class FakeResult:
+        content = [{"type": "image_url", "image_url": image_url}]
+
+    with pytest.raises(MCPToolError, match="image_url block"):
+        mcp_io._coerce_tool_result_content(FakeResult())
+
+
+def test_coerce_content_rejects_image_url_invalid_base64_data_uri() -> None:
+    """Canonical data URI image_url blocks still need valid base64 payloads."""
+
+    class FakeResult:
+        content = [{"type": "image_url", "image_url": {"url": "data:image/png;base64,not-base64!!!"}}]
+
+    with pytest.raises(MCPToolError, match="invalid base64"):
+        mcp_io._coerce_tool_result_content(FakeResult())
+
+
 def test_coerce_content_bare_image_object() -> None:
     """Test coercing a bare MCP-like image object."""
 
@@ -307,7 +347,7 @@ def test_coerce_content_mixed_text_image_preserves_order() -> None:
 
     class ImageItem:
         type = "image"
-        data = "abc123"
+        data = "YWJjMTIz"
         mimeType = "image/jpeg"
 
     class TextItem:
@@ -319,14 +359,14 @@ class FakeResult:
             {"type": "text", "text": "before"},
             ImageItem(),
             TextItem(),
-            {"type": "image", "data": "def456", "mime_type": "image/webp"},
+            {"type": "image", "data": "ZGVmNDU2", "mime_type": "image/webp"},
         ]
 
     assert mcp_io._coerce_tool_result_content(FakeResult()) == [
         {"type": "text", "text": "before"},
-        {"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,abc123"}},
+        {"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,YWJjMTIz"}},
         {"type": "text", "text": "after"},
-        {"type": "image_url", "image_url": {"url": "data:image/webp;base64,def456"}},
+        {"type": "image_url", "image_url": {"url": "data:image/webp;base64,ZGVmNDU2"}},
     ]
 
 
@@ -336,13 +376,13 @@ def test_coerce_content_real_mcp_text_and_image_objects() -> None:
     class FakeResult:
         content = [
             TextContent(type="text", text="before"),
-            ImageContent(type="image", data="abc123", mimeType="image/png"),
+            ImageContent(type="image", data="YWJjMTIz", mimeType="image/png"),
             TextContent(type="text", text="after"),
         ]
 
     assert mcp_io._coerce_tool_result_content(FakeResult()) == [
         {"type": "text", "text": "before"},
-        {"type": "image_url", "image_url": {"url": "data:image/png;base64,abc123"}},
+        {"type": "image_url", "image_url": {"url": "data:image/png;base64,YWJjMTIz"}},
         {"type": "text", "text": "after"},
     ]
 
@@ -357,6 +397,16 @@ class FakeResult:
         mcp_io._coerce_tool_result_content(FakeResult())
 
 
+def test_coerce_content_image_with_invalid_base64_fails_clearly() -> None:
+    """Explicit MCP image content must contain valid base64 data."""
+
+    class FakeResult:
+        content = [{"type": "image", "data": "not-base64!!!", "mimeType": "image/png"}]
+
+    with pytest.raises(MCPToolError, match="invalid base64"):
+        mcp_io._coerce_tool_result_content(FakeResult())
+
+
 def test_coerce_content_image_with_non_image_mime_type_fails_clearly() -> None:
     """Explicit MCP image content must not produce non-image image_url data URIs."""