Replace deprecated pydantic-ai MCP classes with MCPToolset in common.ai (#69006)

Author: kaxil

providers/common/ai/docs/toolsets.rst

@@ -42,11 +42,10 @@ passed to any pydantic-ai ``Agent``, including via
 .. note::
 
     ``AgentOperator`` accepts **any** ``AbstractToolset`` implementation — not
-    just the Airflow-native toolsets above. PydanticAI's own MCP server
-    classes (``MCPServerStreamableHTTP``, ``MCPServerSSE``, ``MCPServerStdio``)
-    and third-party toolsets work too. The Airflow-native toolsets add
-    connection management, secret backend integration, and the connection UI,
-    but you are not locked in.
+    just the Airflow-native toolsets above. PydanticAI's own ``MCPToolset``
+    (built over a FastMCP transport) and third-party toolsets work too. The
+    Airflow-native toolsets add connection management, secret backend
+    integration, and the connection UI, but you are not locked in.
 
 
 Using Toolsets Directly with PydanticAI
@@ -336,29 +335,30 @@ Using Multiple MCP Servers
         ],
     )
 
-Direct PydanticAI MCP Servers
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+Direct PydanticAI MCP Toolsets
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
-For prototyping or when you want full PydanticAI control, you can pass MCP
-server instances directly — no Airflow connection needed:
+For prototyping or when you want full PydanticAI control, you can pass
+``MCPToolset`` instances directly — no Airflow connection needed:
 
 .. code-block:: python
 
-    from pydantic_ai.mcp import MCPServerStreamableHTTP, MCPServerStdio
+    from fastmcp.client.transports import StdioTransport
+    from pydantic_ai.mcp import MCPToolset
 
     AgentOperator(
         task_id="direct_mcp",
         prompt="What tools are available?",
         llm_conn_id="pydanticai_default",
         toolsets=[
-            MCPServerStreamableHTTP("http://localhost:3001/mcp"),
-            MCPServerStdio("uvx", args=["mcp-run-python"]),
+            MCPToolset("http://localhost:3001/mcp"),
+            MCPToolset(StdioTransport(command="uvx", args=["mcp-run-python"])),
         ],
     )
 
-This works because PydanticAI's MCP server classes implement
-``AbstractToolset``. The tradeoff: URLs and credentials are hardcoded in DAG
-code instead of being managed through Airflow connections and secret backends.
+This works because PydanticAI's ``MCPToolset`` implements ``AbstractToolset``.
+The tradeoff: URLs and credentials are hardcoded in DAG code instead of being
+managed through Airflow connections and secret backends.
 
 
 .. _agent-skills:

providers/common/ai/src/airflow/providers/common/ai/example_dags/example_mcp.py

@@ -74,19 +74,20 @@ def example_mcp_multiple_servers():
 
 
 # ---------------------------------------------------------------------------
-# 3. Direct PydanticAI MCP servers (no Airflow connection needed)
+# 3. Direct PydanticAI MCP toolsets (no Airflow connection needed)
 # ---------------------------------------------------------------------------
-# AgentOperator accepts any PydanticAI AbstractToolset, including MCP servers
+# AgentOperator accepts any PydanticAI AbstractToolset, including MCPToolset
 # directly. Use this for prototyping or when you want full PydanticAI control.
 #
-#   from pydantic_ai.mcp import MCPServerStreamableHTTP, MCPServerStdio
+#   from fastmcp.client.transports import StdioTransport
+#   from pydantic_ai.mcp import MCPToolset
 #
 #   AgentOperator(
 #       task_id="direct_mcp",
 #       prompt="What tools are available?",
 #       llm_conn_id="pydanticai_default",
 #       toolsets=[
-#           MCPServerStreamableHTTP("http://localhost:3001/mcp"),
-#           MCPServerStdio("uvx", args=["mcp-run-python"]),
+#           MCPToolset("http://localhost:3001/mcp"),
+#           MCPToolset(StdioTransport(command="uvx", args=["mcp-run-python"])),
 #       ],
 #   )

providers/common/ai/src/airflow/providers/common/ai/hooks/mcp.py

@@ -47,7 +47,7 @@ class MCPHook(BaseHook):
         - **Extra.transport**: Transport type — ``http`` (default), ``sse``, or ``stdio``
         - **Extra.command**: Command to run for stdio transport (e.g. ``uvx``)
         - **Extra.args**: Command arguments for stdio transport (e.g. ``["mcp-run-python"]``)
-        - **Extra.timeout**: Connection timeout in seconds for stdio (default: 10)
+        - **Extra.timeout**: Connection init timeout in seconds for stdio (default: 10)
 
     For HTTP/SSE transports the ``Authorization`` header is, by default, a static
     ``Bearer`` token taken from the connection ``password``. Endpoints that require
