Make FastMCPToolset work with Temporal

Author: DouweM

pydantic_ai_slim/pydantic_ai/durable_exec/temporal/_fastmcp_toolset.py

@@ -0,0 +1,146 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any, Literal
+
+from pydantic import ConfigDict, with_config
+from temporalio import activity, workflow
+from temporalio.workflow import ActivityConfig
+from typing_extensions import Self
+
+from pydantic_ai import ToolsetTool
+from pydantic_ai.exceptions import UserError
+from pydantic_ai.tools import AgentDepsT, RunContext, ToolDefinition
+from pydantic_ai.toolsets.fastmcp import FastMCPToolset
+
+from ._run_context import TemporalRunContext
+from ._toolset import (
+    CallToolParams,
+    CallToolResult,
+    TemporalWrapperToolset,
+)
+
+
+@dataclass
+@with_config(ConfigDict(arbitrary_types_allowed=True))
+class _GetToolsParams:
+    serialized_run_context: Any
+
+
+class TemporalFastMCPToolset(TemporalWrapperToolset[AgentDepsT]):
+    def __init__(
+        self,
+        server: FastMCPToolset[AgentDepsT],
+        *,
+        activity_name_prefix: str,
+        activity_config: ActivityConfig,
+        tool_activity_config: dict[str, ActivityConfig | Literal[False]],
+        deps_type: type[AgentDepsT],
+        run_context_type: type[TemporalRunContext[AgentDepsT]] = TemporalRunContext[AgentDepsT],
+    ):
+        super().__init__(server)
+        self.activity_config = activity_config
+
+        self.tool_activity_config: dict[str, ActivityConfig] = {}
+        for tool_name, tool_config in tool_activity_config.items():
+            if tool_config is False:
+                raise UserError(
+                    f'Temporal activity config for MCP tool {tool_name!r} has been explicitly set to `False` (activity disabled), '
+                    'but MCP tools require the use of IO and so cannot be run outside of an activity.'
+                )
+            self.tool_activity_config[tool_name] = tool_config
+
+        self.run_context_type = run_context_type
+
+        async def get_tools_activity(params: _GetToolsParams, deps: AgentDepsT) -> dict[str, ToolDefinition]:
+            run_context = self.run_context_type.deserialize_run_context(params.serialized_run_context, deps=deps)
+            tools = await self.wrapped.get_tools(run_context)
+            # 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 {name: tool.tool_def for name, tool in tools.items()}
+
+        # Set type hint explicitly so that Temporal can take care of serialization and deserialization
+        get_tools_activity.__annotations__['deps'] = deps_type
+
+        self.get_tools_activity = activity.defn(name=f'{activity_name_prefix}__mcp_server__{self.id}__get_tools')(
+            get_tools_activity
+        )
+
+        async def call_tool_activity(params: CallToolParams, deps: AgentDepsT) -> CallToolResult:
+            run_context = self.run_context_type.deserialize_run_context(params.serialized_run_context, deps=deps)
+            assert isinstance(params.tool_def, ToolDefinition)
+            return await self._wrap_call_tool_result(
+                self.wrapped.call_tool(
+                    params.name,
+                    params.tool_args,
+                    run_context,
+                    self.tool_for_tool_def(params.tool_def),
+                )
+            )
+
+        # Set type hint explicitly so that Temporal can take care of serialization and deserialization
+        call_tool_activity.__annotations__['deps'] = deps_type
+
+        self.call_tool_activity = activity.defn(name=f'{activity_name_prefix}__mcp_server__{self.id}__call_tool')(
+            call_tool_activity
+        )
+
+    def tool_for_tool_def(self, tool_def: ToolDefinition) -> ToolsetTool[AgentDepsT]:
+        assert isinstance(self.wrapped, FastMCPToolset)
+        return self.wrapped.tool_for_tool_def(tool_def)
+
+    @property
+    def temporal_activities(self) -> list[Callable[..., Any]]:
+        return [self.get_tools_activity, self.call_tool_activity]
+
+    async def __aenter__(self) -> Self:
+        # The wrapped MCPServer enters itself around listing and calling tools
+        # so we don't need to enter it here (nor could we because we're not inside a Temporal activity).
+        return self
+
+    async def __aexit__(self, *args: Any) -> bool | None:
+        return None
+
+    async def get_tools(self, ctx: RunContext[AgentDepsT]) -> dict[str, ToolsetTool[AgentDepsT]]:
+        if not workflow.in_workflow():
+            return await super().get_tools(ctx)
+
+        serialized_run_context = self.run_context_type.serialize_run_context(ctx)
+        tool_defs = await workflow.execute_activity(  # pyright: ignore[reportUnknownMemberType]
+            activity=self.get_tools_activity,
+            args=[
+                _GetToolsParams(serialized_run_context=serialized_run_context),
+                ctx.deps,
+            ],
+            **self.activity_config,
+        )
+        return {name: self.tool_for_tool_def(tool_def) for name, tool_def in tool_defs.items()}
+
+    async def call_tool(
+        self,
+        name: str,
+        tool_args: dict[str, Any],
+        ctx: RunContext[AgentDepsT],
+        tool: ToolsetTool[AgentDepsT],
+    ) -> CallToolResult:
+        if not workflow.in_workflow():
+            return await super().call_tool(name, tool_args, ctx, tool)
+
+        tool_activity_config = self.activity_config | self.tool_activity_config.get(name, {})
+        serialized_run_context = self.run_context_type.serialize_run_context(ctx)
+        return self._unwrap_call_tool_result(
+            await workflow.execute_activity(  # pyright: ignore[reportUnknownMemberType]
+                activity=self.call_tool_activity,
+                args=[
+                    CallToolParams(
+                        name=name,
+                        tool_args=tool_args,
+                        serialized_run_context=serialized_run_context,
+                        tool_def=tool.tool_def,
+                    ),
+                    ctx.deps,
+                ],
+                **tool_activity_config,
+            )
+        )

