60e0ffc959
Upgrade checks / Notify on failure (push) Has been cancelled
Upgrade checks / Close issue on success (push) Has been cancelled
Schema Crash Test / Real-world schema crash test (232K schemas) (push) Has been cancelled
Run static analysis / static_analysis (push) Has been cancelled
Tests / Tests: Python 3.10 on ubuntu-latest (push) Has been cancelled
Tests / Tests: Python 3.13 on ubuntu-latest (push) Has been cancelled
Tests / Tests: Python 3.10 on windows-latest (push) Has been cancelled
Tests / Tests with lowest-direct dependencies (push) Has been cancelled
Tests / MCP conformance tests (push) Has been cancelled
Tests / Integration tests (push) Has been cancelled
Tests / Package install smoke (push) Has been cancelled
Upgrade checks / Static analysis (push) Has been cancelled
Upgrade checks / Tests: Python 3.10 on ubuntu-latest (push) Has been cancelled
Upgrade checks / Tests: Python 3.13 on ubuntu-latest (push) Has been cancelled
Upgrade checks / Tests: Python 3.10 on windows-latest (push) Has been cancelled
Upgrade checks / Integration tests (push) Has been cancelled
Update MCPServerConfig Schema / update-config-schema (push) Has been cancelled
Update SDK Documentation / update-sdk-docs (push) Has been cancelled
849 lines
30 KiB
Python
849 lines
30 KiB
Python
"""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)
|