AbstractToolset.call_tool now takes a ToolCallPart

Author: DouweM

docs/mcp/client.md

@@ -189,10 +189,12 @@ async def process_tool_call(
     ctx: RunContext[int],
     call_tool: CallToolFunc,
     name: str,
-    tool_args: dict[str, Any],
+    tool_args: str | dict[str, Any] None,
+    *args: Any,
+    **kwargs: Any
 ) -> ToolResult:
     """A tool call processor that passes along the deps."""
-    return await call_tool(name, tool_args, metadata={'deps': ctx.deps})
+    return await call_tool(name, tool_args, *args, metadata={'deps': ctx.deps}, **kwargs)
 
 
 server = MCPServerStdio('python', ['mcp_server.py'], process_tool_call=process_tool_call)

pydantic_ai_slim/pydantic_ai/_agent_graph.py

@@ -618,7 +618,7 @@ async def process_function_tools(  # noqa: C901
             output_parts.append(part)
         else:
             try:
-                result_data = await _call_tool(toolset, call, run_context)
+                result_data = await toolset.call_tool(call, run_context)
             except exceptions.UnexpectedModelBehavior as e:
                 ctx.state.increment_retries(ctx.deps.max_result_retries, e)
                 raise e
@@ -755,7 +755,7 @@ async def _call_function_tool(
 
     with tracer.start_as_current_span('running tool', attributes=span_attributes) as span:
         try:
-            tool_result = await _call_tool(toolset, tool_call, run_context)
+            tool_result = await toolset.call_tool(tool_call, run_context)
         except ToolRetryError as e:
             part = e.tool_retry
             if include_content and span.is_recording():
@@ -827,14 +827,6 @@ def process_content(content: Any) -> Any:
         return (part, extra_parts)
 
 
-async def _call_tool(
-    toolset: AbstractToolset[DepsT], tool_call: _messages.ToolCallPart, run_context: RunContext[DepsT]
-) -> Any:
-    run_context = dataclasses.replace(run_context, tool_call_id=tool_call.tool_call_id)
-    args_dict = toolset.validate_tool_args(run_context, tool_call.tool_name, tool_call.args)
-    return await toolset.call_tool(run_context, tool_call.tool_name, args_dict)
-
-
 @dataclasses.dataclass
 class _RunMessages:
     messages: list[_messages.ModelMessage]

pydantic_ai_slim/pydantic_ai/_output.py

@@ -29,7 +29,7 @@
     _OutputSpecItem,  # type: ignore[reportPrivateUsage]
 )
 from .tools import GenerateToolJsonSchema, ObjectJsonSchema, ToolDefinition
-from .toolsets import AbstractToolset
+from .toolsets._callable import CallableToolset
 from .toolsets._run import RunToolset
 
 if TYPE_CHECKING:
@@ -829,7 +829,7 @@ async def process(
 
 
 @dataclass(init=False)
-class OutputToolset(AbstractToolset[AgentDepsT]):
+class OutputToolset(CallableToolset[AgentDepsT]):
     """A toolset that contains output tools."""
 
     _tool_defs: list[ToolDefinition]
@@ -930,9 +930,7 @@ def _get_tool_args_validator(self, ctx: RunContext[AgentDepsT], name: str) -> Sc
     def _max_retries_for_tool(self, name: str) -> int:
         return self.max_retries
 
-    async def call_tool(
-        self, ctx: RunContext[AgentDepsT], name: str, tool_args: dict[str, Any], *args: Any, **kwargs: Any
-    ) -> Any:
+    async def _call_tool(self, ctx: RunContext[AgentDepsT], name: str, tool_args: dict[str, Any]) -> Any:
         output = await self.processors[name].call(tool_args, ctx)
         for validator in self.output_validators:
             output = await validator.validate(output, ctx, wrap_validation_errors=False)

pydantic_ai_slim/pydantic_ai/agent.py

@@ -166,7 +166,7 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
     _prepare_output_tools: ToolsPrepareFunc[AgentDepsT] | None = dataclasses.field(repr=False)
     _max_result_retries: int = dataclasses.field(repr=False)
 
-    _running_count: int = dataclasses.field(repr=False)
+    _entered_count: int = dataclasses.field(repr=False)
     _exit_stack: AsyncExitStack | None = dataclasses.field(repr=False)
 
     @overload
@@ -430,7 +430,7 @@ def __init__(
         self._override_model: ContextVar[_utils.Option[models.Model]] = ContextVar('_override_model', default=None)
 
         self._exit_stack = None
-        self._running_count = 0
+        self._entered_count = 0
 
     @staticmethod
     def instrument_all(instrument: InstrumentationSettings | bool = True) -> None:
@@ -1788,18 +1788,17 @@ def is_end_node(
 
     async def __aenter__(self) -> Self:
         """Enter the agent. This will start all [`MCPServerStdio`s][pydantic_ai.mcp.MCPServerStdio] registered with the agent so they can be used in a run."""
-        if self._running_count == 0:
+        if self._entered_count == 0:
             self._exit_stack = AsyncExitStack()
             await self._exit_stack.enter_async_context(self._toolset)
-        self._running_count += 1
+        self._entered_count += 1
         return self
 
     async def __aexit__(self, *args: Any) -> bool | None:
-        self._running_count -= 1
-        if self._running_count <= 0 and self._exit_stack is not None:
+        self._entered_count -= 1
+        if self._entered_count <= 0 and self._exit_stack is not None:
             await self._exit_stack.aclose()
             self._exit_stack = None
-        return None
 
     def set_mcp_sampling_model(self, model: models.Model | models.KnownModelName | str | None = None) -> None:
         """Set the sampling model on all MCP servers registered with the agent.

pydantic_ai_slim/pydantic_ai/mcp.py

@@ -3,11 +3,10 @@
 import base64
 import functools
 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
-from types import TracebackType
 from typing import Any, Callable
 
 import anyio
@@ -19,10 +18,9 @@
 from pydantic_ai._run_context import RunContext
 from pydantic_ai.tools import ToolDefinition
 
-from .toolsets import AbstractToolset
+from .toolsets._callable import CallableToolset
 from .toolsets._run import RunToolset
 from .toolsets.prefixed import PrefixedToolset
-from .toolsets.processed import ProcessedToolset, ToolProcessFunc
 
 try:
     from mcp import types as mcp_types
@@ -45,7 +43,7 @@
 __all__ = 'MCPServer', 'MCPServerStdio', 'MCPServerHTTP', 'MCPServerSSE', 'MCPServerStreamableHTTP'
 
 
-class MCPServer(AbstractToolset[Any], ABC):
+class MCPServer(CallableToolset[Any], ABC):
     """Base class for attaching agents to MCP servers.
 
     See <https://modelcontextprotocol.io> for more information.
@@ -56,7 +54,7 @@ class MCPServer(AbstractToolset[Any], ABC):
     log_level: mcp_types.LoggingLevel | None = None
     log_handler: LoggingFnT | None = None
     timeout: float = 5
-    process_tool_call: ToolProcessFunc[Any] | None = None
+    process_tool_call: ProcessToolCallback | None = None
     allow_sampling: bool = True
     max_retries: int = 1
     sampling_model: models.Model | None = None
@@ -102,14 +100,11 @@ async def list_tools(self) -> list[mcp_types.Tool]:
             result = await self._client.list_tools()
         return result.tools
 
-    async def call_tool(
+    async def _call_tool(
         self,
         ctx: RunContext[Any],
         name: str,
         tool_args: dict[str, Any],
-        *args: Any,
-        metadata: dict[str, Any] | None = None,
-        **kwargs: Any,
     ) -> ToolResult:
         """Call a tool on the server.
 
@@ -127,36 +122,41 @@ async def call_tool(
         Raises:
             ModelRetry: If the tool call fails.
         """
-        async with self:  # Ensure server is running
-            try:
-                result = await self._client.send_request(
-                    mcp_types.ClientRequest(
-                        mcp_types.CallToolRequest(
-                            method='tools/call',
-                            params=mcp_types.CallToolRequestParams(
-                                name=name,
-                                arguments=tool_args,
-                                _meta=mcp_types.RequestParams.Meta(**metadata) if metadata else None,
-                            ),
-                        )
-                    ),
-                    mcp_types.CallToolResult,
-                )
-            except McpError as e:
-                raise exceptions.ModelRetry(e.error.message)
 
-        content = [self._map_tool_result_part(part) for part in result.content]
+        async def _call(name: str, args: dict[str, Any], metadata: dict[str, Any] | None = None) -> ToolResult:
+            async with self:  # Ensure server is running
+                try:
+                    result = await self._client.send_request(
+                        mcp_types.ClientRequest(
+                            mcp_types.CallToolRequest(
+                                method='tools/call',
+                                params=mcp_types.CallToolRequestParams(
+                                    name=name,
+                                    arguments=args,
+                                    _meta=mcp_types.RequestParams.Meta(**metadata) if metadata else None,
+                                ),
+                            )
+                        ),
+                        mcp_types.CallToolResult,
+                    )
+                except McpError as e:
+                    raise exceptions.ModelRetry(e.error.message)
+
+            content = [self._map_tool_result_part(part) for part in result.content]
+
+            if result.isError:
+                text = '\n'.join(str(part) for part in content)
+                raise exceptions.ModelRetry(text)
+            else:
+                return content[0] if len(content) == 1 else content
 
-        if result.isError:
-            text = '\n'.join(str(part) for part in content)
-            raise exceptions.ModelRetry(text)
+        if self.process_tool_call is not None:
+            return await self.process_tool_call(ctx, _call, name, tool_args)
         else:
-            return content[0] if len(content) == 1 else content
+            return await _call(name, tool_args)
 
     async def prepare_for_run(self, ctx: RunContext[Any]) -> RunToolset[Any]:
         frozen_toolset = RunToolset(self, ctx, await self.list_tool_defs())
-        if self.process_tool_call:
-            frozen_toolset = await ProcessedToolset(frozen_toolset, self.process_tool_call).prepare_for_run(ctx)
         if self.tool_prefix:
             frozen_toolset = await PrefixedToolset(frozen_toolset, self.tool_prefix).prepare_for_run(ctx)
         return RunToolset(frozen_toolset, ctx, original=self)
@@ -208,12 +208,7 @@ async def __aenter__(self) -> Self:
         self._running_count += 1
         return self
 
-    async def __aexit__(
-        self,
-        exc_type: type[BaseException] | None,
-        exc_value: BaseException | None,
-        traceback: TracebackType | None,
-    ) -> bool | None:
+    async def __aexit__(self, *args: Any) -> bool | None:
         self._running_count -= 1
         if self._running_count <= 0:
             await self._exit_stack.aclose()
@@ -364,7 +359,7 @@ async def main():
     timeout: float = 5
     """The timeout in seconds to wait for the client to initialize."""
 
-    process_tool_call: ToolProcessFunc[Any] | None = None
+    process_tool_call: ProcessToolCallback | None = None
     """Hook to customize tool calling and optionally pass extra metadata."""
 
     allow_sampling: bool = True
@@ -465,7 +460,7 @@ class _MCPServerHTTP(MCPServer):
     If the connection cannot be established within this time, the operation will fail.
     """
 
-    process_tool_call: ToolProcessFunc[Any] | None = None
+    process_tool_call: ProcessToolCallback | None = None
     """Hook to customize tool calling and optionally pass extra metadata."""
 
     allow_sampling: bool = True
@@ -642,3 +637,23 @@ def _transport_client(self):
     | Sequence[str | messages.BinaryContent | dict[str, Any] | list[Any]]
 )
 """The result type of an MCP 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.
+"""

pydantic_ai_slim/pydantic_ai/result.py

@@ -3,7 +3,7 @@
 import warnings
 from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable
 from copy import copy
-from dataclasses import dataclass, field, replace
+from dataclasses import dataclass, field
 from datetime import datetime
 from typing import Generic, cast
 
@@ -106,11 +106,7 @@ async def _validate_response(
                 raise exceptions.UnexpectedModelBehavior(  # pragma: no cover
                     f'Invalid response, unable to find tool call for {output_tool_name!r}'
                 )
-            run_context = replace(self._run_ctx, tool_call_id=tool_call.tool_call_id)
-            args_dict = self._toolset.validate_tool_args(
-                run_context, tool_call.tool_name, tool_call.args, allow_partial=allow_partial
-            )
-            return await self._toolset.call_tool(run_context, tool_call.tool_name, args_dict)
+            return await self._toolset.call_tool(tool_call, self._run_ctx, allow_partial=allow_partial)
         elif deferred_tool_calls := self._toolset.get_deferred_tool_calls(message.parts):
             if not self._output_schema.allows_deferred_tool_calls:
                 raise exceptions.UserError(  # pragma: no cover
@@ -442,11 +438,7 @@ async def validate_structured_output(
                 raise exceptions.UnexpectedModelBehavior(  # pragma: no cover
                     f'Invalid response, unable to find tool call for {self._output_tool_name!r}'
                 )
-            run_context = replace(self._run_ctx, tool_call_id=tool_call.tool_call_id)
-            args_dict = self._toolset.validate_tool_args(
-                run_context, tool_call.tool_name, tool_call.args, allow_partial=allow_partial
-            )
-            return await self._toolset.call_tool(run_context, tool_call.tool_name, args_dict)
+            return await self._toolset.call_tool(tool_call, self._run_ctx, allow_partial=allow_partial)
         elif deferred_tool_calls := self._toolset.get_deferred_tool_calls(message.parts):
             if not self._output_schema.allows_deferred_tool_calls:
                 raise exceptions.UserError(

pydantic_ai_slim/pydantic_ai/toolsets/__init__.py

@@ -1,13 +1,12 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from types import TracebackType
-from typing import TYPE_CHECKING, Any, Callable, Generic, Literal
+from typing import TYPE_CHECKING, Any, Callable, Generic
 
-from pydantic_core import SchemaValidator
 from typing_extensions import Self
 
 from .._run_context import AgentDepsT, RunContext
+from ..messages import ToolCallPart
 from ..tools import ToolDefinition
 
 if TYPE_CHECKING:
@@ -34,9 +33,7 @@ def tool_name_conflict_hint(self) -> str:
     async def __aenter__(self) -> Self:
         return self
 
-    async def __aexit__(
-        self, exc_type: type[BaseException] | None, exc_value: BaseException | None, traceback: TracebackType | None
-    ) -> bool | None:
+    async def __aexit__(self, *args: Any) -> bool | None:
         return None
 
     @abstractmethod
@@ -55,28 +52,12 @@ def tool_names(self) -> list[str]:
     def get_tool_def(self, name: str) -> ToolDefinition | None:
         return next((tool_def for tool_def in self.tool_defs if tool_def.name == name), None)
 
-    @abstractmethod
-    def _get_tool_args_validator(self, ctx: RunContext[AgentDepsT], name: str) -> SchemaValidator:
-        raise NotImplementedError()
-
-    def validate_tool_args(
-        self, ctx: RunContext[AgentDepsT], name: str, args: str | dict[str, Any] | None, allow_partial: bool = False
-    ) -> dict[str, Any]:
-        pyd_allow_partial: Literal['off', 'trailing-strings'] = 'trailing-strings' if allow_partial else 'off'
-        validator = self._get_tool_args_validator(ctx, name)
-        if isinstance(args, str):
-            return validator.validate_json(args or '{}', allow_partial=pyd_allow_partial)
-        else:
-            return validator.validate_python(args or {}, allow_partial=pyd_allow_partial)
-
     @abstractmethod
     def _max_retries_for_tool(self, name: str) -> int:
         raise NotImplementedError()
 
     @abstractmethod
-    async def call_tool(
-        self, ctx: RunContext[AgentDepsT], name: str, tool_args: dict[str, Any], *args: Any, **kwargs: Any
-    ) -> Any:
+    async def call_tool(self, call: ToolCallPart, ctx: RunContext[AgentDepsT], allow_partial: bool = False) -> Any:
         raise NotImplementedError()
 
     def accept(self, visitor: Callable[[AbstractToolset[AgentDepsT]], Any]) -> Any: