Fix Temporal and DBOS MCP to use cached tools instead of fetching each time (#4331)

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

Author: DouweM

pydantic_ai_slim/pydantic_ai/durable_exec/dbos/_mcp_server.py

@@ -2,14 +2,18 @@
 
 from pydantic_ai import ToolsetTool
 from pydantic_ai.mcp import MCPServer
-from pydantic_ai.tools import AgentDepsT, ToolDefinition
+from pydantic_ai.tools import AgentDepsT, RunContext, ToolDefinition
 
 from ._mcp import DBOSMCPToolset
 from ._utils import StepConfig
 
 
 class DBOSMCPServer(DBOSMCPToolset[AgentDepsT]):
-    """A wrapper for MCPServer that integrates with DBOS, turning call_tool and get_tools to DBOS steps."""
+    """A wrapper for MCPServer that integrates with DBOS, turning call_tool and get_tools into DBOS steps.
+
+    Tool definitions are cached across steps to avoid redundant MCP server round-trips,
+    respecting the wrapped server's `cache_tools` setting.
+    """
 
     def __init__(
         self,
@@ -23,7 +27,24 @@ def __init__(
             step_name_prefix=step_name_prefix,
             step_config=step_config,
         )
+        # Cached across steps to avoid redundant MCP connections per step.
+        # Not invalidated by `tools/list_changed` notifications — users who need
+        # dynamic tools during a workflow should set `cache_tools=False`.
+        self._cached_tool_defs: dict[str, ToolDefinition] | None = None
 
-    def tool_for_tool_def(self, tool_def: ToolDefinition) -> ToolsetTool[AgentDepsT]:
+    @property
+    def _server(self) -> MCPServer:
         assert isinstance(self.wrapped, MCPServer)
-        return self.wrapped.tool_for_tool_def(tool_def)
+        return self.wrapped
+
+    def tool_for_tool_def(self, tool_def: ToolDefinition) -> ToolsetTool[AgentDepsT]:
+        return self._server.tool_for_tool_def(tool_def)
+
+    async def get_tools(self, ctx: RunContext[AgentDepsT]) -> dict[str, ToolsetTool[AgentDepsT]]:
+        if self._server.cache_tools and self._cached_tool_defs is not None:
+            return {name: self.tool_for_tool_def(td) for name, td in self._cached_tool_defs.items()}
+
+        result = await super().get_tools(ctx)
+        if self._server.cache_tools:
+            self._cached_tool_defs = {name: tool.tool_def for name, tool in result.items()}
+        return result

pydantic_ai_slim/pydantic_ai/durable_exec/temporal/_logfire.py

@@ -15,7 +15,7 @@ def _default_setup_logfire() -> Logfire:
     import logfire
 
     instance = logfire.configure()
-    logfire.instrument_pydantic_ai()
+    instance.instrument_pydantic_ai()
     return instance
 
 

pydantic_ai_slim/pydantic_ai/durable_exec/temporal/_mcp_server.py

@@ -6,13 +6,19 @@
 
 from pydantic_ai import ToolsetTool
 from pydantic_ai.mcp import MCPServer
-from pydantic_ai.tools import AgentDepsT, ToolDefinition
+from pydantic_ai.tools import AgentDepsT, RunContext, ToolDefinition
 
 from ._mcp import TemporalMCPToolset
 from ._run_context import TemporalRunContext
 
 
 class TemporalMCPServer(TemporalMCPToolset[AgentDepsT]):
+    """A wrapper for MCPServer that integrates with Temporal, turning get_tools and call_tool into activities.
+
+    Tool definitions are cached across activities to avoid redundant MCP server round-trips,
+    respecting the wrapped server's `cache_tools` setting.
+    """
+
     def __init__(
         self,
         server: MCPServer,
@@ -31,7 +37,24 @@ def __init__(
             deps_type=deps_type,
             run_context_type=run_context_type,
         )
+        # Cached across activities to avoid redundant MCP connections per activity.
+        # Not invalidated by `tools/list_changed` notifications — users who need
+        # dynamic tools during a workflow should set `cache_tools=False`.
+        self._cached_tool_defs: dict[str, ToolDefinition] | None = None
 
-    def tool_for_tool_def(self, tool_def: ToolDefinition) -> ToolsetTool[AgentDepsT]:
+    @property
+    def _server(self) -> MCPServer:
         assert isinstance(self.wrapped, MCPServer)
-        return self.wrapped.tool_for_tool_def(tool_def)
+        return self.wrapped
+
+    def tool_for_tool_def(self, tool_def: ToolDefinition) -> ToolsetTool[AgentDepsT]:
+        return self._server.tool_for_tool_def(tool_def)
+
+    async def get_tools(self, ctx: RunContext[AgentDepsT]) -> dict[str, ToolsetTool[AgentDepsT]]:
+        if self._server.cache_tools and self._cached_tool_defs is not None:
+            return {name: self.tool_for_tool_def(td) for name, td in self._cached_tool_defs.items()}
+
+        result = await super().get_tools(ctx)
+        if self._server.cache_tools:  # pragma: no branch
+            self._cached_tool_defs = {name: tool.tool_def for name, tool in result.items()}
+        return result

pydantic_ai_slim/pydantic_ai/mcp.py

@@ -339,17 +339,22 @@ class MCPServer(AbstractToolset[Any], ABC):
 
     When enabled (default), tools are fetched once and cached until either:
     - The server sends a `notifications/tools/list_changed` notification
-    - The connection is closed
+    - [`MCPServer.__aexit__`][pydantic_ai.mcp.MCPServer.__aexit__] is called (when the last context exits)
 
     Set to `False` for servers that change tools dynamically without sending notifications.
+
+    Note: When using durable execution (Temporal, DBOS), tool definitions are additionally cached
+    at the wrapper level across activities/steps, to avoid redundant MCP connections. This
+    wrapper-level cache is not invalidated by `tools/list_changed` notifications.
+    Set to `False` to disable all caching if tools may change during a workflow.
     """
 
     cache_resources: bool
     """Whether to cache the list of resources.
 
     When enabled (default), resources are fetched once and cached until either:
     - The server sends a `notifications/resources/list_changed` notification
-    - The connection is closed
+    - [`MCPServer.__aexit__`][pydantic_ai.mcp.MCPServer.__aexit__] is called (when the last context exits)
 
     Set to `False` for servers that change resources dynamically without sending notifications.
     """
@@ -479,20 +484,18 @@ async def list_tools(self) -> list[mcp_types.Tool]:
 
         Tools are cached by default, with cache invalidation on:
         - `notifications/tools/list_changed` notifications from the server
-        - Connection close (cache is cleared in `__aexit__`)
+        - `__aexit__` when the last context exits
 
         Set `cache_tools=False` for servers that change tools without sending notifications.
         """
+        if self.cache_tools and self._cached_tools is not None:
+            return self._cached_tools
+
         async with self:
+            result = await self._client.list_tools()
             if self.cache_tools:
-                if self._cached_tools is not None:
-                    return self._cached_tools
-                result = await self._client.list_tools()
                 self._cached_tools = result.tools
-                return result.tools
-            else:
-                result = await self._client.list_tools()
-                return result.tools
+            return result.tools
 
     async def direct_call_tool(
         self,
@@ -600,27 +603,25 @@ async def list_resources(self) -> list[Resource]:
 
         Resources are cached by default, with cache invalidation on:
         - `notifications/resources/list_changed` notifications from the server
-        - Connection close (cache is cleared in `__aexit__`)
+        - `__aexit__` when the last context exits
 
         Set `cache_resources=False` for servers that change resources without sending notifications.
 
         Raises:
             MCPError: If the server returns an error.
         """
+        if self.cache_resources and self._cached_resources is not None:
+            return self._cached_resources
+
         async with self:
             if not self.capabilities.resources:
                 return []
             try:
+                result = await self._client.list_resources()
+                resources = [Resource.from_mcp_sdk(r) for r in result.resources]
                 if self.cache_resources:
-                    if self._cached_resources is not None:
-                        return self._cached_resources
-                    result = await self._client.list_resources()
-                    resources = [Resource.from_mcp_sdk(r) for r in result.resources]
                     self._cached_resources = resources
-                    return resources
-                else:
-                    result = await self._client.list_resources()
-                    return [Resource.from_mcp_sdk(r) for r in result.resources]
+                return resources
             except mcp_exceptions.McpError as e:
                 raise MCPError.from_mcp_sdk(e) from e