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
Xubin Ren
+101
-10
@@ -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."""
|
||||
|
||||
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user