From 8c4e754a50cd6cc29b6d5269163e5b3658381a4e Mon Sep 17 00:00:00 2001 From: Alex Castillo Date: Tue, 3 Nov 2026 09:16:35 -0500 Subject: [PATCH 1/2] fix(pydantic): import restate.ext.pydantic without the pydantic-ai mcp 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. --- python/restate/ext/pydantic/_agent.py | 2 +- python/restate/ext/pydantic/_mcp_toolset.py | 100 ++++++++++++++++++++ python/restate/ext/pydantic/_toolset.py | 86 +---------------- tests/pydantic_ai_without_mcp.py | 18 ++++ 4 files changed, 120 insertions(+), 86 deletions(-) create mode 100644 python/restate/ext/pydantic/_mcp_toolset.py create mode 100644 tests/pydantic_ai_without_mcp.py diff --git a/python/restate/ext/pydantic/_agent.py b/python/restate/ext/pydantic/_agent.py index cc7888cf..dfe159df 100644 --- a/python/restate/ext/pydantic/_agent.py +++ b/python/restate/ext/pydantic/_agent.py @@ -121,7 +121,7 @@ def set_context(toolset: AbstractToolset[AgentDepsT]) -> AbstractToolset[AgentDe try: from pydantic_ai.mcp import MCPToolset - from ._toolset import RestateMCPToolset + from ._mcp_toolset import RestateMCPToolset except ImportError: pass else: diff --git a/python/restate/ext/pydantic/_mcp_toolset.py b/python/restate/ext/pydantic/_mcp_toolset.py new file mode 100644 index 00000000..d303b300 --- /dev/null +++ b/python/restate/ext/pydantic/_mcp_toolset.py @@ -0,0 +1,100 @@ +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass, replace +from typing import Any + +from restate import RunOptions, SdkInternalBaseException +from restate.extensions import current_context + +from pydantic_ai import ToolDefinition +from pydantic_ai._run_context import AgentDepsT +from pydantic_ai.exceptions import UserError +from pydantic_ai.mcp import MCPToolset +from pydantic_ai.tools import RunContext +from pydantic_ai.toolsets.abstract import AbstractToolset, ToolsetTool +from pydantic_ai.toolsets.wrapper import WrapperToolset + +from ._serde import PydanticTypeAdapter + + +@dataclass +class RestateMCPGetToolsContextRunResult: + """A simple wrapper for tool results to be used with Restate's run_typed.""" + + output: dict[str, ToolDefinition] + + +MCP_GET_TOOLS_SERDE = PydanticTypeAdapter(RestateMCPGetToolsContextRunResult) + + +@dataclass +class RestateMCPToolRunResult: + """A simple wrapper for tool results to be used with Restate's run_typed.""" + + output: Any + + +MCP_RUN_SERDE = PydanticTypeAdapter(RestateMCPToolRunResult) + + +class RestateMCPToolset(WrapperToolset[AgentDepsT]): + """A wrapper for MCPToolset that integrates with Restate.""" + + def __init__(self, wrapped: MCPToolset[AgentDepsT], run_options: RunOptions): + super().__init__(wrapped) + self._wrapped = wrapped + self.get_tools_options = replace(run_options, serde=MCP_GET_TOOLS_SERDE) + self.run_tools_options = replace(run_options, serde=MCP_RUN_SERDE) + + def visit_and_replace( + self, visitor: Callable[[AbstractToolset[AgentDepsT]], AbstractToolset[AgentDepsT]] + ) -> AbstractToolset[AgentDepsT]: + return visitor(self) + + async def get_tools(self, ctx: RunContext[AgentDepsT]) -> dict[str, ToolsetTool[AgentDepsT]]: + async def get_tools_in_context() -> RestateMCPGetToolsContextRunResult: + res = await self._wrapped.get_tools(ctx) + # ToolsetTool is not serializable as it holds a SchemaValidator + # (which is also the same for every MCP tool so unnecessary to pass along the wire every time), + # so we just return the ToolDefinitions and wrap them in ToolsetTool outside of the activity. + return RestateMCPGetToolsContextRunResult(output={name: tool.tool_def for name, tool in res.items()}) + + context = current_context() + if context is None: + raise UserError( + "A toolset cannot be used without a Restate context. Make sure to run it within an agent or a run context." + ) + + try: + tool_defs = await context.run_typed("get mcp tools", get_tools_in_context, self.get_tools_options) + return {name: self.tool_for_tool_def(tool_def) for name, tool_def in tool_defs.output.items()} + except SdkInternalBaseException as e: + raise Exception("Internal error during get_tools call") from e + + def tool_for_tool_def(self, tool_def: ToolDefinition) -> ToolsetTool[AgentDepsT]: + assert isinstance(self.wrapped, MCPToolset) + return self.wrapped.tool_for_tool_def(tool_def) + + async def call_tool( + self, + name: str, + tool_args: dict[str, Any], + ctx: RunContext[AgentDepsT], + tool: ToolsetTool[AgentDepsT], + ) -> Any: + async def call_tool_in_context() -> RestateMCPToolRunResult: + res = await self._wrapped.call_tool(name, tool_args, ctx, tool) + return RestateMCPToolRunResult(output=res) + + context = current_context() + if context is None: + raise UserError( + "A toolset cannot be used without a Restate context. Make sure to run it within an agent or a run context." + ) + try: + res = await context.run_typed(f"Calling mcp tool {name}", call_tool_in_context, self.run_tools_options) + except SdkInternalBaseException as e: + raise Exception("Internal error during tool call") from e + + return res.output diff --git a/python/restate/ext/pydantic/_toolset.py b/python/restate/ext/pydantic/_toolset.py index 2f62db73..7852f4d9 100644 --- a/python/restate/ext/pydantic/_toolset.py +++ b/python/restate/ext/pydantic/_toolset.py @@ -4,14 +4,12 @@ from dataclasses import dataclass, replace from typing import Any, Literal -from restate import RunOptions, SdkInternalBaseException, TerminalError +from restate import RunOptions, TerminalError from restate.exceptions import SdkInternalException from restate.extensions import current_context -from pydantic_ai import ToolDefinition from pydantic_ai._run_context import AgentDepsT from pydantic_ai.exceptions import ApprovalRequired, CallDeferred, ModelRetry, UserError -from pydantic_ai.mcp import MCPToolset from pydantic_ai.tools import RunContext from pydantic_ai.toolsets.abstract import AbstractToolset, ToolsetTool from pydantic_ai.toolsets.wrapper import WrapperToolset @@ -32,26 +30,6 @@ class RestateContextRunResult: CONTEXT_RUN_SERDE = PydanticTypeAdapter(RestateContextRunResult) -@dataclass -class RestateMCPGetToolsContextRunResult: - """A simple wrapper for tool results to be used with Restate's run_typed.""" - - output: dict[str, ToolDefinition] - - -MCP_GET_TOOLS_SERDE = PydanticTypeAdapter(RestateMCPGetToolsContextRunResult) - - -@dataclass -class RestateMCPToolRunResult: - """A simple wrapper for tool results to be used with Restate's run_typed.""" - - output: Any - - -MCP_RUN_SERDE = PydanticTypeAdapter(RestateMCPToolRunResult) - - class RestateContextRunToolSet(WrapperToolset[AgentDepsT]): """A toolset that automatically wraps tool calls with restate's `ctx.run_typed()`.""" @@ -114,65 +92,3 @@ def visit_and_replace( self, visitor: Callable[[AbstractToolset[AgentDepsT]], AbstractToolset[AgentDepsT]] ) -> AbstractToolset[AgentDepsT]: return visitor(self) - - -class RestateMCPToolset(WrapperToolset[AgentDepsT]): - """A wrapper for MCPToolset that integrates with Restate.""" - - def __init__(self, wrapped: MCPToolset[AgentDepsT], run_options: RunOptions): - super().__init__(wrapped) - self._wrapped = wrapped - self.get_tools_options = replace(run_options, serde=MCP_GET_TOOLS_SERDE) - self.run_tools_options = replace(run_options, serde=MCP_RUN_SERDE) - - def visit_and_replace( - self, visitor: Callable[[AbstractToolset[AgentDepsT]], AbstractToolset[AgentDepsT]] - ) -> AbstractToolset[AgentDepsT]: - return visitor(self) - - async def get_tools(self, ctx: RunContext[AgentDepsT]) -> dict[str, ToolsetTool[AgentDepsT]]: - async def get_tools_in_context() -> RestateMCPGetToolsContextRunResult: - res = await self._wrapped.get_tools(ctx) - # ToolsetTool is not serializable as it holds a SchemaValidator - # (which is also the same for every MCP tool so unnecessary to pass along the wire every time), - # so we just return the ToolDefinitions and wrap them in ToolsetTool outside of the activity. - return RestateMCPGetToolsContextRunResult(output={name: tool.tool_def for name, tool in res.items()}) - - context = current_context() - if context is None: - raise UserError( - "A toolset cannot be used without a Restate context. Make sure to run it within an agent or a run context." - ) - - try: - tool_defs = await context.run_typed("get mcp tools", get_tools_in_context, self.get_tools_options) - return {name: self.tool_for_tool_def(tool_def) for name, tool_def in tool_defs.output.items()} - except SdkInternalBaseException as e: - raise Exception("Internal error during get_tools call") from e - - def tool_for_tool_def(self, tool_def: ToolDefinition) -> ToolsetTool[AgentDepsT]: - assert isinstance(self.wrapped, MCPToolset) - return self.wrapped.tool_for_tool_def(tool_def) - - async def call_tool( - self, - name: str, - tool_args: dict[str, Any], - ctx: RunContext[AgentDepsT], - tool: ToolsetTool[AgentDepsT], - ) -> Any: - async def call_tool_in_context() -> RestateMCPToolRunResult: - res = await self._wrapped.call_tool(name, tool_args, ctx, tool) - return RestateMCPToolRunResult(output=res) - - context = current_context() - if context is None: - raise UserError( - "A toolset cannot be used without a Restate context. Make sure to run it within an agent or a run context." - ) - try: - res = await context.run_typed(f"Calling mcp tool {name}", call_tool_in_context, self.run_tools_options) - except SdkInternalBaseException as e: - raise Exception("Internal error during tool call") from e - - return res.output diff --git a/tests/pydantic_ai_without_mcp.py b/tests/pydantic_ai_without_mcp.py new file mode 100644 index 00000000..a2a27c03 --- /dev/null +++ b/tests/pydantic_ai_without_mcp.py @@ -0,0 +1,18 @@ +import subprocess +import sys + +import pytest + +pytest.importorskip("pydantic_ai") + + +def test_import_without_pydantic_ai_mcp(): + """restate.ext.pydantic must be importable without the pydantic-ai ``mcp`` extra.""" + code = ( + "import sys\n" + "sys.modules['pydantic_ai.mcp'] = None # simulate the missing optional dependency\n" + "import restate.ext.pydantic as ext\n" + "assert ext.RestateAgent is not None\n" + ) + result = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, check=False) + assert result.returncode == 0, result.stderr From a13a2601d31fc0b06b364b2dfd29bbafdf357044 Mon Sep 17 00:00:00 2001 From: Alex Castillo Date: Mon, 9 Nov 2026 07:15:25 -0400 Subject: [PATCH 2/2] test: remove pydantic_ai_without_mcp test as requested in review --- tests/pydantic_ai_without_mcp.py | 18 ------------------ 1 file changed, 18 deletions(-) delete mode 100644 tests/pydantic_ai_without_mcp.py diff --git a/tests/pydantic_ai_without_mcp.py b/tests/pydantic_ai_without_mcp.py deleted file mode 100644 index a2a27c03..00000000 --- a/tests/pydantic_ai_without_mcp.py +++ /dev/null @@ -1,18 +0,0 @@ -import subprocess -import sys - -import pytest - -pytest.importorskip("pydantic_ai") - - -def test_import_without_pydantic_ai_mcp(): - """restate.ext.pydantic must be importable without the pydantic-ai ``mcp`` extra.""" - code = ( - "import sys\n" - "sys.modules['pydantic_ai.mcp'] = None # simulate the missing optional dependency\n" - "import restate.ext.pydantic as ext\n" - "assert ext.RestateAgent is not None\n" - ) - result = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, check=False) - assert result.returncode == 0, result.stderr