"""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 ( ANONYMOUS_AUTH_KEY, 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=1, hit=0, 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=1, hit=0, 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 def test_wrap_and_unwrap_preserves_is_error(self): tool_result = ToolResult("boom", is_error=True) cached_tool_result = CachableToolResult.wrap(tool_result).unwrap() assert cached_tool_result.is_error is True 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="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) def test_call_tool_key_partitions_by_auth(self): msg = mcp_types.CallToolRequestParams(name="t", arguments={"a": 1}) anon = _make_call_tool_cache_key(msg) user_a = _make_call_tool_cache_key(msg, auth_key="user_a") user_b = _make_call_tool_cache_key(msg, auth_key="user_b") assert anon == _make_call_tool_cache_key(msg, auth_key=ANONYMOUS_AUTH_KEY) assert user_a != user_b assert user_a != anon def test_read_resource_key_partitions_by_auth(self): msg = mcp_types.ReadResourceRequestParams(uri="file:///tmp/x") user_a = _make_read_resource_cache_key(msg, auth_key="user_a") user_b = _make_read_resource_cache_key(msg, auth_key="user_b") assert user_a != user_b def test_get_prompt_key_partitions_by_auth(self): msg = mcp_types.GetPromptRequestParams(name="p", arguments={"a": "1"}) user_a = _make_get_prompt_cache_key(msg, auth_key="user_a") user_b = _make_get_prompt_cache_key(msg, auth_key="user_b") assert user_a != user_b class TestAuthAwareCaching: """Cached responses must not leak across users with different auth tokens. Regression tests for issue #4037: ResponseCachingMiddleware was caching list/call responses with a global key, so a list filtered by per-component auth checks for one user was served back to other users. """ @staticmethod def _make_token(scopes: list[str]): from fastmcp.server.auth import AccessToken return AccessToken( token=f"token-{'-'.join(scopes) or 'none'}", client_id="test-client", scopes=scopes, expires_at=None, claims={}, ) @staticmethod def _set_token(token): from mcp.server.auth.middleware.auth_context import auth_context_var from mcp.server.auth.middleware.bearer_auth import AuthenticatedUser if token is None: return auth_context_var.set(None) return auth_context_var.set(AuthenticatedUser(token)) async def test_list_tools_cache_does_not_leak_across_tokens(self): from fastmcp.server.auth import require_scopes mcp_server = FastMCP("test") mcp_server.add_middleware(ResponseCachingMiddleware()) @mcp_server.tool(auth=require_scopes("read")) def reader() -> str: return "ok" @mcp_server.tool(auth=require_scopes("read", "write")) def writer() -> str: return "ok" from mcp.server.auth.middleware.auth_context import auth_context_var # Privileged user lists tools first - both visible, gets cached. privileged = self._make_token(["read", "write"]) tok = self._set_token(privileged) try: tools = await mcp_server.list_tools() names = {t.name for t in tools} assert names == {"reader", "writer"} finally: auth_context_var.reset(tok) # Lower-privileged user must not see the cached privileged list. limited = self._make_token(["read"]) tok = self._set_token(limited) try: tools = await mcp_server.list_tools() names = {t.name for t in tools} assert names == {"reader"} finally: auth_context_var.reset(tok) # And same-token repeats still hit the cache (sanity check). tok = self._set_token(limited) try: tools = await mcp_server.list_tools() assert {t.name for t in tools} == {"reader"} finally: auth_context_var.reset(tok) async def test_list_resources_cache_does_not_leak_across_tokens(self): from fastmcp.server.auth import require_scopes mcp_server = FastMCP("test") mcp_server.add_middleware(ResponseCachingMiddleware()) @mcp_server.resource("data://public", auth=require_scopes("read")) def public() -> str: return "public" @mcp_server.resource("data://secret", auth=require_scopes("read", "admin")) def secret() -> str: return "secret" from mcp.server.auth.middleware.auth_context import auth_context_var privileged = self._make_token(["read", "admin"]) tok = self._set_token(privileged) try: resources = await mcp_server.list_resources() uris = {str(r.uri) for r in resources} assert uris == {"data://public", "data://secret"} finally: auth_context_var.reset(tok) limited = self._make_token(["read"]) tok = self._set_token(limited) try: resources = await mcp_server.list_resources() uris = {str(r.uri) for r in resources} assert uris == {"data://public"} finally: auth_context_var.reset(tok) async def test_list_prompts_cache_does_not_leak_across_tokens(self): from fastmcp.server.auth import require_scopes mcp_server = FastMCP("test") mcp_server.add_middleware(ResponseCachingMiddleware()) @mcp_server.prompt(auth=require_scopes("read")) def public_prompt() -> str: return "public" @mcp_server.prompt(auth=require_scopes("read", "admin")) def admin_prompt() -> str: return "admin" from mcp.server.auth.middleware.auth_context import auth_context_var privileged = self._make_token(["read", "admin"]) tok = self._set_token(privileged) try: prompts = await mcp_server.list_prompts() assert {p.name for p in prompts} == {"public_prompt", "admin_prompt"} finally: auth_context_var.reset(tok) limited = self._make_token(["read"]) tok = self._set_token(limited) try: prompts = await mcp_server.list_prompts() assert {p.name for p in prompts} == {"public_prompt"} finally: auth_context_var.reset(tok)