"""Tests for response caching middleware.""" import sys import tempfile import warnings from pathlib import Path from unittest.mock import AsyncMock, MagicMock import mcp.types import pytest from inline_snapshot import snapshot from key_value.aio.stores.filetree import ( FileTreeStore, FileTreeV1CollectionSanitizationStrategy, FileTreeV1KeySanitizationStrategy, ) from key_value.aio.stores.memory import MemoryStore from key_value.aio.wrappers.statistics.wrapper import ( GetStatistics, KVStoreCollectionStatistics, PutStatistics, ) from mcp.server.lowlevel.helper_types import ReadResourceContents from mcp.types import TextContent, TextResourceContents from pydantic import AnyUrl, BaseModel from fastmcp import Context, FastMCP from fastmcp.client.client import CallToolResult, Client from fastmcp.client.transports import FastMCPTransport from fastmcp.prompts.base import Message, Prompt from fastmcp.prompts.function_prompt import FunctionPrompt from fastmcp.resources.base import Resource from fastmcp.server.middleware.caching import ( CachableToolResult, CallToolSettings, ResponseCachingMiddleware, ResponseCachingStatistics, _make_call_tool_cache_key, _make_get_prompt_cache_key, _make_read_resource_cache_key, ) from fastmcp.server.middleware.middleware import CallNext, MiddlewareContext from fastmcp.tools.base import Tool, ToolResult TEST_URI = AnyUrl("https://test_uri") SAMPLE_READ_RESOURCE_CONTENTS = ReadResourceContents( content="test_text", mime_type="text/plain", ) def sample_resource_fn() -> list[ReadResourceContents]: return [SAMPLE_READ_RESOURCE_CONTENTS] def sample_prompt_fn() -> Message: return Message("test_text") SAMPLE_RESOURCE = Resource.from_function( fn=sample_resource_fn, uri=TEST_URI, name="test_resource" ) SAMPLE_PROMPT = Prompt.from_function(fn=sample_prompt_fn, name="test_prompt") SAMPLE_GET_PROMPT_RESULT = mcp.types.GetPromptResult( messages=[Message("test_text").to_mcp_prompt_message()] ) SAMPLE_TOOL = Tool(name="test_tool", parameters={"param1": "value1", "param2": 42}) SAMPLE_TOOL_RESULT = ToolResult( content=[TextContent(type="text", text="test_text")], structured_content={"result": "test_result"}, ) SAMPLE_TOOL_RESULT_LARGE = ToolResult( content=[TextContent(type="text", text="test_text" * 100)], structured_content={"result": "test_result"}, ) class CrazyModel(BaseModel): a: int b: int c: str d: float e: bool f: list[int] g: dict[str, int] h: list[dict[str, int]] i: dict[str, list[int]] @pytest.fixture def crazy_model() -> CrazyModel: return CrazyModel( a=5, b=10, c="test", d=1.0, e=True, f=[1, 2, 3], g={"a": 1, "b": 2}, h=[{"a": 1, "b": 2}], i={"a": [1, 2]}, ) class TrackingCalculator: add_calls: int multiply_calls: int crazy_calls: int very_large_response_calls: int def __init__(self): self.add_calls = 0 self.multiply_calls = 0 self.crazy_calls = 0 self.very_large_response_calls = 0 def add(self, a: int, b: int) -> int: self.add_calls += 1 return a + b def multiply(self, a: int, b: int) -> int: self.multiply_calls += 1 return a * b def very_large_response(self) -> str: self.very_large_response_calls += 1 return "istenchars" * 100000 # 1,000,000 characters, 1mb def crazy(self, a: CrazyModel) -> CrazyModel: self.crazy_calls += 1 return a def how_to_calculate(self, a: int, b: int) -> str: return f"To calculate {a} + {b}, you need to add {a} and {b} together." def get_add_calls(self) -> str: return str(self.add_calls) def get_multiply_calls(self) -> str: return str(self.multiply_calls) def get_crazy_calls(self) -> str: return str(self.crazy_calls) async def update_tool_list(self, context: Context): import mcp.types await context.send_notification(mcp.types.ToolListChangedNotification()) def add_tools(self, fastmcp: FastMCP, prefix: str = ""): _ = fastmcp.add_tool(tool=Tool.from_function(fn=self.add, name=f"{prefix}add")) _ = fastmcp.add_tool( tool=Tool.from_function(fn=self.multiply, name=f"{prefix}multiply") ) _ = fastmcp.add_tool( tool=Tool.from_function(fn=self.crazy, name=f"{prefix}crazy") ) _ = fastmcp.add_tool( tool=Tool.from_function( fn=self.very_large_response, name=f"{prefix}very_large_response" ) ) _ = fastmcp.add_tool( tool=Tool.from_function( fn=self.update_tool_list, name=f"{prefix}update_tool_list" ) ) def add_prompts(self, fastmcp: FastMCP, prefix: str = ""): _ = fastmcp.add_prompt( prompt=FunctionPrompt.from_function( fn=self.how_to_calculate, name=f"{prefix}how_to_calculate" ) ) def add_resources(self, fastmcp: FastMCP, prefix: str = ""): _ = fastmcp.add_resource( resource=Resource.from_function( fn=self.get_add_calls, uri="resource://add_calls", name=f"{prefix}add_calls", ) ) _ = fastmcp.add_resource( resource=Resource.from_function( fn=self.get_multiply_calls, uri="resource://multiply_calls", name=f"{prefix}multiply_calls", ) ) _ = fastmcp.add_resource( resource=Resource.from_function( fn=self.get_crazy_calls, uri="resource://crazy_calls", name=f"{prefix}crazy_calls", ) ) @pytest.fixture def tracking_calculator() -> TrackingCalculator: return TrackingCalculator() @pytest.fixture def mock_context() -> MiddlewareContext[mcp.types.CallToolRequestParams]: """Create a mock middleware context for tool calls.""" context = MagicMock(spec=MiddlewareContext[mcp.types.CallToolRequestParams]) context.message = mcp.types.CallToolRequestParams( name="test_tool", arguments={"param1": "value1", "param2": 42} ) context.method = "tools/call" return context @pytest.fixture def mock_call_next() -> CallNext[mcp.types.CallToolRequestParams, ToolResult]: """Create a mock call_next function.""" return AsyncMock( return_value=ToolResult( content=[TextContent(type="text", text="test result")], structured_content={"result": "success", "value": 123}, ) ) @pytest.fixture def sample_tool_result() -> ToolResult: """Create a sample tool result for testing.""" return ToolResult( content=[TextContent(type="text", text="cached result")], structured_content={"cached": True, "data": "test"}, ) class TestResponseCachingMiddleware: """Test ResponseCachingMiddleware functionality.""" def test_initialization(self): """Test middleware initialization.""" assert ResponseCachingMiddleware( call_tool_settings=CallToolSettings( included_tools=["tool1"], excluded_tools=["tool2"], ), ) @pytest.mark.parametrize( ("tool_name", "included_tools", "excluded_tools", "result"), [ ("tool", ["tool", "tool2"], [], True), ("tool", ["second tool", "third tool"], [], False), ("tool", [], ["tool"], False), ("tool", [], ["second tool"], True), ("tool", ["tool", "second tool"], ["tool"], False), ("tool", ["tool", "second tool"], ["second tool"], True), ], ids=[ "tool is included", "tool is not included", "tool is excluded", "tool is not excluded", "tool is included and excluded (excluded takes precedence)", "tool is included and not excluded", ], ) def test_tool_call_filtering( self, tool_name: str, included_tools: list[str], excluded_tools: list[str], result: bool, ): """Test tool filtering logic.""" middleware1 = ResponseCachingMiddleware( call_tool_settings=CallToolSettings( included_tools=included_tools, excluded_tools=excluded_tools ), ) assert middleware1._matches_tool_cache_settings(tool_name=tool_name) is result @pytest.mark.skipif( sys.platform == "win32", reason="SQLite caching tests are flaky on Windows due to temp directory issues.", ) class TestResponseCachingMiddlewareIntegration: """Integration tests with real FastMCP server.""" @pytest.fixture(params=["memory", "filetree"]) async def caching_server( self, tracking_calculator: TrackingCalculator, request: pytest.FixtureRequest, ): """Create a FastMCP server for caching tests.""" mcp = FastMCP("CachingTestServer", dereference_schemas=False) with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as temp_dir: with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) file_store = FileTreeStore( data_directory=Path(temp_dir), key_sanitization_strategy=FileTreeV1KeySanitizationStrategy( Path(temp_dir) ), collection_sanitization_strategy=FileTreeV1CollectionSanitizationStrategy( Path(temp_dir) ), ) response_caching_middleware = ResponseCachingMiddleware( cache_storage=file_store if request.param == "filetree" else MemoryStore(), ) mcp.add_middleware(middleware=response_caching_middleware) tracking_calculator.add_tools(fastmcp=mcp) tracking_calculator.add_resources(fastmcp=mcp) tracking_calculator.add_prompts(fastmcp=mcp) yield mcp @pytest.fixture def non_caching_server(self, tracking_calculator: TrackingCalculator): """Create a FastMCP server for non-caching tests.""" mcp = FastMCP("NonCachingTestServer") tracking_calculator.add_tools(fastmcp=mcp) return mcp async def test_list_tools( self, caching_server: FastMCP, tracking_calculator: TrackingCalculator ): """Test that tool list caching works with a real FastMCP server.""" async with Client(caching_server) as client: pre_tool_list: list[mcp.types.Tool] = await client.list_tools() assert len(pre_tool_list) == 5 # Add a tool and make sure it's missing from the list tool response _ = caching_server.add_tool( tool=Tool.from_function(fn=tracking_calculator.add, name="add_2") ) post_tool_list: list[mcp.types.Tool] = await client.list_tools() assert len(post_tool_list) == 5 assert pre_tool_list == post_tool_list async def test_call_tool( self, caching_server: FastMCP, tracking_calculator: TrackingCalculator, ): """Test that caching works with a real FastMCP server.""" tracking_calculator.add_tools(fastmcp=caching_server) async with Client[FastMCPTransport](transport=caching_server) as client: call_tool_result_one: CallToolResult = await client.call_tool( "add", {"a": 5, "b": 3} ) assert tracking_calculator.add_calls == 1 call_tool_result_two: CallToolResult = await client.call_tool( "add", {"a": 5, "b": 3} ) assert call_tool_result_one == call_tool_result_two async def test_call_tool_very_large_value( self, caching_server: FastMCP, tracking_calculator: TrackingCalculator, ): """Test that caching works with a real FastMCP server.""" tracking_calculator.add_tools(fastmcp=caching_server) async with Client[FastMCPTransport](transport=caching_server) as client: call_tool_result_one: CallToolResult = await client.call_tool( "very_large_response", {} ) assert tracking_calculator.very_large_response_calls == 1 call_tool_result_two: CallToolResult = await client.call_tool( "very_large_response", {} ) assert call_tool_result_one == call_tool_result_two assert tracking_calculator.very_large_response_calls == 2 async def test_call_tool_crazy_value( self, caching_server: FastMCP, tracking_calculator: TrackingCalculator, crazy_model: CrazyModel, ): """Test that caching works with a real FastMCP server.""" tracking_calculator.add_tools(fastmcp=caching_server) async with Client[FastMCPTransport](transport=caching_server) as client: call_tool_result_one: CallToolResult = await client.call_tool( "crazy", {"a": crazy_model} ) assert tracking_calculator.crazy_calls == 1 call_tool_result_two: CallToolResult = await client.call_tool( "crazy", {"a": crazy_model} ) assert call_tool_result_one == call_tool_result_two assert tracking_calculator.crazy_calls == 1 async def test_list_resources( self, caching_server: FastMCP, tracking_calculator: TrackingCalculator ): """Test that list resources caching works with a real FastMCP server.""" async with Client[FastMCPTransport](transport=caching_server) as client: pre_resource_list: list[mcp.types.Resource] = await client.list_resources() assert len(pre_resource_list) == 3 tracking_calculator.add_resources(fastmcp=caching_server) post_resource_list: list[mcp.types.Resource] = await client.list_resources() assert len(post_resource_list) == 3 assert pre_resource_list == post_resource_list async def test_read_resource( self, caching_server: FastMCP, tracking_calculator: TrackingCalculator ): """Test that get resources caching works with a real FastMCP server.""" async with Client[FastMCPTransport](transport=caching_server) as client: pre_resource = await client.read_resource(uri="resource://add_calls") assert isinstance(pre_resource[0], TextResourceContents) assert pre_resource[0].text == "0" tracking_calculator.add_calls = 1 post_resource = await client.read_resource(uri="resource://add_calls") assert isinstance(post_resource[0], TextResourceContents) assert post_resource[0].text == "0" assert pre_resource == post_resource async def test_list_prompts( self, caching_server: FastMCP, tracking_calculator: TrackingCalculator ): """Test that list prompts caching works with a real FastMCP server.""" async with Client[FastMCPTransport](transport=caching_server) as client: pre_prompt_list: list[mcp.types.Prompt] = await client.list_prompts() assert len(pre_prompt_list) == 1 tracking_calculator.add_prompts(fastmcp=caching_server) post_prompt_list: list[mcp.types.Prompt] = await client.list_prompts() assert len(post_prompt_list) == 1 assert pre_prompt_list == post_prompt_list async def test_get_prompts( self, caching_server: FastMCP, tracking_calculator: TrackingCalculator ): """Test that get prompts caching works with a real FastMCP server.""" async with Client[FastMCPTransport](transport=caching_server) as client: pre_prompt = await client.get_prompt( name="how_to_calculate", arguments={"a": 5, "b": 3} ) pre_prompt_content = pre_prompt.messages[0].content assert isinstance(pre_prompt_content, TextContent) assert ( pre_prompt_content.text == "To calculate 5 + 3, you need to add 5 and 3 together." ) tracking_calculator.add_prompts(fastmcp=caching_server) post_prompt = await client.get_prompt( name="how_to_calculate", arguments={"a": 5, "b": 3} ) assert pre_prompt == post_prompt async def test_statistics( self, caching_server: FastMCP, ): """Test that statistics are collected correctly.""" caching_middleware = caching_server.middleware[0] assert isinstance(caching_middleware, ResponseCachingMiddleware) async with Client[FastMCPTransport](transport=caching_server) as client: statistics = caching_middleware.statistics() assert statistics == snapshot(ResponseCachingStatistics()) _ = await client.call_tool("add", {"a": 5, "b": 3}) statistics = caching_middleware.statistics() assert statistics == snapshot( ResponseCachingStatistics( list_tools=KVStoreCollectionStatistics( get=GetStatistics(count=2, hit=1, miss=1), put=PutStatistics(count=1), ), call_tool=KVStoreCollectionStatistics( get=GetStatistics(count=1, miss=1), put=PutStatistics(count=1) ), ) ) _ = await client.call_tool("add", {"a": 5, "b": 3}) statistics = caching_middleware.statistics() assert statistics == snapshot( ResponseCachingStatistics( list_tools=KVStoreCollectionStatistics( get=GetStatistics(count=2, hit=1, miss=1), put=PutStatistics(count=1), ), call_tool=KVStoreCollectionStatistics( get=GetStatistics(count=2, hit=1, miss=1), put=PutStatistics(count=1), ), ) ) class TestCachableToolResult: def test_wrap_and_unwrap(self): tool_result = ToolResult( "unstructured content", structured_content={"structured": "content"}, meta={"meta": "data"}, ) cached_tool_result = CachableToolResult.wrap(tool_result).unwrap() assert cached_tool_result.content == tool_result.content assert cached_tool_result.structured_content == tool_result.structured_content assert cached_tool_result.meta == tool_result.meta class TestCachingWithImportedServerPrefixes: """Test that caching preserves prefixes from imported servers. Regression tests for issue #2300: ResponseCachingMiddleware was losing prefix information when caching components from imported servers. """ @pytest.fixture async def parent_with_imported_child(self, tracking_calculator: TrackingCalculator): """Create a parent server with an imported child server (prefixed).""" child = FastMCP("child") tracking_calculator.add_tools(fastmcp=child) tracking_calculator.add_resources(fastmcp=child) tracking_calculator.add_prompts(fastmcp=child) parent = FastMCP("parent") parent.add_middleware(ResponseCachingMiddleware()) parent.mount(child, namespace="child") return parent async def test_tool_prefixes_preserved_after_cache_hit( self, parent_with_imported_child: FastMCP ): """Tool names should retain prefix after being served from cache.""" async with Client(parent_with_imported_child) as client: # First call populates cache tools_first = await client.list_tools() tool_names_first = [t.name for t in tools_first] # Second call should come from cache tools_cached = await client.list_tools() tool_names_cached = [t.name for t in tools_cached] # All tools should have prefix in both calls assert all(name.startswith("child_") for name in tool_names_first) assert all(name.startswith("child_") for name in tool_names_cached) assert tool_names_first == tool_names_cached async def test_resource_prefixes_preserved_after_cache_hit( self, parent_with_imported_child: FastMCP ): """Resource URIs should retain prefix after being served from cache.""" async with Client(parent_with_imported_child) as client: # First call populates cache resources_first = await client.list_resources() resource_uris_first = [str(r.uri) for r in resources_first] # Second call should come from cache resources_cached = await client.list_resources() resource_uris_cached = [str(r.uri) for r in resources_cached] # All resources should have prefix in URI path in both calls # Resources get path-style prefix: resource://child/path assert all("://child/" in uri for uri in resource_uris_first) assert all("://child/" in uri for uri in resource_uris_cached) assert resource_uris_first == resource_uris_cached async def test_prompt_prefixes_preserved_after_cache_hit( self, parent_with_imported_child: FastMCP ): """Prompt names should retain prefix after being served from cache.""" async with Client(parent_with_imported_child) as client: # First call populates cache prompts_first = await client.list_prompts() prompt_names_first = [p.name for p in prompts_first] # Second call should come from cache prompts_cached = await client.list_prompts() prompt_names_cached = [p.name for p in prompts_cached] # All prompts should have prefix in both calls assert all(name.startswith("child_") for name in prompt_names_first) assert all(name.startswith("child_") for name in prompt_names_cached) assert prompt_names_first == prompt_names_cached async def test_prefixed_tool_callable_after_cache_hit( self, parent_with_imported_child: FastMCP, tracking_calculator: TrackingCalculator, ): """Prefixed tools should be callable after cache populates.""" async with Client(parent_with_imported_child) as client: # Trigger cache population await client.list_tools() await client.list_tools() # From cache # Tool should be callable with prefixed name result = await client.call_tool("child_add", {"a": 5, "b": 3}) assert not result.is_error assert tracking_calculator.add_calls == 1 class TestCacheKeyGeneration: def test_call_tool_key_is_hashed_and_does_not_include_raw_input(self): msg = mcp.types.CallToolRequestParams( name="toolX", arguments={"password": "secret", "path": "../../etc/passwd"}, ) key = _make_call_tool_cache_key(msg) assert len(key) == 64 assert "secret" not in key assert "../../etc/passwd" not in key def test_read_resource_key_is_hashed_and_does_not_include_raw_uri(self): msg = mcp.types.ReadResourceRequestParams( uri=AnyUrl("file:///tmp/../../etc/shadow?token=abcd") ) key = _make_read_resource_cache_key(msg) assert len(key) == 64 assert "shadow" not in key assert "token=abcd" not in key def test_get_prompt_key_is_hashed_and_stable(self): msg = mcp.types.GetPromptRequestParams( name="promptY", arguments={"api_key": "ABC123", "scope": "admin"}, ) key = _make_get_prompt_cache_key(msg) assert len(key) == 64 assert "ABC123" not in key assert key == _make_get_prompt_cache_key(msg)