pydantic_ai_slim/pydantic_ai/durable_exec/temporal/_toolset.py

@@ -142,4 +142,21 @@ def temporalize_toolset(
                 run_context_type=run_context_type,
             )
 
+    try:
+        from pydantic_ai.toolsets.fastmcp import FastMCPToolset
+
+        from ._fastmcp_toolset import TemporalFastMCPToolset
+    except ImportError:
+        pass
+    else:
+        if isinstance(toolset, FastMCPToolset):
+            return TemporalFastMCPToolset(
+                toolset,
+                activity_name_prefix=activity_name_prefix,
+                activity_config=activity_config,
+                tool_activity_config=tool_activity_config,
+                deps_type=deps_type,
+                run_context_type=run_context_type,
+            )
+
     return toolset

pydantic_ai_slim/pydantic_ai/toolsets/fastmcp.py

@@ -32,7 +32,6 @@
         ResourceLink,
         TextContent,
         TextResourceContents,
-        Tool as MCPTool,
     )
 
     from pydantic_ai.mcp import TOOL_SCHEMA_VALIDATOR
@@ -131,11 +130,20 @@ async def __aexit__(self, *args: Any) -> bool | None:
 
     async def get_tools(self, ctx: RunContext[AgentDepsT]) -> dict[str, ToolsetTool[AgentDepsT]]:
         async with self:
-            mcp_tools: list[MCPTool] = await self.client.list_tools()
-
             return {
-                tool.name: _convert_mcp_tool_to_toolset_tool(toolset=self, mcp_tool=tool, retries=self.max_retries)
-                for tool in mcp_tools
+                mcp_tool.name: self.tool_for_tool_def(
+                    ToolDefinition(
+                        name=mcp_tool.name,
+                        description=mcp_tool.description,
+                        parameters_json_schema=mcp_tool.inputSchema,
+                        metadata={
+                            'meta': mcp_tool.meta,
+                            'annotations': mcp_tool.annotations.model_dump() if mcp_tool.annotations else None,
+                            'output_schema': mcp_tool.outputSchema or None,
+                        },
+                    )
+                )
+                for mcp_tool in await self.client.list_tools()
             }
 
     async def call_tool(
@@ -157,28 +165,13 @@ async def call_tool(
         # Otherwise, return the content
         return _map_fastmcp_tool_results(parts=call_tool_result.content)
 
-
-def _convert_mcp_tool_to_toolset_tool(
-    toolset: FastMCPToolset[AgentDepsT],
-    mcp_tool: MCPTool,
-    retries: int,
-) -> ToolsetTool[AgentDepsT]:
-    """Convert an MCP tool to a toolset tool."""
-    return ToolsetTool[AgentDepsT](
-        tool_def=ToolDefinition(
-            name=mcp_tool.name,
-            description=mcp_tool.description,
-            parameters_json_schema=mcp_tool.inputSchema,
-            metadata={
-                'meta': mcp_tool.meta,
-                'annotations': mcp_tool.annotations.model_dump() if mcp_tool.annotations else None,
-                'output_schema': mcp_tool.outputSchema or None,
-            },
-        ),
-        toolset=toolset,
-        max_retries=retries,
-        args_validator=TOOL_SCHEMA_VALIDATOR,
-    )
+    def tool_for_tool_def(self, tool_def: ToolDefinition) -> ToolsetTool[AgentDepsT]:
+        return ToolsetTool[AgentDepsT](
+            tool_def=tool_def,
+            toolset=self,
+            max_retries=self.max_retries,
+            args_validator=TOOL_SCHEMA_VALIDATOR,
+        )
 
 
 def _map_fastmcp_tool_results(parts: list[ContentBlock]) -> list[FastMCPToolResult] | FastMCPToolResult: