add tool and resource caching for mcp servers that support change notifications

Author: dsfaccini

pydantic_ai_slim/pydantic_ai/mcp.py

@@ -232,6 +232,12 @@ class ServerCapabilities:
     tools: bool = False
     """Whether the server offers any tools to call."""
 
+    tools_list_changed: bool = False
+    """Whether the server will emit notifications when the list of tools changes."""
+
+    resources_list_changed: bool = False
+    """Whether the server will emit notifications when the list of resources changes."""
+
     completions: bool = False
     """Whether the server offers autocompletion suggestions for prompts and resources."""
 
@@ -244,12 +250,16 @@ def from_mcp_sdk(cls, mcp_capabilities: mcp_types.ServerCapabilities) -> ServerC
         Args:
             mcp_capabilities: The MCP SDK ServerCapabilities object.
         """
+        tools_cap = mcp_capabilities.tools
+        resources_cap = mcp_capabilities.resources
         return cls(
             experimental=list(mcp_capabilities.experimental.keys()) if mcp_capabilities.experimental else None,
             logging=mcp_capabilities.logging is not None,
             prompts=mcp_capabilities.prompts is not None,
-            resources=mcp_capabilities.resources is not None,
-            tools=mcp_capabilities.tools is not None,
+            resources=resources_cap is not None,
+            tools=tools_cap is not None,
+            tools_list_changed=bool(tools_cap.listChanged) if tools_cap else False,
+            resources_list_changed=bool(resources_cap.listChanged) if resources_cap else False,
             completions=mcp_capabilities.completions is not None,
         )
 
@@ -332,6 +342,11 @@ class MCPServer(AbstractToolset[Any], ABC):
     _server_capabilities: ServerCapabilities
     _instructions: str | None
 
+    _cached_tools: list[mcp_types.Tool] | None
+    _tools_cache_valid: bool
+    _cached_resources: list[Resource] | None
+    _resources_cache_valid: bool
+
     def __init__(
         self,
         tool_prefix: str | None = None,
@@ -366,6 +381,10 @@ def __post_init__(self):
         self._enter_lock = Lock()
         self._running_count = 0
         self._exit_stack = None
+        self._cached_tools = None
+        self._tools_cache_valid = False
+        self._cached_resources = None
+        self._resources_cache_valid = False
 
     @abstractmethod
     @asynccontextmanager
@@ -430,13 +449,23 @@ def instructions(self) -> str | None:
     async def list_tools(self) -> list[mcp_types.Tool]:
         """Retrieve tools that are currently active on the server.
 
-        Note:
-        - We don't cache tools as they might change.
-        - We also don't subscribe to the server to avoid complexity.
+        Tools are cached when the server advertises `tools.listChanged` capability,
+        with cache invalidation on tool change notifications and reconnection.
         """
         async with self:  # Ensure server is running
-            result = await self._client.list_tools()
-        return result.tools
+            # Only cache if server supports listChanged notifications
+            if self._server_capabilities.tools_list_changed:
+                if self._cached_tools is not None and self._tools_cache_valid:
+                    return self._cached_tools
+
+                result = await self._client.list_tools()
+                self._cached_tools = result.tools
+                self._tools_cache_valid = True
+                return result.tools
+            else:
+                # Server doesn't support notifications, always fetch fresh
+                result = await self._client.list_tools()
+                return result.tools
 
     async def direct_call_tool(
         self,
@@ -542,9 +571,8 @@ def tool_for_tool_def(self, tool_def: ToolDefinition) -> ToolsetTool[Any]:
     async def list_resources(self) -> list[Resource]:
         """Retrieve resources that are currently present on the server.
 
-        Note:
-        - We don't cache resources as they might change.
-        - We also don't subscribe to resource changes to avoid complexity.
+        Resources are cached when the server advertises `resources.listChanged` capability,
+        with cache invalidation on resource change notifications and reconnection.
 
         Raises:
             MCPError: If the server returns an error.
@@ -553,10 +581,21 @@ async def list_resources(self) -> list[Resource]:
             if not self.capabilities.resources:
                 return []
             try:
-                result = await self._client.list_resources()
+                # caching logic same as list_tools
+                if self._server_capabilities.resources_list_changed:
+                    if self._cached_resources is not None and self._resources_cache_valid:
+                        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
+                    self._resources_cache_valid = True
+                    return resources
+                else:
+                    result = await self._client.list_resources()
+                    return [Resource.from_mcp_sdk(r) for r in result.resources]
             except mcp_exceptions.McpError as e:
                 raise MCPError.from_mcp_sdk(e) from e
