Added a mechanism to extract metadata from MCP tool call response

pydantic_ai_slim/pydantic_ai/mcp.py

@@ -254,9 +254,20 @@ 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:
-                return structured['result']
-            return 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=structured['_meta'])
+                    if structured.get('_meta', None) is not None
+                    else structured['result']
+                )
+            return (
+                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]
         return mapped[0] if len(mapped) == 1 else mapped
@@ -374,35 +385,87 @@ 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)
+            binary_response = messages.BinaryContent(data=base64.b64decode(part.data), media_type=part.mimeType)
+            return (
+                binary_response
+                if metadata is None
+                else messages.ToolReturn(
+                    return_value=f'See file {binary_response.identifier}',
+                    content=[f'This is file {binary_response.identifier}:', binary_response],
+                    metadata=metadata,
+                )
+            )
         elif isinstance(part, mcp_types.AudioContent):
             # NOTE: The FastMCP server doesn't support audio content.
             # See <https://github.com/modelcontextprotocol/python-sdk/issues/952> for more details.
-            return messages.BinaryContent(
-                data=base64.b64decode(part.data), media_type=part.mimeType
-            )  # pragma: no cover
+            binary_response = messages.BinaryContent(data=base64.b64decode(part.data), media_type=part.mimeType)
+            return (  # pragma: no cover
+                binary_response
+                if metadata is None
+                else messages.ToolReturn(
+                    return_value=f'See file {binary_response.identifier}',
+                    content=[f'This is file {binary_response.identifier}:', binary_response],
+                    metadata=metadata,
+                )
+            )
         elif isinstance(part, mcp_types.EmbeddedResource):
             resource = part.resource
-            return self._get_content(resource)
-        elif isinstance(part, mcp_types.ResourceLink):
-            resource_result: mcp_types.ReadResourceResult = await self._client.read_resource(part.uri)
+            response = self._get_content(resource)
             return (
-                self._get_content(resource_result.contents[0])
-                if len(resource_result.contents) == 1
-                else [self._get_content(resource) for resource in resource_result.contents]
+                response
+                if metadata is None
+                else messages.ToolReturn(
+                    return_value=response if isinstance(response, str) else f'See file {response.identifier}',
+                    content=None if isinstance(response, str) else [f'This is file {response.identifier}:', response],
+                    metadata=metadata,
+                )
             )
+        elif isinstance(part, mcp_types.ResourceLink):
+            resource_result: mcp_types.ReadResourceResult = await self._client.read_resource(part.uri)
+            if len(resource_result.contents) == 1:
+                response = self._get_content(resource_result.contents[0])
+                return (
+                    response
+                    if metadata is None
+                    else messages.ToolReturn(
+                        return_value=response if isinstance(response, str) else f'See file {response.identifier}',
+                        content=None
+                        if isinstance(response, str)
+                        else [f'This is file {response.identifier}:', response],
+                        metadata=metadata,
+                    )
+                )
+            else:
+                responses = [self._get_content(resource) for resource in resource_result.contents]  # pragma: no cover
+                return [  # pragma: no cover
+                    response
+                    if isinstance(response, str)
+                    else messages.ToolReturn(
+                        return_value=response if isinstance(response, str) else f'See file {response.identifier}',
+                        content=None
+                        if isinstance(response, str)
+                        else [f'This is file {response.identifier}:', response],
+                    )
+                    for response in responses
+                ]
         else:
             assert_never(part)
 
@@ -875,9 +938,10 @@ def __eq__(self, value: object, /) -> bool:
 ToolResult = (
     str
     | messages.BinaryContent
+    | messages.ToolReturn
     | dict[str, Any]
     | list[Any]
-    | Sequence[str | messages.BinaryContent | dict[str, Any] | list[Any]]
+    | Sequence[str | messages.BinaryContent | messages.ToolReturn | dict[str, Any] | list[Any]]
 )
 """The result type of an MCP tool call."""
 

tests/mcp_server.py

@@ -68,6 +68,45 @@ async def get_image_resource_link() -> ResourceLink:
     )
 
 
+@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 with attached metadata.
+    """
+    if n <= 0:
+        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:
+        if n % 2 == 0:
+            n = n // 2
+        else:
+            n = 3 * n + 1
+        sequence.append(n)
+
+    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')
 async def kiwi_resource() -> bytes:
     return Path(__file__).parent.joinpath('assets/kiwi.png').read_bytes()

tests/test_mcp.py

@@ -26,6 +26,7 @@
 from pydantic_ai.agent import Agent
 from pydantic_ai.exceptions import ModelRetry, UnexpectedModelBehavior, UserError
 from pydantic_ai.mcp import MCPServerStreamableHTTP, load_mcp_servers
+from pydantic_ai.messages import ToolReturn
 from pydantic_ai.models import Model
 from pydantic_ai.models.test import TestModel
 from pydantic_ai.tools import RunContext
@@ -77,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(18)
+        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.')
@@ -87,6 +88,36 @@ async def test_stdio_server(run_context: RunContext[int]):
         assert result == snapshot(32.0)
 
 
+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(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.')
+
+        result = await server.direct_call_tool('get_collatz_conjecture', {'n': 7})
+        assert isinstance(result, ToolReturn)
+        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():
     server = MCPServerStdio('python', ['-m', 'tests.mcp_server'])
     async with server:
@@ -138,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(18)
+        assert len(tools) == snapshot(20)
 
 
 async def test_process_tool_call(run_context: RunContext[int]) -> int: