From be6e0c12a9a2eebe29af4a7ec057ac37abeb2802 Mon Sep 17 00:00:00 2001
From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com>
Date: Fri, 29 Nov 2024 22:27:42 -0500
Subject: [PATCH] Improve handling of image results from tools

---
 src/fastmcp/server.py | 17 +++++++++++++++++
 tests/test_server.py  | 33 +++++++++++++++++++++++++++++++++
 2 files changed, 50 insertions(+)

diff --git a/src/fastmcp/server.py b/src/fastmcp/server.py
index 15403de..9564f5c 100644
--- a/src/fastmcp/server.py
+++ b/src/fastmcp/server.py
@@ -160,6 +160,23 @@ def _convert_to_content(
         if isinstance(value, (list, tuple)):
             if all(isinstance(x, (TextContent, ImageContent)) for x in value):
                 return value
+            # Handle mixed content including Image objects
+            result = []
+            for item in value:
+                if isinstance(item, (TextContent, ImageContent)):
+                    result.append(item)
+                elif isinstance(item, Image):
+                    result.append(item.to_image_content())
+                else:
+                    result.append(
+                        TextContent(
+                            type="text",
+                            text=json.dumps(
+                                item, indent=2, default=pydantic.json.pydantic_encoder
+                            ),
+                        )
+                    )
+            return result
 
         # Single content type
         if isinstance(value, (TextContent, ImageContent)):
diff --git a/tests/test_server.py b/tests/test_server.py
index 92d1949..0a491a2 100644
--- a/tests/test_server.py
+++ b/tests/test_server.py
@@ -161,6 +161,39 @@ async def test_tool_mixed_content(self):
             assert result.content[1].mimeType == "image/png"
             assert result.content[1].data == "abc"
 
+    async def test_tool_mixed_list_with_image(self, tmp_path: Path):
+        """Test that lists containing Image objects and other types are handled correctly"""
+        # Create a test image
+        image_path = tmp_path / "test.png"
+        image_path.write_bytes(b"test image data")
+
+        def mixed_list_fn() -> list:
+            return [
+                "text message",
+                Image(image_path),
+                {"key": "value"},
+                TextContent(type="text", text="direct content"),
+            ]
+
+        mcp = FastMCP()
+        mcp.add_tool(mixed_list_fn)
+        async with client_session(mcp._mcp_server) as client:
+            result = await client.call_tool("mixed_list_fn", {})
+            assert len(result.content) == 4
+            # Check text conversion
+            assert result.content[0].type == "text"
+            assert '"text message"' in result.content[0].text
+            # Check image conversion
+            assert result.content[1].type == "image"
+            assert result.content[1].mimeType == "image/png"
+            assert base64.b64decode(result.content[1].data) == b"test image data"
+            # Check dict conversion
+            assert result.content[2].type == "text"
+            assert '"key": "value"' in result.content[2].text
+            # Check direct TextContent
+            assert result.content[3].type == "text"
+            assert result.content[3].text == "direct content"
+
 
 class TestServerResources:
     async def test_text_resource(self):