-        return [Resource.from_mcp_sdk(r) for r in result.resources]
 
     async def list_resource_templates(self) -> list[ResourceTemplate]:
         """Retrieve resource templates that are currently present on the server.
@@ -619,6 +658,12 @@ async def __aenter__(self) -> Self:
         """
         async with self._enter_lock:
             if self._running_count == 0:
+                # Invalidate caches on fresh connection
+                self._cached_tools = None
+                self._tools_cache_valid = False
+                self._cached_resources = None
+                self._resources_cache_valid = False
+
                 async with AsyncExitStack() as exit_stack:
                     self._read_stream, self._write_stream = await exit_stack.enter_async_context(self.client_streams())
                     client = ClientSession(
@@ -628,6 +673,7 @@ async def __aenter__(self) -> Self:
                         elicitation_callback=self.elicitation_callback,
                         logging_callback=self.log_handler,
                         read_timeout_seconds=timedelta(seconds=self.read_timeout),
+                        message_handler=self._handle_notification,
                     )
                     self._client = await exit_stack.enter_async_context(client)
 
@@ -680,6 +726,13 @@ async def _sampling_callback(
             model=self.sampling_model.model_name,
         )
 
+    async def _handle_notification(self, message: Any) -> None:
+        """Handle notifications from the MCP server, invalidating caches as needed."""
+        if isinstance(message, mcp_types.ToolListChangedNotification):
+            self._tools_cache_valid = False
+        elif isinstance(message, mcp_types.ResourceListChangedNotification):
+            self._resources_cache_valid = False
+
     async def _map_tool_result_part(
         self, part: mcp_types.ContentBlock
     ) -> str | messages.BinaryContent | dict[str, Any] | list[Any]:

tests/test_mcp.py

@@ -54,7 +54,9 @@
         ElicitRequestParams,
         ElicitResult,
         ImageContent,
+        ResourceListChangedNotification,
         TextContent,
+        ToolListChangedNotification,
     )
 
     from pydantic_ai._mcp import map_from_mcp_params, map_from_model_response
@@ -1987,3 +1989,165 @@ async def test_custom_http_client_not_closed():
     assert len(tools) > 0
 
     assert not custom_http_client.is_closed
