Add support for MCP servers (#1100)

Co-authored-by: David Montague <35119617+dmontagu@users.noreply.github.com>
Co-authored-by: Samuel Colvin <s@muelcolvin.com>
Co-authored-by: Victorien <65306057+Viicos@users.noreply.github.com>

Author: 4 people

.github/workflows/ci.yml

@@ -157,6 +157,7 @@ jobs:
 
       # this must run last as it modifies the environment!
       - name: test lowest versions
+        if: matrix.python-version != '3.9'
         run: |
           unset UV_FROZEN
           uv run --all-extras --resolution lowest-direct coverage run -m pytest

docs/api/mcp.md

@@ -0,0 +1 @@
+::: pydantic_ai.mcp

docs/mcp_servers.md

@@ -0,0 +1,122 @@
+# MCP Servers
+
+**PydanticAI** supports integration with
+[MCP (Model Control Protocol) Servers](https://modelcontextprotocol.io/introduction),
+allowing you to extend agent capabilities through external services. This integration enables
+dynamic tool discovery.
+
+## Install
+
+To use MCP servers, you need to either install [`pydantic-ai`](install.md), or install
+[`pydantic-ai-slim`](install.md#slim-install) with the `mcp` optional group:
+
+```bash
+pip/uv-add 'pydantic-ai-slim[mcp]'
+```
+
+!!! note
+    MCP integration requires Python 3.10 or higher.
+
+## Running MCP Servers
+
+Before diving into how to use MCP servers with PydanticAI, let's look at how to run MCP servers
+with different transports.
+
+To run MCP servers, you'll need to install the MCP CLI package:
+
+```bash
+pip/uv-add 'mcp[cli]'
+```
+
+Here's a simple MCP server that provides a temperature conversion tool. We will later assume this is the server we connect to from our agent:
+
+```python {title="temperature_mcp_server.py" py="3.10"}
+from mcp.server.fastmcp import FastMCP
+
+mcp = FastMCP('Temperature Conversion Server')
+
+
+@mcp.tool()
+async def celsius_to_fahrenheit(celsius: float) -> float:
+    """Convert Celsius to Fahrenheit.
+
+    Args:
+        celsius: Temperature in Celsius
+    """
+    return (celsius * 9 / 5) + 32
+
+
+if __name__ == '__main__':
+    mcp.run('stdio')  # (1)!
+```
+
+1. Run with stdio transport (for subprocess communication).
+
+The same server can be run with [SSE transport](https://modelcontextprotocol.io/docs/concepts/transports#server-sent-events-sse)
+for HTTP communication:
+
+```python {title="temperature_mcp_server_sse.py" py="3.10"}
+from temperature_mcp_server import mcp
+
+if __name__ == '__main__':
+    mcp.run('sse', port=8000)  # (1)!
+```
+
+1. Run with SSE transport on port 8000.
+
+## Usage
+
+PydanticAI comes with two ways to connect to MCP servers:
+
+- [`MCPServerSSE`][pydantic_ai.mcp.MCPServerSSE] which connects to an MCP server using the [HTTP SSE](https://modelcontextprotocol.io/docs/concepts/transports#server-sent-events-sse) transport
+- [`MCPServerStdio`][pydantic_ai.mcp.MCPServerStdio] which runs the server as a subprocess and connects to it using the [stdio](https://modelcontextprotocol.io/docs/concepts/transports#standard-input%2Foutput-stdio) transport
+
+Examples of both are shown below.
+
+### MCP Remote Server
+
+You can have a MCP server running on a remote server. In this case, you'd use the
+[`MCPServerSSE`][pydantic_ai.mcp.MCPServerSSE] class:
+
+```python {title="mcp_remote_server.py" py="3.10"}
+from pydantic_ai import Agent
+from pydantic_ai.mcp import MCPServerSSE
+
+server = MCPServerSSE(url='http://localhost:8005/sse')
+agent = Agent('openai:gpt-4o', mcp_servers=[server])
+
+
+async def main():
+    async with agent.run_mcp_servers():
+        result = await agent.run('Can you convert 30 degrees celsius to fahrenheit?')
+    print(result.data)
+    #> 30 degrees Celsius is equal to 86 degrees Fahrenheit.
+```
+
+This will connect to the MCP server at the given URL and use the
+[SSE transport](https://modelcontextprotocol.io/docs/concepts/transports#server-sent-events-sse).
+
+### MCP Subprocess Server
+
+We also have a subprocess-based server that can be used to run the MCP server in a separate process.
+In this case, you'd use the [`MCPServerStdio`][pydantic_ai.mcp.MCPServerStdio] class,
+when using `MCPServerStdio` you need to run the server with the [`run_mcp_servers`][pydantic_ai.Agent.run_mcp_servers]
+context manager before running the server.
+
+```python {title="mcp_subprocess_server.py" py="3.10"}
+from pydantic_ai.agent import Agent
+from pydantic_ai.mcp import MCPServerStdio
+
+server = MCPServerStdio('python', ['-m', 'pydantic_ai_examples.mcp_server'])
+agent = Agent('openai:gpt-4o', mcp_servers=[server])
+
+
+async def main():
+    async with agent.run_mcp_servers():
+        result = await agent.run('Can you convert 30 degrees celsius to fahrenheit?')
+    print(result.data)
+    #> 30 degrees Celsius is equal to 86 degrees Fahrenheit.
+```
+
+This will start the MCP server in a separate process and connect to it using the
+[stdio transport](https://modelcontextprotocol.io/docs/concepts/transports#standard-input%2Foutput-stdio).

examples/pydantic_ai_examples/mcp_server.py

@@ -0,0 +1,37 @@
+"""Simple MCP Server that can be used to test the MCP protocol.
+
+Run with:
+
+    uv run -m pydantic_ai_examples.mcp_server --transport <TRANSPORT>
+
+TRANSPORT can be either `sse` or `stdio`.
+"""
+
+import argparse
+
+from mcp.server.fastmcp import FastMCP
+
+mcp = FastMCP('PydanticAI MCP Server', port=8005)
+
+
+@mcp.tool()
+async def celsius_to_fahrenheit(celsius: float) -> float:
+    """Convert Celsius to Fahrenheit.
+
+    Args:
+        celsius: Temperature in Celsius
+
+    Returns:
+        Temperature in Fahrenheit
+    """
+    return (celsius * 9 / 5) + 32
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        '--transport', type=str, default='stdio', choices=('sse', 'stdio')
+    )
+    args = parser.parse_args()
+
+    mcp.run(transport=args.transport)

examples/pyproject.toml

@@ -41,6 +41,7 @@ dependencies = [
     "uvicorn>=0.32.0",
     "devtools>=0.12.2",
     "gradio>=5.9.0; python_version>'3.9'",
+    "mcp[cli]>=1.4.1; python_version >= '3.10'"
 ]
 
 [tool.hatch.build.targets.wheel]

mkdocs.yml

@@ -28,6 +28,7 @@ nav:
       - multi-agent-applications.md
       - graph.md
       - input.md
+      - mcp_servers.md
       - cli.md
   - Examples:
       - examples/index.md
@@ -64,6 +65,7 @@ nav:
       - api/models/function.md
       - api/models/fallback.md
       - api/providers.md
+      - api/mcp.md
       - api/pydantic_graph/graph.md
       - api/pydantic_graph/nodes.md
       - api/pydantic_graph/persistence.md

pydantic_ai_slim/pydantic_ai/_agent_graph.py

@@ -7,7 +7,7 @@
 from contextlib import asynccontextmanager, contextmanager
 from contextvars import ContextVar
 from dataclasses import field
-from typing import Any, Generic, Literal, Union, cast
+from typing import TYPE_CHECKING, Any, Generic, Literal, Union, cast
 
 from opentelemetry.trace import Span, Tracer
 from typing_extensions import TypeGuard, TypeVar, assert_never
@@ -27,11 +27,10 @@
 from .models.instrumented import InstrumentedModel
 from .result import ResultDataT
 from .settings import ModelSettings, merge_model_settings
-from .tools import (
-    RunContext,
-    Tool,
-    ToolDefinition,
-)
+from .tools import RunContext, Tool, ToolDefinition
+
+if TYPE_CHECKING:
+    from .mcp import MCPServer
 
 __all__ = (
     'GraphAgentState',
@@ -94,6 +93,7 @@ class GraphAgentDeps(Generic[DepsT, ResultDataT]):
     result_validators: list[_result.ResultValidator[DepsT, ResultDataT]]
 
     function_tools: dict[str, Tool[DepsT]] = dataclasses.field(repr=False)
+    mcp_servers: Sequence[MCPServer] = dataclasses.field(repr=False)
 
     run_span: Span
     tracer: Tracer
@@ -219,7 +219,17 @@ async def add_tool(tool: Tool[DepsT]) -> None:
         if tool_def := await tool.prepare_tool_def(ctx):
             function_tool_defs.append(tool_def)
 
-    await asyncio.gather(*map(add_tool, ctx.deps.function_tools.values()))
+    async def add_mcp_server_tools(server: MCPServer) -> None:
+        if not server.is_running:
+            raise exceptions.UserError(f'MCP server is not running: {server}')
+        tool_defs = await server.list_tools()
+        # TODO(Marcelo): We should check if the tool names are unique. If not, we should raise an error.
+        function_tool_defs.extend(tool_defs)
+
+    await asyncio.gather(
+        *map(add_tool, ctx.deps.function_tools.values()),
+        *map(add_mcp_server_tools, ctx.deps.mcp_servers),
+    )
 
     result_schema = ctx.deps.result_schema
     return models.ModelRequestParameters(
@@ -594,6 +604,21 @@ async def process_function_tools(
                 yield event
                 call_index_to_event_id[len(calls_to_run)] = event.call_id
                 calls_to_run.append((tool, call))
+        elif mcp_tool := await _tool_from_mcp_server(call.tool_name, ctx):
+            if stub_function_tools:
+                # TODO(Marcelo): We should add coverage for this part of the code.
+                output_parts.append(  # pragma: no cover
+                    _messages.ToolReturnPart(
+                        tool_name=call.tool_name,
+                        content='Tool not executed - a final result was already processed.',
+                        tool_call_id=call.tool_call_id,
+                    )
+                )
+            else:
+                event = _messages.FunctionToolCallEvent(call)
+                yield event
+                call_index_to_event_id[len(calls_to_run)] = event.call_id
+                calls_to_run.append((mcp_tool, call))
         elif result_schema is not None and call.tool_name in result_schema.tools:
             # if tool_name is in _result_schema, it means we found a result tool but an error occurred in
             # validation, we don't add another part here
@@ -641,6 +666,35 @@ async def process_function_tools(
         output_parts.append(results_by_index[k])
 
 
+async def _tool_from_mcp_server(
+    tool_name: str,
+    ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, NodeRunEndT]],
+) -> Tool[DepsT] | None:
+    """Call each MCP server to find the tool with the given name.
+
+    Args:
+        tool_name: The name of the tool to find.
+        ctx: The current run context.
+
+    Returns:
+        The tool with the given name, or `None` if no tool with the given name is found.
+    """
+
+    async def run_tool(ctx: RunContext[DepsT], **args: Any) -> Any:
+        # There's no normal situation where the server will not be running at this point, we check just in case
+        # 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)
+        return result
+
+    for server in ctx.deps.mcp_servers:
+        tools = await server.list_tools()
+        if tool_name in {tool.name for tool in tools}:
+            return Tool(name=tool_name, function=run_tool, takes_ctx=True)
+    return None
+
+
 def _unknown_tool(
     tool_name: str,
     ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, NodeRunEndT]],

pydantic_ai_slim/pydantic_ai/agent.py

@@ -3,10 +3,10 @@
 import dataclasses
 import inspect
 from collections.abc import AsyncIterator, Awaitable, Iterator, Sequence
-from contextlib import AbstractAsyncContextManager, asynccontextmanager, contextmanager
+from contextlib import AbstractAsyncContextManager, AsyncExitStack, asynccontextmanager, contextmanager
 from copy import deepcopy
 from types import FrameType
-from typing import Any, Callable, ClassVar, Generic, cast, final, overload
+from typing import TYPE_CHECKING, Any, Callable, ClassVar, Generic, cast, final, overload
 
 from opentelemetry.trace import NoOpTracer, use_span
 from typing_extensions import TypeGuard, TypeVar, deprecated
@@ -47,6 +47,9 @@
 ModelRequestNode = _agent_graph.ModelRequestNode
 UserPromptNode = _agent_graph.UserPromptNode
 
+if TYPE_CHECKING:
+    from pydantic_ai.mcp import MCPServer
+
 __all__ = (
     'Agent',
     'AgentRun',
@@ -129,6 +132,7 @@ class Agent(Generic[AgentDepsT, ResultDataT]):
         repr=False
     )
     _function_tools: dict[str, Tool[AgentDepsT]] = dataclasses.field(repr=False)
+    _mcp_servers: Sequence[MCPServer] = dataclasses.field(repr=False)
     _default_retries: int = dataclasses.field(repr=False)
     _max_result_retries: int = dataclasses.field(repr=False)
     _override_deps: _utils.Option[AgentDepsT] = dataclasses.field(default=None, repr=False)
@@ -148,6 +152,7 @@ def __init__(
         result_tool_description: str | None = None,
         result_retries: int | None = None,
         tools: Sequence[Tool[AgentDepsT] | ToolFuncEither[AgentDepsT, ...]] = (),
+        mcp_servers: Sequence[MCPServer] = (),
         defer_model_check: bool = False,
         end_strategy: EndStrategy = 'early',
         instrument: InstrumentationSettings | bool | None = None,
@@ -173,6 +178,8 @@ def __init__(
             result_retries: The maximum number of retries to allow for result validation, defaults to `retries`.
             tools: Tools to register with the agent, you can also register tools via the decorators
                 [`@agent.tool`][pydantic_ai.Agent.tool] and [`@agent.tool_plain`][pydantic_ai.Agent.tool_plain].
+            mcp_servers: MCP servers to register with the agent. You should register a [`MCPServer`][pydantic_ai.mcp.MCPServer]
+                for each server you want the agent to connect to.
             defer_model_check: by default, if you provide a [named][pydantic_ai.models.KnownModelName] model,
                 it's evaluated to create a [`Model`][pydantic_ai.models.Model] instance immediately,
                 which checks for the necessary environment variables. Set this to `false`
@@ -215,6 +222,7 @@ def __init__(
 
         self._default_retries = retries
         self._max_result_retries = result_retries if result_retries is not None else retries
+        self._mcp_servers = mcp_servers
         for tool in tools:
             if isinstance(tool, Tool):
                 self._register_tool(tool)
@@ -461,6 +469,7 @@ async def main():
             result_tools=self._result_schema.tool_defs() if self._result_schema else [],
             result_validators=result_validators,
             function_tools=self._function_tools,
+            mcp_servers=self._mcp_servers,
             run_span=run_span,
             tracer=tracer,
         )
@@ -1253,6 +1262,20 @@ def is_end_node(
         """
         return isinstance(node, End)
 
+    @asynccontextmanager
+    async def run_mcp_servers(self) -> AsyncIterator[None]:
+        """Run [`MCPServerStdio`s][pydantic_ai.mcp.MCPServerStdio] so they can be used by the agent.
+
+        Returns: a context manager to start and shutdown the servers.
+        """
+        exit_stack = AsyncExitStack()
+        try:
+            for mcp_server in self._mcp_servers:
+                await exit_stack.enter_async_context(mcp_server)
+            yield
+        finally:
+            await exit_stack.aclose()
+
 
 @dataclasses.dataclass(repr=False)
 class AgentRun(Generic[AgentDepsT, ResultDataT]):