chore: Improved metadata parsing for both structured content and TextContent.

Author: anirbanbasu

pydantic_ai_slim/pydantic_ai/mcp.py

@@ -270,12 +270,6 @@ async def direct_call_tool(
             )
 
         mapped = [await self._map_tool_result_part(part) for part in result.content]
-        if result.meta:
-            return (
-                messages.ToolReturn(return_value=mapped[0], metadata=result.meta)
-                if len(mapped) == 1
-                else [messages.ToolReturn(return_value=mapped_item, metadata=result.meta) for mapped_item in mapped]
-            )
         return mapped[0] if len(mapped) == 1 else mapped
 
     async def call_tool(
@@ -391,17 +385,23 @@ async def _sampling_callback(
 
     async def _map_tool_result_part(
         self, part: mcp_types.ContentBlock
-    ) -> str | messages.BinaryContent | dict[str, Any] | list[Any]:
+    ) -> str | messages.ToolReturn | messages.BinaryContent | dict[str, Any] | list[Any]:
         # See https://github.com/jlowin/fastmcp/blob/main/docs/servers/tools.mdx#return-values
 
+        # Let's also check for metadata but it can be present in not just TextContent
+        metadata: dict[str, Any] | None = part.meta
         if isinstance(part, mcp_types.TextContent):
             text = part.text
             if text.startswith(('[', '{')):
                 try:
-                    return pydantic_core.from_json(text)
+                    return (
+                        pydantic_core.from_json(text)
+                        if metadata is None
+                        else messages.ToolReturn(return_value=pydantic_core.from_json(text), metadata=metadata)
+                    )
                 except ValueError:
                     pass
-            return text
+            return text if metadata is None else messages.ToolReturn(return_value=text, metadata=metadata)
         elif isinstance(part, mcp_types.ImageContent):
             return messages.BinaryContent(data=base64.b64decode(part.data), media_type=part.mimeType)
         elif isinstance(part, mcp_types.AudioContent):

tests/mcp_server.py

@@ -68,18 +68,20 @@ async def get_image_resource_link() -> ResourceLink:
     )
 
 
-@mcp.tool(annotations=ToolAnnotations(title='Collatz Conjecture sequence generator'))
-async def collatz_conjecture(n: int) -> dict[str, Any]:
+@mcp.tool(structured_output=False, annotations=ToolAnnotations(title='Collatz Conjecture sequence generator'))
+async def get_collatz_conjecture(n: int) -> TextContent:
     """Generate the Collatz conjecture sequence for a given number.
     This tool attaches response metadata.
 
     Args:
         n: The starting number for the Collatz sequence.
     Returns:
-        A list representing the Collatz sequence.
+        A list representing the Collatz sequence with attached metadata.
     """
     if n <= 0:
-        raise ValueError('Startig number for the Collatz conjecture must be a positive integer.')
+        raise ValueError('Starting number for the Collatz conjecture must be a positive integer.')
+
+    input_param_n = n  # store the original input value
 
     sequence = [n]
     while n != 1:
@@ -88,10 +90,21 @@ async def collatz_conjecture(n: int) -> dict[str, Any]:
         else:
             n = 3 * n + 1
         sequence.append(n)
-    response: dict[str, Any] = {'result': str(sequence)}
-    # attach metadata to the response
-    response['_meta'] = {'pydantic_ai': {'tool': 'collatz_conjecture', 'length': len(sequence)}}
-    return response
+
+    return TextContent(
+        type='text',
+        text=str(sequence),
+        _meta={'pydantic_ai': {'tool': 'collatz_conjecture', 'n': input_param_n, 'length': len(sequence)}},
+    )
+
+
+@mcp.tool()
+async def get_structured_text_content_with_metadata() -> dict[str, Any]:
+    """Return structured dict with metadata."""
+    return {
+        'result': 'This is some text content.',
+        '_meta': {'pydantic_ai': {'source': 'get_structured_text_content_with_metadata'}},
+    }
 
 
 @mcp.resource('resource://kiwi.png', mime_type='image/png')

tests/test_mcp.py

@@ -78,7 +78,7 @@ async def test_stdio_server(run_context: RunContext[int]):
     server = MCPServerStdio('python', ['-m', 'tests.mcp_server'])
     async with server:
         tools = [tool.tool_def for tool in (await server.get_tools(run_context)).values()]
-        assert len(tools) == snapshot(19)
+        assert len(tools) == snapshot(20)
         assert tools[0].name == 'celsius_to_fahrenheit'
         assert isinstance(tools[0].description, str)
         assert tools[0].description.startswith('Convert Celsius to Fahrenheit.')
@@ -92,16 +92,30 @@ async def test_tool_response_metadata(run_context: RunContext[int]):
     server = MCPServerStdio('python', ['-m', 'tests.mcp_server'])
     async with server:
         tools = [tool.tool_def for tool in (await server.get_tools(run_context)).values()]
-        assert len(tools) == snapshot(19)
-        assert tools[4].name == 'collatz_conjecture'
+        assert len(tools) == snapshot(20)
+        assert tools[4].name == 'get_collatz_conjecture'
         assert isinstance(tools[4].description, str)
         assert tools[4].description.startswith('Generate the Collatz conjecture sequence for a given number.')
 
-        # Test calling the Collatz conjecture generator tool
-        result = await server.direct_call_tool('collatz_conjecture', {'n': 7})
+        result = await server.direct_call_tool('get_collatz_conjecture', {'n': 7})
         assert isinstance(result, ToolReturn)
-        assert result.return_value == '[7, 22, 11, 34, 17, 52, 26, 13, 40, 20, 10, 5, 16, 8, 4, 2, 1]'
-        assert result.metadata == {'pydantic_ai': {'tool': 'collatz_conjecture', 'length': 17}}
+        assert result.return_value == snapshot([7, 22, 11, 34, 17, 52, 26, 13, 40, 20, 10, 5, 16, 8, 4, 2, 1])
+        assert result.metadata == snapshot({'pydantic_ai': {'tool': 'collatz_conjecture', 'n': 7, 'length': 17}})
+
+
+async def test_tool_structured_response_metadata(run_context: RunContext[int]):
+    server = MCPServerStdio('python', ['-m', 'tests.mcp_server'])
+    async with server:
+        tools = [tool.tool_def for tool in (await server.get_tools(run_context)).values()]
+        assert len(tools) == snapshot(20)
+        assert tools[5].name == 'get_structured_text_content_with_metadata'
+        assert isinstance(tools[5].description, str)
+        assert tools[5].description.startswith('Return structured dict with metadata.')
+
+        result = await server.direct_call_tool('get_structured_text_content_with_metadata', {})
+        assert isinstance(result, ToolReturn)
+        assert result.return_value == 'This is some text content.'
+        assert result.metadata == snapshot({'pydantic_ai': {'source': 'get_structured_text_content_with_metadata'}})
 
 
 async def test_reentrant_context_manager():
@@ -155,7 +169,7 @@ async def test_stdio_server_with_cwd(run_context: RunContext[int]):
     server = MCPServerStdio('python', ['mcp_server.py'], cwd=test_dir)
     async with server:
         tools = await server.get_tools(run_context)
-        assert len(tools) == snapshot(19)
+        assert len(tools) == snapshot(20)
 
 
 async def test_process_tool_call(run_context: RunContext[int]) -> int: