Make MCPServer a Toolset

pydantic · DouweM · Jul 16, 2025 · Jun 3, 2025 · Jun 3, 2025 · Jun 3, 2025
commit 0f8da7450fb6ca2103e7075aa967ce2ecd154955
diff --git a/pydantic_ai_slim/pydantic_ai/agent.py b/pydantic_ai_slim/pydantic_ai/agent.py
@@ -235,7 +235,9 @@ def __init__(
         output_retries: int | None = None,
         tools: Sequence[Tool[AgentDepsT] | ToolFuncEither[AgentDepsT, ...]] = (),
         prepare_tools: ToolsPrepareFunc[AgentDepsT] | None = None,
-        mcp_servers: Sequence[MCPServer] = (),
+        mcp_servers: Sequence[
+            MCPServer
+        ] = (),  # TODO: Deprecate argument, MCPServers can be passed directly to toolsets
         toolsets: Sequence[AbstractToolset[AgentDepsT]] = (),
         defer_model_check: bool = False,
         end_strategy: EndStrategy = 'early',
@@ -364,11 +366,7 @@ def __init__(
         # This will raise errors for any name conflicts
         CombinedToolset[AgentDepsT]([self._output_toolset, self._function_toolset])
 
-        mcp_toolsets = [
-            cast(AbstractToolset[AgentDepsT], mcp_server.as_toolset(max_retries=self._max_result_retries))
-            for mcp_server in mcp_servers
-        ]
-        toolset = CombinedToolset[AgentDepsT]([self._function_toolset, *toolsets, *mcp_toolsets])
+        toolset = CombinedToolset[AgentDepsT]([self._function_toolset, *toolsets, *mcp_servers])
         if prepare_tools:
             toolset = PreparedToolset[AgentDepsT](toolset, prepare_tools)
         self._toolset = toolset

diff --git a/pydantic_ai_slim/pydantic_ai/mcp.py b/pydantic_ai_slim/pydantic_ai/mcp.py
@@ -16,8 +16,11 @@
 from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
 from typing_extensions import Self, assert_never, deprecated
 
+from pydantic_ai._run_context import RunContext
+from pydantic_ai.tools import ToolDefinition
+
 from .exceptions import UserError
-from .toolset import AbstractToolset, MCPToolset, PrefixedToolset, ProcessedToolset, ToolProcessFunc
+from .toolset import AbstractToolset, FrozenToolset, PrefixedToolset, ProcessedToolset, ToolProcessFunc
 
 try:
     from mcp import types as mcp_types
@@ -40,7 +43,7 @@
 __all__ = 'MCPServer', 'MCPServerStdio', 'MCPServerHTTP', 'MCPServerSSE', 'MCPServerStreamableHTTP'
 
 
-class MCPServer(ABC):
+class MCPServer(AbstractToolset[Any], ABC):
     """Base class for attaching agents to MCP servers.
 
     See <https://modelcontextprotocol.io> for more information.
@@ -51,7 +54,7 @@ class MCPServer(ABC):
     log_level: mcp_types.LoggingLevel | None = None
     log_handler: LoggingFnT | None = None
     init_timeout: float = 5
-    process_tool_call: ToolProcessFunc | None = None
+    process_tool_call: ToolProcessFunc[Any] | None = None
     # } end of "abstract fields"
 
     _running_count: int = 0
@@ -81,6 +84,14 @@ def is_running(self) -> bool:
         """Check if the MCP server is running."""
         return bool(self._running_count)
 
+    @property
+    def name(self) -> str:
+        return repr(self)
+
+    @property
+    def name_conflict_hint(self) -> str:
+        return 'Consider setting `tool_prefix` to avoid name conflicts.'
+
     async def list_tools(self) -> list[mcp_types.Tool]:
         """Retrieve tools that are currently active on the server.
 
@@ -95,16 +106,22 @@ async def list_tools(self) -> list[mcp_types.Tool]:
 
     async def call_tool(
         self,
-        tool_name: str,
-        arguments: dict[str, Any],
+        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.
 
         Args:
-            tool_name: The name of the tool to call.
-            arguments: The arguments to pass to the tool.
+            ctx: The run context of the tool call.
+            name: The name of the tool to call.
+            tool_args: The arguments to pass to the tool.
+            *args: Additional arguments passed by a tool call processor.
             metadata: Request-level metadata (optional)
+            **kwargs: Additional keyword arguments passed by a tool call processor.
 
         Returns:
             The result of the tool call.
@@ -121,8 +138,8 @@ async def call_tool(
                     mcp_types.CallToolRequest(
                         method='tools/call',
                         params=mcp_types.CallToolRequestParams(
-                            name=tool_name,
-                            arguments=arguments,
+                            name=name,
+                            arguments=tool_args,
                             _meta=mcp_types.RequestParams.Meta(**metadata) if metadata else None,
                         ),
                     )
@@ -140,13 +157,42 @@ async def call_tool(
         else:
             return content[0] if len(content) == 1 else content
 
-    def as_toolset(self, max_retries: int = 1) -> AbstractToolset:
-        toolset = MCPToolset(self, max_retries=max_retries)
+    async def freeze_for_run(self, ctx: RunContext[Any]) -> FrozenToolset[Any]:
+        frozen_self = FrozenToolset[Any](self, ctx, await self.list_tool_defs())
+        toolset = frozen_self
         if self.process_tool_call:
             toolset = ProcessedToolset(toolset, self.process_tool_call)
         if self.tool_prefix:
             toolset = PrefixedToolset(toolset, self.tool_prefix)
-        return toolset
+        return FrozenToolset[Any](toolset, ctx, original=self)
+
+    @property
+    def tool_defs(self) -> list[ToolDefinition]:
+        return []
+
+    async def list_tool_defs(self) -> list[ToolDefinition]:
+        mcp_tools = await self.list_tools()
+        return [
+            ToolDefinition(
+                name=mcp_tool.name,
+                description=mcp_tool.description or '',
+                parameters_json_schema=mcp_tool.inputSchema,
+            )
+            for mcp_tool in mcp_tools
+        ]
+
+    def get_tool_args_validator(self, ctx: RunContext[Any], name: str) -> pydantic_core.SchemaValidator:
+        return pydantic_core.SchemaValidator(
+            schema=pydantic_core.core_schema.dict_schema(
+                pydantic_core.core_schema.str_schema(), pydantic_core.core_schema.any_schema()
+            )
+        )
+
+    def max_retries_for_tool(self, name: str) -> int:
+        return 1
+
+    def set_mcp_sampling_model(self, model: models.Model) -> None:
+        self.sampling_model = model
 
     async def __aenter__(self) -> Self:
         if self._running_count == 0:
@@ -323,7 +369,7 @@ async def main():
     timeout: float = 5
     """ The timeout in seconds to wait for the client to initialize."""
 
-    process_tool_call: ToolProcessFunc | None = None
+    process_tool_call: ToolProcessFunc[Any] | None = None
     """Hook to customize tool calling and optionally pass extra metadata."""
 
     @asynccontextmanager
@@ -418,7 +464,7 @@ class _MCPServerHTTP(MCPServer):
     init_timeout: float = 5
     """The timeout in seconds to wait for the client to initialize."""
 
-    process_tool_call: ToolProcessFunc | None = None
+    process_tool_call: ToolProcessFunc[Any] | None = None
     """Hook to customize tool calling and optionally pass extra metadata."""
 
     @property

diff --git a/pydantic_ai_slim/pydantic_ai/toolset.py b/pydantic_ai_slim/pydantic_ai/toolset.py
@@ -11,7 +11,7 @@
 
 from pydantic import ValidationError
 from pydantic.json_schema import GenerateJsonSchema
-from pydantic_core import SchemaValidator, core_schema
+from pydantic_core import SchemaValidator
 from typing_extensions import Self
 
 from ._output import BaseOutputSchema
@@ -29,7 +29,6 @@
 )
 
 if TYPE_CHECKING:
-    from pydantic_ai.mcp import MCPServer
     from pydantic_ai.models import Model
 
 
@@ -98,7 +97,7 @@ def set_mcp_sampling_model(self, model: Model) -> None:
 
 
 @dataclass(init=False)
-class _BareFunctionToolset(AbstractToolset[AgentDepsT]):
+class FunctionToolset(AbstractToolset[AgentDepsT]):
     """A toolset that functions can be registered to as tools."""
 
     max_retries: int = field(default=1)
@@ -259,6 +258,16 @@ def register_tool(self, tool: Tool[AgentDepsT]) -> None:
             tool.max_retries = self.max_retries
         self.tools[tool.name] = tool
 
+    async def freeze_for_run(self, ctx: RunContext[AgentDepsT]) -> FrozenToolset[AgentDepsT]:
+        frozen_self = FrozenToolset[AgentDepsT](self, ctx)
+        toolset = await IndividuallyPreparedToolset(frozen_self, self._prepare_tool_def).freeze_for_run(ctx)
+        return FrozenToolset[AgentDepsT](toolset, ctx, original=self)
+
+    async def _prepare_tool_def(self, ctx: RunContext[AgentDepsT], tool_def: ToolDefinition) -> ToolDefinition | None:
+        tool_name = tool_def.name
+        ctx = replace(ctx, tool_name=tool_name, retry=ctx.retries.get(tool_name, 0))
+        return await self.tools[tool_name].prepare_tool_def(ctx)
+
     @property
     def tool_defs(self) -> list[ToolDefinition]:
         return [tool.tool_def for tool in self.tools.values()]
@@ -301,63 +310,6 @@ async def call_tool(
         return await self.output_schema.tools[name].processor.process(tool_args, ctx)
 
 
-@dataclass
-class MCPToolset(AbstractToolset[AgentDepsT]):
-    """A toolset that contains MCP tools and handles running the server."""
-
-    server: MCPServer
-    max_retries: int = field(default=1)
-
-    @property
-    def name(self) -> str:
-        return repr(self.server)
-
-    @property
-    def name_conflict_hint(self) -> str:
-        return 'Consider setting `tool_prefix` to avoid name conflicts.'
-
-    async def __aenter__(self) -> Self:
-        await self.server.__aenter__()
-        return self
-
-    async def __aexit__(
-        self, exc_type: type[BaseException] | None, exc_value: BaseException | None, traceback: TracebackType | None
-    ) -> bool | None:
-        return await self.server.__aexit__(exc_type, exc_value, traceback)
-
-    async def freeze_for_run(self, ctx: RunContext[AgentDepsT]) -> FrozenToolset[AgentDepsT]:
-        return FrozenToolset[AgentDepsT](self, ctx, await self.list_tool_defs())
-
-    @property
-    def tool_defs(self) -> list[ToolDefinition]:
-        return []
-
-    async def list_tool_defs(self) -> list[ToolDefinition]:
-        mcp_tools = await self.server.list_tools()
-        return [
-            ToolDefinition(
-                name=mcp_tool.name,
-                description=mcp_tool.description or '',
-                parameters_json_schema=mcp_tool.inputSchema,
-            )
-            for mcp_tool in mcp_tools
-        ]
-
-    def get_tool_args_validator(self, ctx: RunContext[AgentDepsT], name: str) -> SchemaValidator:
-        return SchemaValidator(schema=core_schema.dict_schema(core_schema.str_schema(), core_schema.any_schema()))
-
-    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], metadata: dict[str, Any] | None = None
-    ) -> Any:
-        return await self.server.call_tool(name, tool_args, metadata)
-
-    def set_mcp_sampling_model(self, model: Model) -> None:
-        self.server.sampling_model = model
-
-
 @dataclass
 class WrapperToolset(AbstractToolset[AgentDepsT]):
     """A toolset that wraps another toolset and delegates to it."""
@@ -569,21 +521,6 @@ async def freeze_for_run(self, ctx: RunContext[AgentDepsT]) -> FrozenToolset[Age
         return FrozenToolset[AgentDepsT](freezable_prepared, ctx)
 
 
-@dataclass
-class FunctionToolset(IndividuallyPreparedToolset[AgentDepsT]):
-    """A toolset that contains function tools."""
-
-    def __init__(self, tools: Sequence[Tool[AgentDepsT] | ToolFuncEither[AgentDepsT, ...]] = [], max_retries: int = 1):
-        wrapped = _BareFunctionToolset(tools, max_retries)
-
-        async def prepare_tool_def(ctx: RunContex
F42D
t[AgentDepsT], tool_def: ToolDefinition) -> ToolDefinition | None:
-            tool_name = tool_def.name
-            ctx = replace(ctx, tool_name=tool_name, retry=ctx.retries.get(tool_name, 0))
-            return await wrapped.tools[tool_name].prepare_tool_def(ctx)
-
-        super().__init__(wrapped, prepare_tool_def)
-
-
 @dataclass(init=False)
 class _FreezableIndividuallyPreparedToolset(IndividuallyPreparedToolset[AgentDepsT]):
     name_map: dict[str, str]
@@ -676,18 +613,21 @@ class FrozenToolset(WrapperToolset[AgentDepsT]):
     _tool_defs: list[ToolDefinition]
     _tool_names: list[str]
     _retries: dict[str, int]
+    _original: AbstractToolset[AgentDepsT]
 
     def __init__(
         self,
         wrapped: AbstractToolset[AgentDepsT],
         ctx: RunContext[AgentDepsT],
         tool_defs: list[ToolDefinition] | None = None,
+        original: AbstractToolset[AgentDepsT] | None = None,
     ):
         self.wrapped = wrapped
         self.ctx = ctx
         self._tool_defs = wrapped.tool_defs if tool_defs is None else tool_defs
         self._tool_names = [tool_def.name for tool_def in self._tool_defs]
         self._retries = ctx.retries.copy()
+        self._original = original or wrapped
 
     @property
     def name(self) -> str:
@@ -698,7 +638,7 @@ async def freeze_for_run(self, ctx: RunContext[AgentDepsT]) -> FrozenToolset[Age
             return self
         else:
             ctx = replace(ctx, retries=self._retries)
-            return await self.wrapped.freeze_for_run(ctx)
+            return await self._original.freeze_for_run(ctx)
 
     @property
     def tool_defs(self) -> list[ToolDefinition]:

diff --git a/tests/test_examples.py b/tests/test_examples.py
@@ -16,11 +16,13 @@
 import pytest
 from _pytest.mark import ParameterSet
 from devtools import debug
+from pydantic_core import SchemaValidator, core_schema
 from pytest_examples import CodeExample, EvalExample, find_examples
 from pytest_mock import MockerFixture
 from rich.console import Console
 
 from pydantic_ai import ModelHTTPError
+from pydantic_ai._run_context import RunContext
 from pydantic_ai._utils import group_by_temporal
 from pydantic_ai.exceptions import UnexpectedModelBehavior
 from pydantic_ai.messages import (
@@ -36,7 +38,8 @@
 from pydantic_ai.models.fallback import FallbackModel
 from pydantic_ai.models.function import AgentInfo, DeltaToolCall, DeltaToolCalls, FunctionModel
 from pydantic_ai.models.test import TestModel
-from pydantic_ai.toolset import AbstractToolset, MCPToolset
+from pydantic_ai.tools import ToolDefinition
+from pydantic_ai.toolset import AbstractToolset
 
 from .conftest import ClientWithHandler, TestEnv, try_import
 
@@ -253,7 +256,7 @@ def rich_prompt_ask(prompt: str, *_args: Any, **_kwargs: Any) -> str:
            raise ValueError(f'Unexpected prompt: {prompt}')
 
 
-class MockMCPServer:
+class MockMCPServer(AbstractToolset[Any]):
     is_running = True
 
     async def __aenter__(self) -> MockMCPServer:
@@ -262,14 +265,20 @@ async def __aenter__(self) -> MockMCPServer:
     async def __aexit__(self, *args: Any) -> None:
         pass
 
-    async def list_tools(self) -> list[None]:
+    @property
+    def tool_defs(self) -> list[ToolDefinition]:
         return []
 
-    async def call_tool(self, name: str, args: dict[str, Any], metadata: dict[str, Any] | None = None) -> Any:
-        return None
+    def get_tool_args_validator(self, ctx: RunContext[Any], name: str) -> SchemaValidator:
+        return SchemaValidator(core_schema.any_schema())
+
+    def max_retries_for_tool(self, name: str) -> int:
+        return 0
 
-    def as_toolset(self, max_retries: int = 1) -> AbstractToolset:
-        return MCPToolset(self, max_retries=max_retries)  # type: ignore
+    async def call_tool(
+        self, ctx: RunContext[Any], name: str, tool_args: dict[str, Any], *args: Any, **kwargs: Any
+    ) -> Any:
+        return None
 
 
 text_responses: dict[str, str | ToolCallPart] = {