Skip to content

fix(pydantic): import restate.ext.pydantic without the pydantic-ai mcp extra - #257

Closed
pcbeingused333 wants to merge 3 commits into
restatedev:mainfrom
pcbeingused333:fix/pydantic-ext-optional-mcp
Closed

pcbeingused333 wants to merge 3 commits into
restatedev:mainfrom
pcbeingused333:fix/pydantic-ext-optional-mcp

Conversation

@pcbeingused333

@pcbeingused333 pcbeingused333 commented Nov 3, 2026 •

Copy link
Copy Markdown
Contributor

Fixes #221.

python/restate/ext/pydantic/_toolset.py imported pydantic_ai.mcp.MCPToolset at module level. restate.ext.pydantic imports that module from __init__ and _agent, so the whole extension failed to import unless pydantic-ai-slim[mcp] (fastmcp) was installed, even though _agent.py already imports MCPToolset and RestateMCPToolset lazily inside a try/except ImportError and the pydantic_ai extra only pins pydantic-ai-slim.

RestateMCPToolset and its two serde dataclasses now live in a new _mcp_toolset.py, which is only imported from that guarded block. _toolset.py keeps RestateContextRunToolSet unchanged. Nothing else imported RestateMCPToolset from _toolset.


This pull request was prepared with the assistance of an AI coding agent (Claude Code) and submitted by Alex Castillo, who is responsible for its content.

…p extra

_toolset.py imported pydantic_ai.mcp at module level, so importing
restate.ext.pydantic failed unless pydantic-ai-slim[mcp] was installed,
although _agent.py already treats MCP support as optional. Move
RestateMCPToolset into its own module that is only imported lazily.
@github-actions

github-actions Bot commented Nov 3, 2026 •

Copy link
Copy Markdown

All contributors have signed the CLA ✍️ ✅
Posted by the CLA Assistant Lite bot.

@pcbeingused333

Copy link
Copy Markdown
Contributor Author

I have read the CLA Document and I hereby sign the CLA

@igalshilman

Copy link
Copy Markdown
Contributor

looks good thank you for your contribution.
I have left one comment, and other than that looks good.

@igalshilman

Copy link
Copy Markdown
Contributor

@pullfrog review

@pullfrog pullfrog Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

✅ No new issues found.

Reviewed changes

Reviewed the change that moves the MCP toolset wrapper out of _toolset.py, so restate.ext.pydantic can be imported without the pydantic-ai-slim[mcp] extra.

  • MCP wrapper moved to _mcp_toolset.py: RestateMCPToolset, RestateMCPGetToolsContextRunResult, RestateMCPToolRunResult and their serdes are moved over unchanged, and the top-level from pydantic_ai.mcp import MCPToolset goes with them.
  • Lazy import updated: the guarded try/except ImportError in RestateAgent.__init__'s set_context now imports from ._mcp_toolset. pydantic_ai/mcp.py re-raises a missing mcp/fastmcp as ImportError, so the guard catches it.
  • Regression test: tests/pydantic_ai_without_mcp.py sets sys.modules['pydantic_ai.mcp'] = None in a subprocess and imports the extension. On the base commit this fails, because _toolset.py imports pydantic_ai.mcp at module level.

What I checked:

  • main (0d06f2c) still has the module-level import, so the fix is still needed.
  • Nothing else on main imports RestateMCPToolset or the MCP serdes from _toolset.
  • GitHub reports the PR as mergeable against current main.
  • On the PR head, the new test passes, and ruff check, ruff format --check and pyright on python/restate/ext/pydantic/ all pass.

Pullfrog  | View workflow run | Using claude-opus-5-5 | 𝕏

# Conflicts:
#	python/restate/ext/pydantic/_toolset.py
Comment thread tests/pydantic_ai_without_mcp.py Outdated


def test_import_without_pydantic_ai_mcp():
"""restate.ext.pydantic must be importable without the pydantic-ai ``mcp`` extra."""

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Lets remove this test please

Comment thread tests/pydantic_ai_without_mcp.py Outdated
pytest.importorskip("pydantic_ai")


def test_import_without_pydantic_ai_mcp():

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

please remove this entire test.

@igalshilman

Copy link
Copy Markdown
Contributor

Merged with your authorship kept in #307 (the test left out, as the review asked). Thank you!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

restate.ext.pydantic requires pydantic-ai-slim[mcp] because _toolset.py imports pydantic_ai.mcp unguarded

2 participants