Add process_tool_call hook to MCP servers to modify tool args, metadata, and return value (#2000)

Co-authored-by: Hayk Martiros <[email protected]>
Co-authored-by: Douwe Maan <[email protected]>

Author: 3 people

docs/mcp/client.md

@@ -166,6 +166,48 @@ async def main():
 
 1. See [MCP Run Python](run-python.md) for more information.
 
+## Tool call customisation
+
+The MCP servers provide the ability to set a `process_tool_call` which allows
+the customisation of tool call requests and their responses.
+
+A common use case for this is to inject metadata to the requests which the server
+call needs.
+
+```python {title="mcp_process_tool_call.py" py="3.10"}
+from typing import Any
+
+from pydantic_ai import Agent
+from pydantic_ai.mcp import CallToolFunc, MCPServerStdio, ToolResult
+from pydantic_ai.models.test import TestModel
+from pydantic_ai.tools import RunContext
+
+
+async def process_tool_call(
+    ctx: RunContext[int],
+    call_tool: CallToolFunc,
+    tool_name: str,
+    args: dict[str, Any],
+) -> ToolResult:
+    """A tool call processor that passes along the deps."""
+    return await call_tool(tool_name, args, metadata={'deps': ctx.deps})
+
+
+server = MCPServerStdio('python', ['-m', 'tests.mcp_server'], process_tool_call=process_tool_call)
+agent = Agent(
+    model=TestModel(call_tools=['echo_deps']),
+    deps_type=int,
+    mcp_servers=[server]
+)
+
+
+async def main():
+    async with agent.run_mcp_servers():
+        result = await agent.run('Echo with deps set to 42', deps=42)
+    print(result.output)
+    #> {"echo_deps":{"echo":"This is an echo message","deps":42}}
+```
+
 ## Using Tool Prefixes to Avoid Naming Conflicts
 
 When connecting to multiple MCP servers that might provide tools with the same name, you can use the `tool_prefix` parameter to avoid naming conflicts. This parameter adds a prefix to all tool names from a specific server.

pydantic_ai_slim/pydantic_ai/_agent_graph.py

@@ -762,7 +762,12 @@ async def run_tool(ctx: RunContext[DepsT], **args: Any) -> Any:
         # some weird edge case occurs.
         if not server.is_running:  # pragma: no cover
             raise exceptions.UserError(f'MCP server is not running: {server}')
-        result = await server.call_tool(tool_name, args)
+
+        if server.process_tool_call is not None:
+            result = await server.process_tool_call(ctx, server.call_tool, tool_name, args)
+        else:
+            result = await server.call_tool(tool_name, args)
+
         return result
 
     for server in ctx.deps.mcp_servers:

pydantic_ai_slim/pydantic_ai/mcp.py

@@ -4,7 +4,7 @@
 import functools
 import json
 from abc import ABC, abstractmethod
-from collections.abc import AsyncIterator, Sequence
+from collections.abc import AsyncIterator, Awaitable, Sequence
 from contextlib import AbstractAsyncContextManager, AsyncExitStack, asynccontextmanager
 from dataclasses import dataclass
 from pathlib import Path
@@ -20,18 +20,23 @@
 from mcp.types import (
     AudioContent,
     BlobResourceContents,
+    CallToolRequest,
+    CallToolRequestParams,
+    CallToolResult,
+    ClientRequest,
     Content,
     EmbeddedResource,
     ImageContent,
     LoggingLevel,
+    RequestParams,
     TextContent,
     TextResourceContents,
 )
 from typing_extensions import Self, assert_never, deprecated
 
 from pydantic_ai.exceptions import ModelRetry
 from pydantic_ai.messages import BinaryContent
-from pydantic_ai.tools import ToolDefinition
+from pydantic_ai.tools import RunContext, ToolDefinition
 
 try:
     from mcp.client.session import ClientSession
@@ -61,6 +66,9 @@ class MCPServer(ABC):
     e.g. if `tool_prefix='foo'`, then a tool named `bar` will be registered as `foo_bar`
     """
 
+    process_tool_call: ProcessToolCallback | None = None
+    """Hook to customize tool calling and optionally pass extra metadata."""
+
     _client: ClientSession
     _read_stream: MemoryObjectReceiveStream[SessionMessage | Exception]
     _write_stream: MemoryObjectSendStream[SessionMessage]
@@ -114,13 +122,17 @@ async def list_tools(self) -> list[ToolDefinition]:
         ]
 
     async def call_tool(
-        self, tool_name: str, arguments: dict[str, Any]
-    ) -> str | BinaryContent | dict[str, Any] | list[Any] | Sequence[str | BinaryContent | dict[str, Any] | list[Any]]:
+        self,
+        tool_name: str,
+        arguments: dict[str, Any],
+        metadata: dict[str, Any] | None = None,
+    ) -> ToolResult:
         """Call a tool on the server.
 
         Args:
             tool_name: The name of the tool to call.
             arguments: The arguments to pass to the tool.
+            metadata: Request-level metadata (optional)
 
         Returns:
             The result of the tool call.
@@ -129,7 +141,20 @@ async def call_tool(
             ModelRetry: If the tool call fails.
         """
         try:
-            result = await self._client.call_tool(self.get_unprefixed_tool_name(tool_name), arguments)
+            # meta param is not provided by session yet, so build and can send_request directly.
+            result = await self._client.send_request(
+                ClientRequest(
+                    CallToolRequest(
+                        method='tools/call',
+                        params=CallToolRequestParams(
+                            name=self.get_unprefixed_tool_name(tool_name),
+                            arguments=arguments,
+                            _meta=RequestParams.Meta(**metadata) if metadata else None,
+                        ),
+                    )
+                ),
+                CallToolResult,
+            )
         except McpError as e:
             raise ModelRetry(e.error.message)
 
@@ -269,6 +294,9 @@ async def main():
     e.g. if `tool_prefix='foo'`, then a tool named `bar` will be registered as `foo_bar`
     """
 
+    process_tool_call: ProcessToolCallback | None = None
+    """Hook to customize tool calling and optionally pass extra metadata."""
+
     timeout: float = 5
     """ The timeout in seconds to wait for the client to initialize."""
 
@@ -363,6 +391,9 @@ class _MCPServerHTTP(MCPServer):
     For example, if `tool_prefix='foo'`, then a tool named `bar` will be registered as `foo_bar`
     """
 
+    process_tool_call: ProcessToolCallback | None = None
+    """Hook to customize tool calling and optionally pass extra metadata."""
+
     @property
     @abstractmethod
     def _transport_client(
@@ -521,3 +552,29 @@ async def main():
     @property
     def _transport_client(self):
         return streamablehttp_client  # pragma: no cover
+
+
+ToolResult = (
+    str | BinaryContent | dict[str, Any] | list[Any] | Sequence[str | BinaryContent | dict[str, Any] | list[Any]]
+)
+"""The result type of a tool call."""
+
+CallToolFunc = Callable[[str, dict[str, Any], dict[str, Any] | None], Awaitable[ToolResult]]
+"""A function type that represents a tool call."""
+
+ProcessToolCallback = Callable[
+    [
+        RunContext[Any],
+        CallToolFunc,
+        str,
+        dict[str, Any],
+    ],
+    Awaitable[ToolResult],
+]
+"""A process tool callback.
+
+It accepts a run context, the original tool call function, a tool name, and arguments.
+
+Allows wrapping an MCP server tool call to customize it, including adding extra request
+metadata.
+"""

tests/mcp_server.py

@@ -3,6 +3,8 @@
 from typing import Any
 
 from mcp.server.fastmcp import Context, FastMCP, Image
+from mcp.server.session import ServerSessionT
+from mcp.shared.context import LifespanContextT, RequestT
 from mcp.types import BlobResourceContents, EmbeddedResource, TextResourceContents
 from pydantic import AnyUrl
 
@@ -118,6 +120,22 @@ async def get_log_level(ctx: Context) -> str:  # type: ignore
     return log_level
 
 
+@mcp.tool()
+async def echo_deps(ctx: Context[ServerSessionT, LifespanContextT, RequestT]) -> dict[str, Any]:
+    """Echo the run context.
+
+    Args:
+        ctx: Context object containing request and session information.
+
+    Returns:
+        Dictionary with an echo message and the deps.
+    """
+    await ctx.info('This is an info message')
+
+    deps: Any = getattr(ctx.request_context.meta, 'deps')
+    return {'echo': 'This is an echo message', 'deps': deps}
+
+
 @mcp._mcp_server.set_logging_level()  # pyright: ignore[reportPrivateUsage]
 async def set_logging_level(level: str) -> None:
     global log_level

tests/test_mcp.py

@@ -2,6 +2,7 @@
 
 import re
 from pathlib import Path
+from typing import Any, Final
 from unittest.mock import AsyncMock, patch
 
 import pytest
@@ -19,19 +20,22 @@
     ToolReturnPart,
     UserPromptPart,
 )
+from pydantic_ai.models.test import TestModel
+from pydantic_ai.tools import RunContext
 from pydantic_ai.usage import Usage
 
 from .conftest import IsDatetime, IsStr, try_import
 
 with try_import() as imports_successful:
     from mcp import ErrorData, McpError
 
-    from pydantic_ai.mcp import MCPServerSSE, MCPServerStdio
+    from pydantic_ai.mcp import CallToolFunc, MCPServerSSE, MCPServerStdio, ToolResult
     from pydantic_ai.models.google import GoogleModel
     from pydantic_ai.models.openai import OpenAIModel
     from pydantic_ai.providers.google import GoogleProvider
     from pydantic_ai.providers.openai import OpenAIProvider
 
+TOOL_COUNT: Final[int] = 12
 
 pytestmark = [
     pytest.mark.skipif(not imports_successful(), reason='mcp and openai not installed'),
@@ -51,7 +55,7 @@ async def test_stdio_server():
     server = MCPServerStdio('python', ['-m', 'tests.mcp_server'])
     async with server:
         tools = await server.list_tools()
-        assert len(tools) == 11
+        assert len(tools) == TOOL_COUNT
         assert tools[0].name == 'celsius_to_fahrenheit'
         assert tools[0].description.startswith('Convert Celsius to Fahrenheit.')
 
@@ -72,7 +76,29 @@ async def test_stdio_server_with_cwd():
     server = MCPServerStdio('python', ['mcp_server.py'], cwd=test_dir)
     async with server:
         tools = await server.list_tools()
-        assert len(tools) == snapshot(11)
+        assert len(tools) == TOOL_COUNT
+
+
+async def test_process_tool_call() -> None:
+    called: bool = False
+
+    async def process_tool_call(
+        ctx: RunContext[int],
+        call_tool: CallToolFunc,
+        tool_name: str,
+        args: dict[str, Any],
+    ) -> ToolResult:
+        """A process_tool_call that sets a flag and sends deps as metadata."""
+        nonlocal called
+        called = True
+        return await call_tool(tool_name, args, {'deps': ctx.deps})
+
+    server = MCPServerStdio('python', ['-m', 'tests.mcp_server'], process_tool_call=process_tool_call)
+    async with server:
+        agent = Agent(deps_type=int, model=TestModel(call_tools=['echo_deps']), mcp_servers=[server])
+        result = await agent.run('Echo with deps set to 42', deps=42)
+        assert result.output == snapshot('{"echo_deps":{"echo":"This is an echo message","deps":42}}')
+        assert called, 'process_tool_call should have been called'
 
 
 def test_sse_server():
@@ -217,7 +243,7 @@ async def test_log_level_unset():
     assert server._get_log_level() is None  # pyright: ignore[reportPrivateUsage]
     async with server:
         tools = await server.list_tools()
-        assert len(tools) == snapshot(11)
+        assert len(tools) == TOOL_COUNT
         assert tools[10].name == 'get_log_level'
 
         result = await server.call_tool('get_log_level', {})
@@ -945,7 +971,7 @@ async def test_mcp_server_raises_mcp_error(allow_model_requests: None, agent: Ag
     async with agent.run_mcp_servers():
         with patch.object(
             server._client,  # pyright: ignore[reportPrivateUsage]
-            'call_tool',
+            'send_request',
             new=AsyncMock(side_effect=mcp_error),
         ):
             with pytest.raises(ModelRetry, match='Test MCP error conversion'):