feat(mcp): deliver image content from MCP tools as artifacts

MCPToolWrapper.execute only handled TextContent; every other block was
rendered with str(block). An MCP ImageContent block therefore became a
large base64 string embedded in the tool result, which (a) was truncated
by max_tool_result_chars, corrupting the data, and (b) could never reach a
channel because it was plain text, not an image artifact.

Decode ImageContent (and EmbeddedResource blobs with an image/* MIME type)
and persist them via store_generated_image_artifact, returning the same
compact {artifacts, next_step} JSON the built-in image_generation tool
produces. The base64 stays out of the model context; the model delivers the
saved file via the message tool's media parameter.
This commit is contained in:
codedragon 2026-06-26 13:00:40 -04:00 committed by Xubin Ren
parent 194e9d5f5f
commit 67ce6822ca
2 changed files with 164 additions and 10 deletions

View File

@ -1,6 +1,7 @@
"""MCP client: connects to MCP servers and wraps their tools as native nanobot tools."""
import asyncio
import json
import os
import re
import shutil
@ -313,6 +314,52 @@ class _MCPWrapperBase(Tool):
return True
def _image_block_data_url(block: Any, types: Any) -> str | None:
"""Return a base64 ``data:`` URL for an MCP image-bearing content block.
Handles ``ImageContent`` directly and ``EmbeddedResource`` wrapping a binary
blob with an ``image/*`` MIME type. Returns ``None`` for anything else.
``getattr`` guards keep this safe when the installed/faked ``mcp`` SDK does
not expose a given type.
"""
image_cls = getattr(types, "ImageContent", None)
if image_cls is not None and isinstance(block, image_cls):
mime = getattr(block, "mimeType", None) or "image/png"
return f"data:{mime};base64,{block.data}"
embedded_cls = getattr(types, "EmbeddedResource", None)
blob_cls = getattr(types, "BlobResourceContents", None)
if embedded_cls is not None and isinstance(block, embedded_cls):
resource = getattr(block, "resource", None)
if blob_cls is not None and isinstance(resource, blob_cls):
mime = getattr(resource, "mimeType", None) or ""
if isinstance(mime, str) and mime.startswith("image/"):
return f"data:{mime};base64,{resource.blob}"
return None
def _mcp_image_tool_result(text_parts: list[str], artifacts: list[dict[str, Any]]) -> str:
"""Build the compact tool result for an MCP call that returned image(s).
The base64 stays out of the model context entirely only artifact paths and
metadata are returned, so the result is small and the channel can deliver the
saved file via the message tool.
"""
payload: dict[str, Any] = {
"artifacts": artifacts,
"next_step": (
"These images were returned by an MCP tool and saved as local artifacts. "
"Call the message tool with the artifact 'path' values in the media "
"parameter to deliver the images to the user. Do not paste base64 or raw "
"paths into your reply unless the user asks for debug details."
),
}
text = "\n".join(part for part in text_parts if part)
if text:
payload["text"] = text
return json.dumps(payload, ensure_ascii=False)
class MCPToolWrapper(_MCPWrapperBase):
"""Wraps a single MCP server tool as a nanobot Tool."""
@ -340,8 +387,6 @@ class MCPToolWrapper(_MCPWrapperBase):
return self._parameters
async def execute(self, **kwargs: Any) -> str:
from mcp import types
retried_transient = False
refreshed_session = False
while True:
@ -396,17 +441,63 @@ class MCPToolWrapper(_MCPWrapperBase):
)
return f"(MCP tool call failed: {type(exc).__name__})"
else:
# Success — extract result
parts = []
for block in result.content:
if isinstance(block, types.TextContent):
parts.append(block.text)
else:
parts.append(str(block))
return "\n".join(parts) or "(no output)"
# Success — extract text and persist any image content as artifacts.
return self._render_call_result(result.content, kwargs)
return "(MCP tool call failed)" # Unreachable, but satisfies type checkers
def _render_call_result(self, content: Any, arguments: Mapping[str, Any]) -> str:
"""Turn MCP content blocks into a tool result string.
Text is concatenated as before. Image blocks are decoded and saved as
local artifacts (mirroring the built-in image generation tool) so the
model can deliver them via the message tool instead of trying to forward
base64 which would be truncated and bloat the context window.
"""
from mcp import types
text_parts: list[str] = []
artifacts: list[dict[str, Any]] = []
for block in content:
if isinstance(block, types.TextContent):
text_parts.append(block.text)
continue
data_url = _image_block_data_url(block, types)
if data_url is not None:
stored = self._store_image_block(data_url, arguments)
if stored is not None:
artifacts.append(stored)
else:
text_parts.append("(MCP tool returned an image that could not be stored)")
continue
text_parts.append(str(block))
if artifacts:
return _mcp_image_tool_result(text_parts, artifacts)
return "\n".join(text_parts) or "(no output)"
def _store_image_block(
self, data_url: str, arguments: Mapping[str, Any]
) -> dict[str, Any] | None:
"""Persist one image data URL as an artifact; return its metadata or None."""
from nanobot.utils.artifacts import ArtifactError, store_generated_image_artifact
try:
return store_generated_image_artifact(
data_url,
prompt=str(arguments.get("prompt") or ""),
model=str(arguments.get("model") or ""),
save_dir="generated",
provider=f"mcp:{self._server_name}",
)
except (ArtifactError, OSError) as exc:
logger.warning(
"MCP tool '{}' returned an image that could not be stored: {}",
self._name,
exc,
)
return None
class MCPResourceWrapper(_MCPWrapperBase):
"""Wraps an MCP resource URI as a read-only nanobot Tool."""

View File

@ -1,8 +1,10 @@
from __future__ import annotations
import asyncio
import json
import sys
from contextlib import asynccontextmanager
from pathlib import Path
from types import ModuleType, SimpleNamespace
import httpx
@ -36,6 +38,12 @@ class _FakeBlobResourceContents:
self.blob = blob
class _FakeImageContent:
def __init__(self, data: str, mime_type: str = "image/png") -> None:
self.data = data
self.mimeType = mime_type
@pytest.fixture
def fake_mcp_runtime() -> dict[str, object | None]:
return {"session": None}
@ -50,6 +58,7 @@ def _fake_mcp_module(
TextContent=_FakeTextContent,
TextResourceContents=_FakeTextResourceContents,
BlobResourceContents=_FakeBlobResourceContents,
ImageContent=_FakeImageContent,
)
class _FakeStdioServerParameters:
@ -295,6 +304,60 @@ async def test_execute_returns_text_blocks() -> None:
assert result == "hello\n42"
# Smallest valid 1x1 PNG, base64 without the data: prefix.
_PNG_B64 = (
"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8"
"/x8AAwMCAO+/p9sAAAAASUVORK5CYII="
)
@pytest.mark.asyncio
async def test_execute_persists_image_block_as_artifact(tmp_path: Path) -> None:
from nanobot.config.loader import set_config_path
set_config_path(tmp_path / "config.json")
async def call_tool(_name: str, arguments: dict) -> object:
return SimpleNamespace(
content=[
_FakeTextContent("here you go"),
_FakeImageContent(_PNG_B64, "image/png"),
]
)
wrapper = _make_wrapper(SimpleNamespace(call_tool=call_tool))
result = await wrapper.execute(prompt="a cat", model="sdxl")
payload = json.loads(result)
assert payload["text"] == "here you go"
assert len(payload["artifacts"]) == 1
artifact = payload["artifacts"][0]
assert artifact["mime"] == "image/png"
assert artifact["prompt"] == "a cat"
assert artifact["provider"] == "mcp:test"
assert Path(artifact["path"]).is_file()
# The base64 payload must NOT leak into the model-facing result.
assert _PNG_B64 not in result
assert "message tool" in payload["next_step"]
@pytest.mark.asyncio
async def test_execute_notes_unstorable_image_block(tmp_path: Path) -> None:
from nanobot.config.loader import set_config_path
set_config_path(tmp_path / "config.json")
async def call_tool(_name: str, arguments: dict) -> object:
return SimpleNamespace(content=[_FakeImageContent("not-valid-base64!!", "image/png")])
wrapper = _make_wrapper(SimpleNamespace(call_tool=call_tool))
result = await wrapper.execute()
assert result == "(MCP tool returned an image that could not be stored)"
@pytest.mark.asyncio
async def test_execute_returns_timeout_message() -> None:
async def call_tool(_name: str, arguments: dict) -> object: