diff --git a/nanobot/agent/tools/mcp.py b/nanobot/agent/tools/mcp.py index a56d3dd5..ea93422e 100644 --- a/nanobot/agent/tools/mcp.py +++ b/nanobot/agent/tools/mcp.py @@ -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.""" diff --git a/tests/tools/test_mcp_tool.py b/tests/tools/test_mcp_tool.py index bb66ab26..a5a8d585 100644 --- a/tests/tools/test_mcp_tool.py +++ b/tests/tools/test_mcp_tool.py @@ -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: