Skip to content

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

Merged
kaxil merged 2 commits into
apache:mainfrom
astronomer:commonai-mcp-toolset-migration
Jun 27, 2026
Merged

Replace deprecated pydantic-ai MCP classes with MCPToolset in common.ai#69006
kaxil merged 2 commits into
apache:mainfrom
astronomer:commonai-mcp-toolset-migration

Conversation

@kaxil

@kaxil kaxil commented Jun 26, 2026

Copy link
Copy Markdown
Member

pydantic-ai deprecated MCPServerStdio, MCPServerSSE, and MCPServerStreamableHTTP in favor of a unified MCPToolset, and removes all three in pydantic-ai v2. The common.ai MCPHook constructed all three directly, so it emits DeprecationWarnings today and would break on the v2 upgrade.

This migrates MCPHook.get_conn() to build an MCPToolset over the matching FastMCP transport.

What changed

Transport Before After
http MCPServerStreamableHTTP(host, headers=…, tool_prefix=…) MCPToolset(StreamableHttpTransport(host, headers=…))
sse MCPServerSSE(host, headers=…, tool_prefix=…) MCPToolset(SSETransport(host, headers=…))
stdio MCPServerStdio(cmd, args=…, timeout=…, tool_prefix=…) MCPToolset(StdioTransport(command=cmd, args=…), init_timeout=…)
prefix tool_prefix= constructor arg .prefixed(prefix)

Design rationale

  • Explicit transports instead of URL inference. MCPToolset can infer a transport from a bare URL, but FastMCP only picks SSE when the URL ends in /sse. The connection has an explicit transport field, so the hook builds the transport explicitly to honor that choice regardless of URL shape, and to keep attaching the Authorization header.
  • StdioTransport for stdio. The connection supplies an arbitrary command + args, which is the case pydantic-ai's deprecation message prescribes StdioTransport(...) for; the MCPToolset("script.py") shorthand only covers literal script paths.
  • .prefixed() for tool_prefix. MCPToolset has no tool_prefix argument. .prefixed(prefix) produces the same <prefix>_<tool> names as before and matches pydantic-ai's own load_mcp_toolsets. It returns a PrefixedToolset that proxies the async-context-manager protocol, so the connection-backed MCPToolset wrapper keeps working unchanged.
  • init_timeout is the successor to the old stdio timeout; the hook keeps its default of 10s.

Tradeoffs / notes

  • Behavior-preserving: tool names, auth headers, caching defaults, and the no-auth path are unchanged.
  • fastmcp.client.transports is the canonical public import path. pydantic_ai.mcp imports from it, and pydantic-ai's own deprecation message recommends MCPToolset(fastmcp.client.transports.StdioTransport(...)).

Was generative AI tooling used to co-author this PR?
  • Yes (please specify the tool below)

  • Read the Pull Request Guidelines for more information. Note: commit author/co-author name and email in commits become permanently public when merged.
  • For fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
  • When adding dependency, check compliance with the ASF 3rd Party License Policy.
  • For significant user-facing changes create newsfragment: {pr_number}.significant.rst, in airflow-core/newsfragments. You can add this file in a follow-up commit after the PR is created so you know the PR number.

…common.ai

pydantic-ai deprecated MCPServerStdio, MCPServerSSE, and MCPServerStreamableHTTP
in favor of a unified MCPToolset (all three are removed in pydantic-ai v2). The
common.ai MCPHook constructed all three, so it emitted DeprecationWarnings and
would break on the v2 upgrade.

Build an MCPToolset over the matching FastMCP transport instead:

- http  -> MCPToolset(StreamableHttpTransport(host, headers=...))
- sse   -> MCPToolset(SSETransport(host, headers=...))
- stdio -> MCPToolset(StdioTransport(command=..., args=...), init_timeout=...)

Transports are built explicitly rather than relying on FastMCP's URL inference,
so the connection's explicit 'transport' field is honored (inference only picks
SSE for URLs ending in /sse) and the Authorization header is preserved.

MCPToolset has no tool_prefix argument, so the prefix is applied via
.prefixed(), matching pydantic-ai's own load_mcp_toolsets and producing the same
'<prefix>_<tool>' names as before.
@kaxil
kaxil requested a review from gopidesupavan as a code owner June 26, 2026 01:12
@kaxil kaxil changed the title Replace deprecated pydantic-ai MCP classes with MCPToolset in common.ai Replace deprecated pydantic-ai MCP classes with MCPToolset in common.ai Jun 26, 2026

@gopidesupavan gopidesupavan left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM thanks kaxil :)

not a blocker can we do autospec for tests ?
@patch(_MCP_TOOLSET, autospec=True)
@patch(_HTTP_TRANSPORT, autospec=True)

@kaxil

kaxil commented Jun 27, 2026

Copy link
Copy Markdown
Member Author

Thanks @gopidesupavan! Good call -- done in 42c8b27. All the @patch decorators (and the mask_secret patch) now use autospec=True.

@kaxil
kaxil merged commit a1daaf1 into apache:main Jun 27, 2026
81 checks passed
@kaxil
kaxil deleted the commonai-mcp-toolset-migration branch June 27, 2026 14:57
karenbraganz pushed a commit to karenbraganz/airflow that referenced this pull request Jun 30, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants