fix: Tests seem to work with the MCP metadata but are these exhaustive?

Author: anirbanbasu

pydantic_ai_slim/pydantic_ai/mcp.py

@@ -254,20 +254,23 @@ async def direct_call_tool(
         ):
             # The MCP SDK wraps primitives and generic types like list in a `result` key, but we want to use the raw value returned by the tool function.
             # See https://github.com/modelcontextprotocol/python-sdk#structured-output
-            if isinstance(structured, dict) and len(structured) == 1 and 'result' in structured:
+            if isinstance(structured, dict) and (
+                (len(structured) == 1 and 'result' in structured)
+                or (len(structured) == 2 and 'result' in structured and '_meta' in structured)
+            ):
                 return (
-                    messages.ToolReturn(return_value=structured['result'], metadata=result.meta)
-                    if getattr(result, '_meta', None) is not None
+                    messages.ToolReturn(return_value=structured['result'], metadata=structured['_meta'])
+                    if structured.get('_meta', None) is not None
                     else structured['result']
                 )
             return (
-                messages.ToolReturn(return_value=structured, metadata=result.meta)
-                if getattr(result, '_meta', None) is not None
+                messages.ToolReturn(return_value=structured, metadata=structured['_meta'])
+                if structured.get('_meta', None) is not None
                 else structured
             )
 
         mapped = [await self._map_tool_result_part(part) for part in result.content]
-        if getattr(result, '_meta', None) is not None:
+        if result.meta:
             return (
                 messages.ToolReturn(return_value=mapped[0], metadata=result.meta)
                 if len(mapped) == 1

tests/mcp_server.py

@@ -69,7 +69,7 @@ async def get_image_resource_link() -> ResourceLink:
 
 
 @mcp.tool(annotations=ToolAnnotations(title='Collatz Conjecture sequence generator'))
-async def collatz_conjecture(n: int) -> TextContent:
+async def collatz_conjecture(n: int) -> dict[str, Any]:
     """Generate the Collatz conjecture sequence for a given number.
     This tool attaches response metadata.
 
@@ -88,11 +88,9 @@ async def collatz_conjecture(n: int) -> TextContent:
         else:
             n = 3 * n + 1
         sequence.append(n)
-    response = TextContent(type='text', text=str(sequence))
+    response: dict[str, Any] = {'result': str(sequence)}
     # attach metadata to the response
-    if response.meta is None:
-        response.meta = {}
-    response.meta['pydantic_ai'] = {'tool': 'collatz_conjecture', 'length': len(sequence)}
+    response['_meta'] = {'pydantic_ai': {'tool': 'collatz_conjecture', 'length': len(sequence)}}
     return response
 
 

tests/test_mcp.py

@@ -100,7 +100,6 @@ async def test_tool_response_metadata(run_context: RunContext[int]):
         # Test calling the Collatz conjecture generator tool
         result = await server.direct_call_tool('collatz_conjecture', {'n': 7})
         assert isinstance(result, ToolReturn)
-        assert isinstance(result.return_value, str)
         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}}