fastmcp/tests/server/middleware/test_caching.py

"""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.elicitation import ElicitResult
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,
    CacheableToolResult,
    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
from fastmcp.utilities.tasks import TaskConfig

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_list_operations_preserve_component_metadata(self):
        """Base component fields should survive conversion through the cache."""
        from fastmcp.server.extensions import ServerExtension
        from fastmcp.utilities.tasks import TASKS_EXTENSION_ID

        class _StubTasksExtension(ServerExtension):
            identifier = TASKS_EXTENSION_ID

        icon = mcp_types.Icon(src="https://example.com/component.png")
        mcp = FastMCP("MetadataServer")
        mcp.add_middleware(ResponseCachingMiddleware())
        # A task-enabled tool requires the tasks extension to serve; register a
        # stub so the metadata (execution.task_support) can be verified end-to-end.
        mcp.add_extension(_StubTasksExtension())

        @mcp.tool(icons=[icon], task=TaskConfig(mode="optional"))
        async def greet() -> str:
            return "hello"

        @mcp.resource("resource://metadata", icons=[icon])
        def metadata_resource() -> str:
            return "resource"

        @mcp.prompt(icons=[icon])
        def metadata_prompt() -> str:
            return "prompt"

        cached_tools = await mcp.list_tools()
        cached_resources = await mcp.list_resources()
        cached_prompts = await mcp.list_prompts()
        assert type(cached_tools[0]) is Tool
        assert type(cached_resources[0]) is Resource
        assert type(cached_prompts[0]) is Prompt
        assert not hasattr(cached_tools[0], "fn")
        assert not hasattr(cached_resources[0], "fn")
        assert not hasattr(cached_prompts[0], "fn")

        # Pinned to legacy: the tool's `execution.task_support` (SEP-1686) is
        # advertised in the handshake-era tool listing; the modern listing omits it.
        async with Client(mcp, mode="legacy") as client:
            for _ in range(2):
                tools = await client.list_tools()
                resources = await client.list_resources()
                prompts = await client.list_prompts()

                assert tools[0].icons == [icon]
                assert tools[0].execution == mcp_types.ToolExecution(
                    task_support="optional"
                )
                assert resources[0].icons == [icon]
                assert prompts[0].icons == [icon]

    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 TestCacheableToolResult:
    def test_wrap_and_unwrap(self):
        tool_result = ToolResult(
            "unstructured content",
            structured_content={"structured": "content"},
            meta={"meta": "data"},
        )

        cached_tool_result = CacheableToolResult.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 = CacheableToolResult.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:
            with warnings.catch_warnings():
                warnings.simplefilter("error", UserWarning)

                # 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)


class TestCachingWithInputRequiredResults:
    """A multi-round-trip ask must survive `ResponseCachingMiddleware`.

    An ask carries no content of its own, so caching one would store an empty
    result and the client would never see the question. Continuation legs must
    also bypass the cache, since they share a cache key with a fresh call.
    """

    @staticmethod
    def _answer(responses: mcp_types.InputResponses) -> str:
        """The accepted value for the single question these guards ask."""
        result = responses["q"]
        assert isinstance(result, mcp_types.ElicitResult)
        assert result.content is not None
        return str(result.content["q"])

    @staticmethod
    def _ask() -> mcp_types.InputRequiredResult:
        params = mcp_types.ElicitRequestFormParams(
            message="Which quarter?",
            requested_schema={
                "type": "object",
                "properties": {"q": {"type": "string"}},
                "required": ["q"],
            },
        )
        request = mcp_types.ElicitRequest(method="elicitation/create", params=params)
        return mcp_types.InputRequiredResult(
            result_type="input_required", input_requests={"q": request}
        )

    @classmethod
    def _server(cls) -> FastMCP:
        mcp = FastMCP("cached-guards")
        mcp.add_middleware(ResponseCachingMiddleware())

        @mcp.tool
        async def summarize_tool(ctx: Context) -> str | mcp_types.InputRequiredResult:
            if ctx.input_responses is None:
                return cls._ask()
            return f"Summary for {cls._answer(ctx.input_responses)}"

        @mcp.prompt
        async def summarize(ctx: Context) -> str | mcp_types.InputRequiredResult:
            if ctx.input_responses is None:
                return cls._ask()
            return f"Summary for {cls._answer(ctx.input_responses)}"

        @mcp.resource("report://x")
        async def report(ctx: Context) -> str | mcp_types.InputRequiredResult:
            if ctx.input_responses is None:
                return cls._ask()
            return f"Report for {cls._answer(ctx.input_responses)}"

        return mcp

    @staticmethod
    async def _handler(message, response_type, params, ctx):
        return ElicitResult(action="accept", content=response_type(q="Q3"))

    def _client(self) -> Client:
        return Client(self._server(), mode="auto", elicitation_handler=self._handler)

    async def test_tool_guard_completes_under_caching(self):
        async with self._client() as client:
            result = await client.call_tool("summarize_tool", {})
        assert result.data == "Summary for Q3"

    async def test_prompt_guard_completes_under_caching(self):
        async with self._client() as client:
            result = await client.get_prompt("summarize")
        assert result.messages[0].content.text == "Summary for Q3"

    async def test_resource_guard_completes_under_caching(self):
        async with self._client() as client:
            result = await client.read_resource("report://x")
        assert result[0].text == "Report for Q3"

    async def test_guard_still_asks_on_a_second_fresh_call(self):
        """The ask must not be cached away for the next caller.

        A second fresh flow has to be asked the same question; serving it a
        cached final answer would skip the component's own per-round logic.
        """
        async with self._client() as client:
            first = await client.get_prompt("summarize")
            second = await client.get_prompt("summarize")
        assert first.messages[0].content.text == "Summary for Q3"
        assert second.messages[0].content.text == "Summary for Q3"