@@ -121,22 +121,27 @@ def _auth_headers(self, conn: Any) -> dict[str, str] | None:
 
     def get_conn(self) -> Any:
         """
-        Return a configured PydanticAI MCP server instance.
+        Return a configured PydanticAI MCP toolset instance.
 
-        Creates the appropriate MCP server based on the transport type
-        in the connection's extra field:
+        Builds a :class:`~pydantic_ai.mcp.MCPToolset` over the FastMCP transport
+        matching the transport type in the connection's extra field:
 
-        - ``http`` (default): :class:`~pydantic_ai.mcp.MCPServerStreamableHTTP`
-        - ``sse``: :class:`~pydantic_ai.mcp.MCPServerSSE`
-        - ``stdio``: :class:`~pydantic_ai.mcp.MCPServerStdio`
+        - ``http`` (default): ``fastmcp.client.transports.StreamableHttpTransport``
+        - ``sse``: ``fastmcp.client.transports.SSETransport``
+        - ``stdio``: ``fastmcp.client.transports.StdioTransport``
+
+        When ``tool_prefix`` is set the toolset is wrapped via
+        :meth:`~pydantic_ai.toolsets.abstract.AbstractToolset.prefixed`, so a
+        prefix of ``"weather"`` yields tool names like ``weather_get_forecast``.
 
         The result is cached for the lifetime of this hook instance.
         """
         if self._server is not None:
             return self._server
 
         try:
-            from pydantic_ai.mcp import MCPServerSSE, MCPServerStdio, MCPServerStreamableHTTP
+            from fastmcp.client.transports import SSETransport, StdioTransport, StreamableHttpTransport
+            from pydantic_ai.mcp import MCPToolset
         except ImportError:
             raise ImportError(
                 'MCP support requires the `mcp` package. Install it with: pip install "pydantic-ai-slim[mcp]"'
@@ -149,15 +154,11 @@ def get_conn(self) -> Any:
         if transport == "http":
             if not conn.host:
                 raise ValueError(f"Connection {self.mcp_conn_id!r} requires a host URL for HTTP transport.")
-            self._server = MCPServerStreamableHTTP(
-                conn.host, headers=self._auth_headers(conn), tool_prefix=self.tool_prefix
-            )
+            toolset = MCPToolset(StreamableHttpTransport(conn.host, headers=self._auth_headers(conn)))
         elif transport == "sse":
             if not conn.host:
                 raise ValueError(f"Connection {self.mcp_conn_id!r} requires a host URL for SSE transport.")
-            self._server = MCPServerSSE(
-                conn.host, headers=self._auth_headers(conn), tool_prefix=self.tool_prefix
-            )
+            toolset = MCPToolset(SSETransport(conn.host, headers=self._auth_headers(conn)))
         elif transport == "stdio":
             command = extra.get("command")
             if not command:
@@ -168,13 +169,14 @@ def get_conn(self) -> Any:
             if isinstance(args, str):
                 args = [args]
             timeout = extra.get("timeout", 10)
-            self._server = MCPServerStdio(command, args=args, timeout=timeout, tool_prefix=self.tool_prefix)
+            toolset = MCPToolset(StdioTransport(command=command, args=args), init_timeout=timeout)
         else:
             raise ValueError(
                 f"Unknown transport {transport!r} in connection {self.mcp_conn_id!r}. "
                 "Supported: 'http', 'sse', 'stdio'."
             )
 
+        self._server = toolset.prefixed(self.tool_prefix) if self.tool_prefix else toolset
         return self._server
 
     def test_connection(self) -> tuple[bool, str]:

providers/common/ai/src/airflow/providers/common/ai/toolsets/langchain_bridge.py

@@ -97,9 +97,9 @@ def airflow_toolset_to_langchain_tools(
         and every ``call_tool`` each run under their own event loop, and pydantic-ai
         opens and tears the connection down around each one. For ``MCPToolset`` this
         means the server is reconnected on every tool call. That is fine for
-        stateless tools (and for HTTP/SSE servers, modulo per-call latency), but an
-        ``MCPServerStdio`` server, or any server that keeps state between calls,
-        will lose that state because each call starts a fresh process/session.
+        stateless tools (and for HTTP/SSE servers, modulo per-call latency), but a
+        stdio server, or any server that keeps state between calls, will lose that
+        state because each call starts a fresh process/session.
 
     .. note::
         A pydantic-ai toolset is normally driven inside an agent run, where a

providers/common/ai/src/airflow/providers/common/ai/toolsets/mcp.py

@@ -35,18 +35,17 @@ class MCPToolset(AbstractToolset[Any]):
 
     Reads MCP server transport type, URL, command, and credentials from the
     connection via :class:`~airflow.providers.common.ai.hooks.mcp.MCPHook` and
-    creates the appropriate PydanticAI MCP server instance.
-    All ``AbstractToolset`` methods delegate to the underlying MCP server.
+    builds the matching PydanticAI :class:`~pydantic_ai.mcp.MCPToolset`.
+    All ``AbstractToolset`` methods delegate to the underlying MCP toolset.
 
     This is the recommended way to use MCP servers in Airflow — it stores
     server configuration in Airflow connections (and secret backends) rather
     than hard-coding URLs and credentials in DAG code.
 
-    If you prefer full PydanticAI control, you can pass MCP server instances
-    directly to ``AgentOperator(toolsets=[...])``, since
-    :class:`~pydantic_ai.mcp.MCPServerStreamableHTTP`,
-    :class:`~pydantic_ai.mcp.MCPServerSSE`, and
-    :class:`~pydantic_ai.mcp.MCPServerStdio` all implement ``AbstractToolset``.
+    If you prefer full PydanticAI control, you can pass a
+    :class:`~pydantic_ai.mcp.MCPToolset` (built over a FastMCP transport)
+    directly to ``AgentOperator(toolsets=[...])``, since it implements
+    ``AbstractToolset``.
 
     For MCP endpoints that need a freshly minted or short-lived token (e.g. a
     Snowflake managed MCP server authenticated with a key-pair JWT, or OAuth /