+
+
+# ============================================================================
+# Tool and Resource Caching Tests
+# ============================================================================
+
+
+async def test_tools_caching_with_list_changed_capability(mcp_server: MCPServerStdio) -> None:
+    """Test that list_tools() caches results when server supports listChanged notifications."""
+    async with mcp_server:
+        # Mock the server capabilities to indicate listChanged is supported
+        mcp_server._server_capabilities.tools_list_changed = True  # pyright: ignore[reportPrivateUsage]
+
+        # First call - should fetch from server and cache
+        tools1 = await mcp_server.list_tools()
+        assert len(tools1) > 0
+        assert mcp_server._cached_tools is not None  # pyright: ignore[reportPrivateUsage]
+        assert mcp_server._tools_cache_valid is True  # pyright: ignore[reportPrivateUsage]
+
+        # Mock _client.list_tools to track if it's called again
+        original_list_tools = mcp_server._client.list_tools  # pyright: ignore[reportPrivateUsage]
+        call_count = 0
+
+        async def mock_list_tools():  # pragma: no cover
+            nonlocal call_count
+            call_count += 1
+            return await original_list_tools()
+
+        mcp_server._client.list_tools = mock_list_tools  # pyright: ignore[reportPrivateUsage,reportAttributeAccessIssue]
+
+        # Second call - should return cached value without calling server
+        tools2 = await mcp_server.list_tools()
+        assert tools2 == tools1
+        assert call_count == 0  # list_tools should not have been called
+
+
+async def test_tools_no_caching_without_list_changed_capability(mcp_server: MCPServerStdio) -> None:
+    """Test that list_tools() always fetches fresh when server doesn't support listChanged."""
+    async with mcp_server:
+        # Verify the server doesn't advertise listChanged by default
+        # (this depends on the test MCP server implementation)
+        mcp_server._server_capabilities.tools_list_changed = False  # pyright: ignore[reportPrivateUsage]
+
+        # First call
+        tools1 = await mcp_server.list_tools()
+        assert len(tools1) > 0
+
+        # Mock _client.list_tools to track calls
+        original_list_tools = mcp_server._client.list_tools  # pyright: ignore[reportPrivateUsage]
+        call_count = 0
+
+        async def mock_list_tools():
+            nonlocal call_count
+            call_count += 1
+            return await original_list_tools()
+
+        mcp_server._client.list_tools = mock_list_tools  # pyright: ignore[reportPrivateUsage,reportAttributeAccessIssue]
+
+        # Second call - should fetch fresh since no listChanged capability
+        tools2 = await mcp_server.list_tools()
+        assert tools2 == tools1
+        assert call_count == 1  # list_tools should have been called
+
+
+async def test_tools_cache_invalidation_on_notification(mcp_server: MCPServerStdio) -> None:
+    """Test that tools cache is invalidated when ToolListChangedNotification is received."""
+    async with mcp_server:
+        # Enable caching
+        mcp_server._server_capabilities.tools_list_changed = True  # pyright: ignore[reportPrivateUsage]
+
+        # Populate cache
+        await mcp_server.list_tools()
+        assert mcp_server._tools_cache_valid is True  # pyright: ignore[reportPrivateUsage]
+
+        # Simulate receiving a tool list changed notification
+        notification = ToolListChangedNotification()
+        await mcp_server._handle_notification(notification)  # pyright: ignore[reportPrivateUsage]
+
+        # Cache should be invalidated
+        assert mcp_server._tools_cache_valid is False  # pyright: ignore[reportPrivateUsage]
+
+        # Cached tools are still present but marked invalid
+        assert mcp_server._cached_tools is not None  # pyright: ignore[reportPrivateUsage]
+
+
+async def test_resources_caching_with_list_changed_capability(mcp_server: MCPServerStdio) -> None:
+    """Test that list_resources() caches results when server supports listChanged notifications."""
+    async with mcp_server:
+        # Mock the server capabilities to indicate listChanged is supported
+        mcp_server._server_capabilities.resources_list_changed = True  # pyright: ignore[reportPrivateUsage]
+
+        # First call - should fetch from server and cache
+        if mcp_server.capabilities.resources:  # pragma: no branch
+            resources1 = await mcp_server.list_resources()
+            assert mcp_server._cached_resources is not None  # pyright: ignore[reportPrivateUsage]
+            assert mcp_server._resources_cache_valid is True  # pyright: ignore[reportPrivateUsage]
+
+            # Mock _client.list_resources to track if it's called again
+            original_list_resources = mcp_server._client.list_resources  # pyright: ignore[reportPrivateUsage]
+            call_count = 0
+
+            async def mock_list_resources():  # pragma: no cover
+                nonlocal call_count
+                call_count += 1
+                return await original_list_resources()
+
+            mcp_server._client.list_resources = mock_list_resources  # pyright: ignore[reportPrivateUsage,reportAttributeAccessIssue]
+
+            # Second call - should return cached value without calling server
+            resources2 = await mcp_server.list_resources()
+            assert resources2 == resources1
+            assert call_count == 0  # list_resources should not have been called
+
+
+async def test_resources_cache_invalidation_on_notification(mcp_server: MCPServerStdio) -> None:
+    """Test that resources cache is invalidated when ResourceListChangedNotification is received."""
+    async with mcp_server:
+        # Enable caching
+        mcp_server._server_capabilities.resources_list_changed = True  # pyright: ignore[reportPrivateUsage]
+
+        # Populate cache (if server supports resources)
+        if mcp_server.capabilities.resources:  # pragma: no branch
+            await mcp_server.list_resources()
+            assert mcp_server._resources_cache_valid is True  # pyright: ignore[reportPrivateUsage]
+
+            # Simulate receiving a resource list changed notification
+            notification = ResourceListChangedNotification()
+            await mcp_server._handle_notification(notification)  # pyright: ignore[reportPrivateUsage]
+
+            # Cache should be invalidated
+            assert mcp_server._resources_cache_valid is False  # pyright: ignore[reportPrivateUsage]
+
+
+async def test_cache_invalidation_on_reconnection() -> None:
+    """Test that caches are cleared when reconnecting to the server."""
+    server = MCPServerStdio('python', ['-m', 'tests.mcp_server'])
+
+    # First connection
+    async with server:
+        server._server_capabilities.tools_list_changed = True  # pyright: ignore[reportPrivateUsage]
+        await server.list_tools()
+        assert server._cached_tools is not None  # pyright: ignore[reportPrivateUsage]
+        assert server._tools_cache_valid is True  # pyright: ignore[reportPrivateUsage]
+
+    # After exiting, the server is no longer running
+    # but cache state persists until next connection
+
+    # Reconnect
+    async with server:
+        # Cache should be cleared on fresh connection
+        assert server._cached_tools is None  # pyright: ignore[reportPrivateUsage]
+        assert server._tools_cache_valid is False  # pyright: ignore[reportPrivateUsage]
+
+
+async def test_server_capabilities_list_changed_fields() -> None:
+    """Test that ServerCapabilities correctly parses listChanged fields."""
+    server = MCPServerStdio('python', ['-m', 'tests.mcp_server'])
+    async with server:
+        # Test that capabilities are accessible
+        caps = server.capabilities
+        assert isinstance(caps.tools_list_changed, bool)
+        assert isinstance(caps.resources_list_changed, bool)