diff --git a/pyproject.toml b/pyproject.toml index 89104173f..c31b62e2a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -15,7 +15,6 @@ dependencies = [ "pydantic[email]>=2.11.7", "pyperclip>=1.9.0", "openapi-core>=0.19.5", - "diskcache>=5.6.3", ] requires-python = ">=3.10" @@ -44,11 +43,13 @@ classifiers = [ [project.optional-dependencies] websockets = ["websockets>=15.0.1"] openai = ["openai>=1.102.0"] +caching = ["diskcache>=5.6.3", "cachetools>=6.2.0"] [dependency-groups] dev = [ "dirty-equals>=0.9.0", "fastmcp[openai]", + "fastmcp[caching]", # add optional dependencies for fastmcp dev "fastapi>=0.115.12", "inline-snapshot[dirty-equals]>=0.27.2", @@ -69,6 +70,7 @@ dev = [ "pytest-xdist>=3.6.1", "ruff", "ty>=0.0.1a19", + "pytest-benchmark>=5.1.0", ] [project.scripts] diff --git a/src/fastmcp/server/middleware/caching.py b/src/fastmcp/server/middleware/caching.py index 5304b9058..0b30fcdcf 100644 --- a/src/fastmcp/server/middleware/caching.py +++ b/src/fastmcp/server/middleware/caching.py @@ -1,53 +1,106 @@ import hashlib import json +from collections import defaultdict from datetime import datetime, timedelta, timezone -from typing import Any, ClassVar, Protocol +from typing import Any, ClassVar, Generic, Protocol, TypedDict, TypeVar, cast -from diskcache import Cache as DiskCacheClient -from mcp.types import CallToolRequestParams, ContentBlock -from pydantic import BaseModel, ConfigDict -from typing_extensions import Self, overload, runtime_checkable +import mcp.types +from mcp.server.lowlevel.helper_types import ReadResourceContents +from pydantic import BaseModel, ConfigDict, Field +from typing_extensions import NotRequired, Self, overload, runtime_checkable +from fastmcp.prompts.prompt import Prompt +from fastmcp.resources.resource import Resource from fastmcp.server.middleware.middleware import CallNext, Middleware, MiddlewareContext -from fastmcp.tools.tool import ToolResult +from fastmcp.tools.tool import Tool, ToolResult +from fastmcp.utilities.logging import get_logger +try: + from cachetools import TLRUCache as MemoryCacheClient + from diskcache import Cache as DiskCacheClient +except ImportError: + raise ImportError( + "fastmcp[caching] is required to use the caching middleware. Please install it with `pip install fastmcp[caching] or `uv add fastmcp[caching]`" + ) + +logger = get_logger(__name__) + +# Constants ONE_HOUR_IN_SECONDS = 3600 +FIVE_MINUTES_IN_SECONDS = 300 + ONE_GB_IN_BYTES = 1024 * 1024 * 1024 +ONE_MB_IN_BYTES = 1024 * 1024 + +GLOBAL_KEY = "__global__" -class CacheEntry(BaseModel): +CachableTypes = ( + ToolResult + | list[Tool] + | list[Resource] + | list[Prompt] + | list[ReadResourceContents] + | mcp.types.GetPromptResult +) + +CachableTypeVar = TypeVar("CachableTypeVar", bound=CachableTypes) + + +def make_collection_key(collection: str, key: str) -> str: + return f"{collection}:{key}" + + +class CachedPrompt(Prompt): + """A cached prompt.""" + + def render( + self, arguments: dict[str, Any] | None = None + ) -> list[mcp.types.PromptMessage]: + raise NotImplementedError( + "Render called on CachedPrompt, this should never happen" + ) + + +class CachedResource(Resource): + """A cached resource.""" + + def read(self) -> str | bytes: + raise NotImplementedError( + "Read called on CachedResource, this should never happen" + ) + + +class CacheEntry(BaseModel, Generic[CachableTypeVar]): """A cache entry.""" - model_config: ClassVar[ConfigDict] = ConfigDict(frozen=True) + model_config: ClassVar[ConfigDict] = ConfigDict( + frozen=True, arbitrary_types_allowed=True + ) key: str - - content: list[ContentBlock] | Any | None - structured_content: str | None + collection: str + value: CachableTypeVar created_at: datetime + + ttl: int + expires_at: datetime def is_expired(self) -> bool: return datetime.now(tz=timezone.utc) > self.expires_at - def to_tool_result(self) -> ToolResult: - return ToolResult( - content=self.content, - structured_content=json.loads(self.structured_content) - if self.structured_content is not None - else None, - ) - @classmethod - def from_tool_result(cls, key: str, value: ToolResult, ttl: int) -> Self: + def from_value( + cls, collection: str, key: str, value: CachableTypeVar, ttl: int + ) -> Self: return cls( + collection=collection, key=key, - content=value.content, - structured_content=json.dumps(value.structured_content) - if value.structured_content is not None - else None, + value=value, created_at=datetime.now(tz=timezone.utc), + ttl=ttl, expires_at=datetime.now(tz=timezone.utc) + timedelta(seconds=ttl), ) @@ -56,79 +109,155 @@ class CacheEntry(BaseModel): class CacheProtocol(Protocol): """A protocol for a cache client.""" - async def get(self, key: str) -> ToolResult | None: + async def get_entry( + self, + collection: str, + key: str, + ) -> CacheEntry[CachableTypes] | None: + """Get a cache entry from the cache.""" + + async def get_value( + self, + collection: str, + key: str, + ) -> CachableTypes | None: """Get a value from the cache.""" - async def set(self, key: str, value: ToolResult, ttl: int) -> None: + if not (cache_entry := await self.get_entry(collection=collection, key=key)): + return None + + return cache_entry.value + + async def set_entry( + self, + cache_entry: CacheEntry[CachableTypes], + ) -> None: """Set a value in the cache.""" - async def delete(self, key: str) -> None: + async def set_value( + self, + collection: str, + key: str, + value: CachableTypes, + ttl: int, + ) -> None: + """Set a value in the cache.""" + + await self.set_entry( + cache_entry=CacheEntry.from_value( + collection=collection, key=key, value=value, ttl=ttl + ) + ) + + async def delete( + self, + collection: str, + key: str, + ) -> None: """Delete a value from the cache.""" + def make_collection_key(self, collection: str, key: str) -> str: + return f"{collection}:{key}" + class DiskCache(CacheProtocol): """A caching client that uses the DiskCache library to cache to disk.""" @overload - def __init__(self, disk_cache: DiskCacheClient): + def __init__(self, *, disk_cache: DiskCacheClient): """Initialize the disk cache with a diskcache client.""" @overload - def __init__(self, path: str, size_limit: int = ONE_GB_IN_BYTES): + def __init__(self, path: str, *, size_limit: int = ONE_GB_IN_BYTES): """Initialize a 1GB disk cache at the provided path.""" def __init__( self, - disk_cache: DiskCacheClient | None = None, path: str | None = None, + *, + disk_cache: DiskCacheClient | None = None, size_limit: int = ONE_GB_IN_BYTES, ): self._cache = disk_cache or DiskCacheClient( directory=path, size_limit=size_limit ) - async def get(self, key: str) -> ToolResult | None: - return self._cache.get(key) + async def get_entry( + self, collection: str, key: str + ) -> CacheEntry[CachableTypes] | None: + collection_key = self.make_collection_key(collection=collection, key=key) - async def set(self, key: str, value: ToolResult, ttl: int) -> None: - self._cache.set(key, value, expire=ttl) + cache_entry = self._cache.get(key=collection_key) - async def delete(self, key: str) -> None: - self._cache.delete(key) + if cache_entry is None: + return None + + return cache_entry # pyright: ignore[reportReturnType] + + async def set_entry( + self, + cache_entry: CacheEntry[CachableTypes], + ) -> None: + collection_key = self.make_collection_key( + collection=cache_entry.collection, key=cache_entry.key + ) + + self._cache.set(key=collection_key, value=cache_entry, expire=cache_entry.ttl) + + async def delete(self, collection: str, key: str) -> None: + collection_key = self.make_collection_key(collection=collection, key=key) + + self._cache.delete(key=collection_key) + + +DEFAULT_MEMORY_CACHE_MAX_ENTRIES = 1000 + + +def _memory_cache_ttu(_key: Any, value: CacheEntry[CachableTypes], now: float) -> float: + return now + value.ttl + + +def _memory_cache_getsizeof(value: CacheEntry[CachableTypes]) -> int: + return 1 class InMemoryCache(CacheProtocol): """A simple in-memory cache.""" - def __init__(self, max_size: int = 1000): - self._cache: dict[str, CacheEntry] = {} - self._max_size = max_size + def __init__(self, max_entries: int = DEFAULT_MEMORY_CACHE_MAX_ENTRIES): + """Initialize the in-memory cache. - async def get(self, key: str) -> ToolResult | None: - cached_entry = self._cache.get(key) - - if cached_entry is None: - return None - - if cached_entry.is_expired(): - self._cache.pop(key, None) - return None - - return ToolResult( - content=cached_entry.content, - structured_content=json.loads(cached_entry.structured_content) - if cached_entry.structured_content is not None - else None, + Args: + max_entries: The maximum number of entries to store in the cache. Defaults to 1000. + """ + self.max_entries = max_entries + self._cache = MemoryCacheClient( + maxsize=max_entries, + ttu=_memory_cache_ttu, + getsizeof=_memory_cache_getsizeof, ) - async def set(self, key: str, value: Any, ttl: int) -> None: - if len(self._cache) >= self._max_size: - self._cache.pop(next(iter(self._cache))) + async def get_entry( + self, collection: str, key: str + ) -> CacheEntry[CachableTypes] | None: + collection_key = self.make_collection_key(collection=collection, key=key) - self._cache[key] = CacheEntry.from_tool_result(key=key, value=value, ttl=ttl) + return self._cache.get(collection_key) - async def delete(self, key: str) -> None: - self._cache.pop(key, None) + async def set_entry( + self, + cache_entry: CacheEntry[CachableTypes], + ) -> None: + collection_key = self.make_collection_key( + collection=cache_entry.collection, key=cache_entry.key + ) + + self._cache[collection_key] = cache_entry + + async def delete(self, collection: str, key: str) -> None: + collection_key = self.make_collection_key(collection=collection, key=key) + + self._cache.pop(collection_key, None) async def setup(self) -> None: return None @@ -137,11 +266,97 @@ class InMemoryCache(CacheProtocol): self._cache.clear() +class CacheMethodStats(BaseModel): + """Stats for a cache method.""" + + hits: int = Field(default=0) + misses: int = Field(default=0) + too_big: int = Field(default=0) + + class CacheStats(BaseModel): """Stats for the cache.""" - hits: int - misses: int + collections: dict[str, CacheMethodStats] = Field( + default_factory=lambda: defaultdict[str, CacheMethodStats](CacheMethodStats) + ) + + def get_misses(self, collection: str) -> int: + return self.collections[collection].misses + + def get_hits(self, collection: str) -> int: + return self.collections[collection].hits + + def get_too_big(self, collection: str) -> int: + return self.collections[collection].too_big + + def mark_miss(self, collection: str) -> None: + self.collections[collection].misses += 1 + + def mark_hit(self, collection: str) -> None: + self.collections[collection].hits += 1 + + def mark_too_big(self, collection: str) -> None: + self.collections[collection].too_big += 1 + + +class SharedMethodSettings(TypedDict): + """Shared config for a cache method.""" + + ttl: NotRequired[int] + + +class CallToolSettings(SharedMethodSettings): + """Extra configuration options for Tool-related caching.""" + + included_tools: NotRequired[list[str]] + excluded_tools: NotRequired[list[str]] + + +class MethodSettings(TypedDict): + """Config for the response caching middleware methods.""" + + list_tools: NotRequired[SharedMethodSettings] + call_tool: NotRequired[CallToolSettings] + + list_resources: NotRequired[SharedMethodSettings] + read_resource: NotRequired[SharedMethodSettings] + + list_prompts: NotRequired[SharedMethodSettings] + get_prompt: NotRequired[SharedMethodSettings] + + +MethodSettingsType = TypeVar("MethodSettingsType", bound=SharedMethodSettings) + +MCP_METHOD_TO_METHOD_SETTINGS_KEY = { + "tools/list": "list_tools", + "tools/call": "call_tool", + "resources/list": "list_resources", + "resources/read": "read_resource", + "prompts/list": "list_prompts", + "prompts/get": "get_prompt", +} + +DEFAULT_METHOD_SETTINGS: MethodSettings = MethodSettings( + list_tools=SharedMethodSettings( + ttl=FIVE_MINUTES_IN_SECONDS, + ), + call_tool=CallToolSettings( + ttl=ONE_HOUR_IN_SECONDS, + ), + list_resources=SharedMethodSettings( + ttl=FIVE_MINUTES_IN_SECONDS, + ), + list_prompts=SharedMethodSettings( + ttl=FIVE_MINUTES_IN_SECONDS, + ), + read_resource=SharedMethodSettings( + ttl=ONE_HOUR_IN_SECONDS, + ), + get_prompt=SharedMethodSettings( + ttl=ONE_HOUR_IN_SECONDS, + ), +) class ResponseCachingMiddleware(Middleware): @@ -152,68 +367,339 @@ class ResponseCachingMiddleware(Middleware): - Cache key derived from tool name and arguments. """ - _stats: CacheStats - def __init__( self, cache_backend: CacheProtocol | None = None, - included_tools: list[str] | None = None, - excluded_tools: list[str] | None = None, + method_settings: MethodSettings | None = None, default_ttl: int = ONE_HOUR_IN_SECONDS, + max_item_size: int | None = None, ): """Initialize the response caching middleware. Args: cache_backend: The cache backend to use. If None, an in-memory cache is used. - included_tools: The tools to cache responses from. If None, all tools are cached. - excluded_tools: The tools to not cache responses from. If None, no tools are excluded. + method_settings: The settings for the middleware. If None, the default settings are used. default_ttl: The default TTL for cached responses. Defaults to one hour. + max_item_size: The maximum size of an item to cache. Defaults to no size limit. """ - self._default_ttl = default_ttl - self._backend = cache_backend or InMemoryCache() + self._default_ttl: int = default_ttl + self._backend: CacheProtocol = cache_backend or InMemoryCache() + self._max_item_size: int | None = max_item_size - self._stats = CacheStats(hits=0, misses=0) + self._stats = CacheStats() - self._included_tools = included_tools - self._excluded_tools = excluded_tools + self.method_settings: MethodSettings = ( + method_settings or DEFAULT_METHOD_SETTINGS + ) + + async def on_list_tools( + self, + context: MiddlewareContext[mcp.types.ListToolsRequest], + call_next: CallNext[mcp.types.ListToolsRequest, list[Tool]], + ) -> list[Tool]: + if self._should_bypass_caching(context=context): + return await call_next(context=context) + + if cached_value := await self._get_cache( + context=context, + call_next=call_next, + key=None, + ): + return cached_value + + result: list[Tool] = await call_next(context) + + # Convert tool subclasses to Tool objects + result = [ + Tool( + name=tool.name, + title=tool.title, + description=tool.description, + parameters=tool.parameters, + output_schema=tool.output_schema, + annotations=tool.annotations, + meta=tool.meta, + tags=tool.tags, + ) + for tool in result + ] + + return await self._store_in_cache_and_return( + context=context, + key=None, + value=result, + ) + + async def on_list_resources( + self, + context: MiddlewareContext[mcp.types.ListResourcesRequest], + call_next: CallNext[mcp.types.ListResourcesRequest, list[Resource]], + ) -> list[Resource]: + if self._should_bypass_caching(context=context): + return await call_next(context) + + if cached_value := await self._get_cache( + context=context, + call_next=call_next, + key=None, + ): + return cached_value + + result: list[Resource] = await call_next(context) + + result = [ + CachedResource( + **resource.model_dump(exclude={"fn"}), + ) + for resource in result + ] + + return await self._store_in_cache_and_return( + context=context, + key=None, + value=result, + ) + + async def on_list_prompts( + self, + context: MiddlewareContext[mcp.types.ListPromptsRequest], + call_next: CallNext[mcp.types.ListPromptsRequest, list[Prompt]], + ) -> list[Prompt]: + if self._should_bypass_caching(context=context): + return await call_next(context) + + if cached_value := await self._get_cache( + context=context, + call_next=call_next, + key=None, + ): + return cached_value + + result: list[Prompt] = await call_next(context) + + result = [ + CachedPrompt( + name=prompt.name, + title=prompt.title, + description=prompt.description, + arguments=prompt.arguments, + meta=prompt.meta, + ) + for prompt in result + ] + + return await self._store_in_cache_and_return( + context=context, + key=None, + value=result, + ) async def on_call_tool( self, - context: MiddlewareContext[CallToolRequestParams], - call_next: CallNext[CallToolRequestParams, Any], + context: MiddlewareContext[mcp.types.CallToolRequestParams], + call_next: CallNext[mcp.types.CallToolRequestParams, ToolResult], ) -> Any: - if not self._should_cache_tool(context.message.name): + if self._should_bypass_caching(context=context): + return await call_next(context=context) + + if not self._matches_tool_cache_settings(context=context): + return await call_next(context=context) + + return await self._cached_call_next( + context=context, + call_next=call_next, + key=self._make_cache_key(msg=context.message), + ) + + async def on_read_resource( + self, + context: MiddlewareContext[mcp.types.ReadResourceRequestParams], + call_next: CallNext[ + mcp.types.ReadResourceRequestParams, list[ReadResourceContents] + ], + ) -> list[ReadResourceContents]: + if self._should_bypass_caching(context=context): + return await call_next(context=context) + + return await self._cached_call_next( + context=context, + call_next=call_next, + ) + + async def on_get_prompt( + self, + context: MiddlewareContext[mcp.types.GetPromptRequestParams], + call_next: CallNext[ + mcp.types.GetPromptRequestParams, mcp.types.GetPromptResult + ], + ) -> mcp.types.GetPromptResult: + if self._should_bypass_caching(context=context): return await call_next(context) - key = self._make_cache_key(context.message) + return await self._cached_call_next( + context=context, + call_next=call_next, + key=None, + ) - if cached_entry := await self._backend.get(key): - self._stats.hits += 1 - return cached_entry + async def on_notification( + self, + context: MiddlewareContext[mcp.types.Notification], + call_next: CallNext[mcp.types.Notification, Any], + ) -> Any: + if isinstance(context.message, mcp.types.ToolListChangedNotification): + await self._backend.delete(collection="tools/list", key=GLOBAL_KEY) - # Cache miss: call downstream - self._stats.misses += 1 - result = await call_next(context) + return await call_next(context) - await self._backend.set(key, result, self._default_ttl) + async def _cached_call_next( + self, + context: MiddlewareContext[Any], + call_next: CallNext[Any, CachableTypeVar], + key: str | None = None, + ) -> CachableTypeVar: + if key is None: + key = GLOBAL_KEY - return result + if cached_value := await self._get_cache( + context=context, + call_next=call_next, + key=key, + ): + return cached_value + + result: CachableTypeVar = await call_next(context) + + return await self._store_in_cache_and_return( + context=context, + key=key, + value=result, + ) + + async def _get_cache( + self, + context: MiddlewareContext[Any], + call_next: CallNext[Any, CachableTypeVar], + key: str | None = None, + ) -> CachableTypeVar | None: + if key is None: + key = GLOBAL_KEY + + if not (collection := context.method): + logger.warning("No method found on context, skipping cache") + return None + + if cached_value := await self._backend.get_value( + collection=collection, key=key + ): + self._stats.mark_hit(collection=collection) + return cast(CachableTypeVar, cached_value) + + self._stats.mark_miss(collection=collection) + + return None + + async def _store_in_cache_and_return( + self, + context: MiddlewareContext[Any], + key: str | None, + value: CachableTypeVar, + ) -> CachableTypeVar: + if key is None: + key = GLOBAL_KEY + + if not (collection := context.method): + logger.warning("No method found on context, skipping cache") + return value + + if self._max_item_size is not None: + size = 0 + + for item in dump_if_base_model(value): + size += len(item.encode("utf-8")) + + if size > self._max_item_size: + self._stats.mark_too_big(collection=collection) + return value + + ttl: int = self._get_cache_ttl(context=context) + + await self._backend.set_value( + collection=collection, + key=key, + value=value, + ttl=ttl, + ) + + return value + + def _matches_tool_cache_settings( + self, context: MiddlewareContext[mcp.types.CallToolRequestParams] + ) -> bool: + tool_name = context.message.name + + tool_call_cache_settings: CallToolSettings | None = self._get_cache_settings( + context=context, + settings_type=CallToolSettings, + ) + + if not tool_call_cache_settings: + return True + + if included_tools := tool_call_cache_settings.get("included_tools"): + if tool_name not in included_tools: + return False + + if excluded_tools := tool_call_cache_settings.get("excluded_tools"): + if tool_name in excluded_tools: + return False - def _should_cache_tool(self, tool_name: str) -> bool: - if self._excluded_tools is not None and tool_name in self._excluded_tools: - return False - if self._included_tools is not None and tool_name not in self._included_tools: - return False return True - def _make_cache_key(self, msg: CallToolRequestParams) -> str: + def _get_cache_settings( + self, + context: MiddlewareContext[Any], + settings_type: type[MethodSettingsType] = SharedMethodSettings, + ) -> MethodSettingsType | None: + if not context.method: + return None + + method_settings_key = MCP_METHOD_TO_METHOD_SETTINGS_KEY.get( + context.method, None + ) + + if ( + method_settings_key is None + or method_settings_key not in self.method_settings + ): + return None + + return cast(MethodSettingsType, self.method_settings[method_settings_key]) + + def _get_cache_ttl(self, context: MiddlewareContext[Any]) -> int: + settings: SharedMethodSettings | None = self._get_cache_settings( + context=context + ) + + if not settings or "ttl" not in settings: + return self._default_ttl + + return settings["ttl"] + + def _should_bypass_caching(self, context: MiddlewareContext[Any]) -> bool: + if not self._get_cache_settings(context=context): + return True + + return False + + def _make_cache_key(self, msg: mcp.types.CallToolRequestParams) -> str: raw = f"{self._get_tool_key(msg)}:{self._get_tool_arguments_str(msg)}" return hashlib.sha256(raw.encode("utf-8")).hexdigest() - def _get_tool_key(self, msg: CallToolRequestParams) -> str: + def _get_tool_key(self, msg: mcp.types.CallToolRequestParams) -> str: return msg.name - def _get_tool_arguments_str(self, msg: CallToolRequestParams) -> str: + def _get_tool_arguments_str(self, msg: mcp.types.CallToolRequestParams) -> str: if msg.arguments is None: return "null" @@ -222,3 +708,17 @@ class ResponseCachingMiddleware(Middleware): except TypeError: return repr(msg.arguments) + + +def dump_if_base_model(value: Any) -> list[str]: + if isinstance(value, BaseModel): + return [value.model_dump_json()] + + if isinstance(value, list): + return [ + item + for sublist in [dump_if_base_model(val) for val in value] + for item in sublist + ] + + return [json.dumps(value, sort_keys=True, separators=(",", ":"))] diff --git a/src/fastmcp/server/middleware/middleware.py b/src/fastmcp/server/middleware/middleware.py index 2eb0894bf..8b262d4f5 100644 --- a/src/fastmcp/server/middleware/middleware.py +++ b/src/fastmcp/server/middleware/middleware.py @@ -15,6 +15,7 @@ from typing import ( ) import mcp.types as mt +from mcp.server.lowlevel.helper_types import ReadResourceContents from typing_extensions import TypeVar from fastmcp.prompts.prompt import Prompt @@ -154,8 +155,8 @@ class Middleware: async def on_read_resource( self, context: MiddlewareContext[mt.ReadResourceRequestParams], - call_next: CallNext[mt.ReadResourceRequestParams, mt.ReadResourceResult], - ) -> mt.ReadResourceResult: + call_next: CallNext[mt.ReadResourceRequestParams, list[ReadResourceContents]], + ) -> list[ReadResourceContents]: return await call_next(context) async def on_get_prompt( diff --git a/tests/server/middleware/test_caching.py b/tests/server/middleware/test_caching.py index 871f2c24e..37bd333d5 100644 --- a/tests/server/middleware/test_caching.py +++ b/tests/server/middleware/test_caching.py @@ -1,6 +1,7 @@ """Tests for response caching middleware.""" import tempfile +from collections.abc import Sequence from datetime import datetime, timedelta, timezone from typing import Any from unittest.mock import AsyncMock, MagicMock @@ -8,20 +9,56 @@ from unittest.mock import AsyncMock, MagicMock import mcp.types import pytest from inline_snapshot import snapshot -from pydantic import BaseModel +from mcp.server.lowlevel.helper_types import ReadResourceContents +from mcp.types import ( + TextContent, + TextResourceContents, +) +from pydantic import AnyUrl, BaseModel from fastmcp import FastMCP from fastmcp.client import Client +from fastmcp.client.client import CallToolResult +from fastmcp.client.transports import FastMCPTransport +from fastmcp.prompts.prompt import FunctionPrompt from fastmcp.server.middleware.caching import ( + CachableTypes, + CachedPrompt, + CachedResource, CacheEntry, + CacheMethodStats, + CacheProtocol, CacheStats, + CallToolSettings, DiskCache, InMemoryCache, + MethodSettings, ResponseCachingMiddleware, ) -from fastmcp.server.middleware.middleware import MiddlewareContext +from fastmcp.server.middleware.middleware import CallNext, MiddlewareContext from fastmcp.tools.tool import Tool, ToolResult +TEST_URI = AnyUrl("https://test_uri") + +SAMPLE_RESOURCE = CachedResource(name="resource", uri=TEST_URI, mime_type="text/plain") +SAMPLE_PROMPT = CachedPrompt(name="prompt") +SAMPLE_READ_RESOURCE_CONTENTS = ReadResourceContents( + content="test_text", + mime_type="text/plain", +) +SAMPLE_GET_PROMPT_RESULT = mcp.types.GetPromptResult( + messages=[ + mcp.types.PromptMessage( + role="user", content=mcp.types.TextContent(type="text", text="test_text") + ) + ] +) +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"}, +) + class CrazyModel(BaseModel): a: int @@ -35,15 +72,43 @@ class CrazyModel(BaseModel): i: dict[str, list[int]] -def extract_content_for_snapshot(result: ToolResult) -> dict[str, Any]: +def extract_content_for_snapshot(result: ToolResult | CallToolResult) -> dict[str, Any]: return { "content": [c.model_dump() for c in result.content], "structured_content": result.structured_content, } +def dump_mcp_type( + model: BaseModel | ToolResult | ReadResourceContents, +) -> dict[str, Any]: + if isinstance(model, ToolResult): + return extract_content_for_snapshot(model) + + if isinstance(model, ReadResourceContents): + return { + "content": model.content, + "mime_type": model.mime_type, + } + + return model.model_dump() + + +def dump_mcp_types( + model: BaseModel + | ToolResult + | Sequence[BaseModel] + | Sequence[ToolResult] + | list[ReadResourceContents], +) -> list[dict[str, Any]]: + if isinstance(model, Sequence): + return [dump_mcp_type(model=m) for m in model] + + return dump_mcp_type(model=model) # type: ignore + + @pytest.fixture -def crazy_model(): +def crazy_model() -> CrazyModel: return CrazyModel( a=5, b=10, @@ -79,19 +144,55 @@ class TrackingCalculator: self.crazy_calls += 1 return a - def add_tools(self, fastmcp: FastMCP): - fastmcp.add_tool(tool=Tool.from_function(fn=self.add)) - fastmcp.add_tool(tool=Tool.from_function(fn=self.multiply)) - fastmcp.add_tool(tool=Tool.from_function(fn=self.crazy)) + 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) -> int: + return self.add_calls + + def get_multiply_calls(self) -> int: + return self.multiply_calls + + def get_crazy_calls(self) -> int: + return self.crazy_calls + + 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")) + + 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_fn( + fn=self.get_add_calls, uri="resource://add_calls", name=f"{prefix}add_calls" + ) + fastmcp.add_resource_fn( + fn=self.get_multiply_calls, + uri="resource://multiply_calls", + name=f"{prefix}multiply_calls", + ) + fastmcp.add_resource_fn( + fn=self.get_crazy_calls, + uri="resource://crazy_calls", + name=f"{prefix}crazy_calls", + ) @pytest.fixture -def tracking_calculator(): +def tracking_calculator() -> TrackingCalculator: return TrackingCalculator() @pytest.fixture -def mock_context(): +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( @@ -102,7 +203,7 @@ def mock_context(): @pytest.fixture -def mock_call_next(): +def mock_call_next() -> CallNext[mcp.types.CallToolRequestParams, ToolResult]: """Create a mock call_next function.""" return AsyncMock( return_value=ToolResult( @@ -113,7 +214,7 @@ def mock_call_next(): @pytest.fixture -def sample_tool_result(): +def sample_tool_result() -> ToolResult: """Create a sample tool result for testing.""" return ToolResult( content=[{"type": "text", "text": "cached result"}], @@ -131,24 +232,32 @@ class TestCacheEntry: past = now - timedelta(seconds=3600) # Test valid entry - entry = CacheEntry( + entry: CacheEntry[ToolResult] = CacheEntry( + collection="test_collection", key="test_key", - content=[{"type": "text", "text": "test"}], - structured_content='{"result": "success"}', + value=ToolResult( + content=[{"type": "text", "text": "success"}], + structured_content={"result": "success"}, + ), created_at=now, expires_at=future, + ttl=3600, ) assert entry.key == "test_key" assert not entry.is_expired() # Test expired entry - expired_entry = CacheEntry( + expired_entry: CacheEntry[ToolResult] = CacheEntry( + collection="test_collection", key="expired_key", - content=None, - structured_content=None, + value=ToolResult( + content=[{"type": "text", "text": "success"}], + structured_content={"result": "success"}, + ), created_at=past, expires_at=past, + ttl=3600, ) assert expired_entry.is_expired() @@ -156,64 +265,144 @@ class TestCacheEntry: def test_serialization(self): """Test cache entry serialization to/from tool result.""" tool_result = ToolResult( - content=[{"type": "text", "text": "test"}], + content=[{"type": "text", "text": "success"}], structured_content={"result": "success"}, ) # Test round-trip conversion - entry = CacheEntry.from_tool_result("test_key", tool_result, 3600) - result = entry.to_tool_result() + entry: CacheEntry[ToolResult] = CacheEntry.from_value( + collection="test_collection", + key="test_key", + value=tool_result, + ttl=3600, + ) - assert result.content == tool_result.content - assert result.structured_content == tool_result.structured_content + retrieved_tool_result: ToolResult = entry.value + + assert retrieved_tool_result.content == tool_result.content + + assert ( + retrieved_tool_result.structured_content == tool_result.structured_content + ) -class TestInMemoryCache: +class TestMemoryCache: """Test InMemoryCache implementation.""" - async def test_basic_operations(self, sample_tool_result): - """Test basic cache operations.""" - cache = InMemoryCache() - - # Test set and get - await cache.set("test_key", sample_tool_result, 3600) - result = await cache.get("test_key") - - assert result is not None - assert result.content == sample_tool_result.content - - # Test delete - await cache.delete("test_key") - assert await cache.get("test_key") is None - - async def test_expiration_and_cleanup(self, sample_tool_result): - """Test cache expiration and cleanup.""" - cache = InMemoryCache() - - # Create an expired entry - entry = CacheEntry.from_tool_result("expired_key", sample_tool_result, -1) - cache._cache["expired_key"] = entry - - # Should return None and remove expired entry - result = await cache.get("expired_key") - assert result is None - assert "expired_key" not in cache._cache - async def test_size_limit(self, sample_tool_result): """Test cache size limit enforcement.""" - cache = InMemoryCache(max_size=2) + cache = InMemoryCache(max_entries=2) # Fill cache to capacity - await cache.set("key1", sample_tool_result, 3600) - await cache.set("key2", sample_tool_result, 3600) + await cache.set_value( + collection="test_collection", key="key1", value=sample_tool_result, ttl=3600 + ) + await cache.set_value( + collection="test_collection", key="key2", value=sample_tool_result, ttl=3600 + ) # Add one more - should evict the first - await cache.set("key3", sample_tool_result, 3600) + await cache.set_value( + collection="test_collection", key="key3", value=sample_tool_result, ttl=3600 + ) assert len(cache._cache) == 2 - assert "key1" not in cache._cache - assert "key2" in cache._cache - assert "key3" in cache._cache + assert "test_collection:key1" not in cache._cache + assert "test_collection:key2" in cache._cache + assert "test_collection:key3" in cache._cache + + +class TestCacheImplementations: + """Test InMemoryCache implementation.""" + + @pytest.fixture(params=["memory", "disk"]) + async def cache(self, request): + if request.param == "memory": + return InMemoryCache() + else: + with tempfile.TemporaryDirectory() as temp_dir: + return DiskCache(path=temp_dir) + + async def test_get_none_if_not_set(self, cache: CacheProtocol): + """Test that we get None if a value is not set.""" + assert ( + await cache.get_value(collection="test_collection", key="test_key") is None + ) + + @pytest.mark.parametrize( + "value", + [ + [SAMPLE_TOOL], + SAMPLE_TOOL_RESULT, + [SAMPLE_RESOURCE], + [SAMPLE_READ_RESOURCE_CONTENTS], + [SAMPLE_PROMPT], + SAMPLE_GET_PROMPT_RESULT, + ], + ids=[ + "tool_list", + "tool_result", + "resource", + "read_resource_contents", + "prompt", + "get_prompt_result", + ], + ) + async def test_set_and_get(self, cache: CacheProtocol, value: CachableTypes): + """Test that we can set and then get back a value from the cache.""" + + await cache.set_value( + collection="test_collection", + key="test_key", + value=value, + ttl=3600, + ) + result = await cache.get_value(collection="test_collection", key="test_key") + + assert result is not None + + assert isinstance(result, type(value)) + + assert dump_mcp_types(model=result) == dump_mcp_types(model=value) + + async def test_set_get_delete_get_value(self, cache: CacheProtocol): + """Test that we can set, get, delete, and get a value from the cache.""" + await cache.set_value( + collection="test_collection", + key="test_key", + value=SAMPLE_TOOL_RESULT, + ttl=3600, + ) + result = await cache.get_value(collection="test_collection", key="test_key") + + assert result is not None + assert dump_mcp_types(model=result) == dump_mcp_types(model=SAMPLE_TOOL_RESULT) + + await cache.delete(collection="test_collection", key="test_key") + + assert ( + await cache.get_value(collection="test_collection", key="test_key") is None + ) + + async def test_expiration_and_cleanup(self, cache: CacheProtocol): + """Test cache expiration and cleanup.""" + # Create an expired entry + await cache.set_value( + collection="test_collection", + key="expired_key", + value=SAMPLE_TOOL_RESULT, + ttl=-1, + ) + + # Should return None and remove expired entry + result = await cache.get_value(collection="test_collection", key="expired_key") + + assert result is None + + assert ( + await cache.get_value(collection="test_collection", key="expired_key") + is None + ) class TestResponseCachingMiddleware: @@ -221,43 +410,66 @@ class TestResponseCachingMiddleware: def test_initialization(self): """Test middleware initialization.""" - cache = InMemoryCache() middleware = ResponseCachingMiddleware( - cache_backend=cache, - included_tools=["tool1"], - excluded_tools=["tool2"], + method_settings=MethodSettings( + call_tool=CallToolSettings( + included_tools=["tool1"], + excluded_tools=["tool2"], + ) + ), default_ttl=1800, ) - assert middleware._backend is cache + assert middleware.method_settings == snapshot( + {"call_tool": {"included_tools": ["tool1"], "excluded_tools": ["tool2"]}} + ) assert middleware._default_ttl == 1800 - assert middleware._included_tools == ["tool1"] - assert middleware._excluded_tools == ["tool2"] - assert middleware._stats.hits == 0 - assert middleware._stats.misses == 0 + assert middleware._max_item_size is None - def test_tool_filtering(self): + @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.""" - cache = InMemoryCache() - # Test included tools only middleware1 = ResponseCachingMiddleware( - cache, included_tools=["tool1", "tool2"] + method_settings=MethodSettings( + call_tool=CallToolSettings( + included_tools=included_tools, excluded_tools=excluded_tools + ) + ), ) - assert middleware1._should_cache_tool("tool1") is True - assert middleware1._should_cache_tool("tool3") is False - - # Test excluded tools - middleware2 = ResponseCachingMiddleware(cache, excluded_tools=["tool1"]) - assert middleware2._should_cache_tool("tool1") is False - assert middleware2._should_cache_tool("tool2") is True - - # Test both (excluded takes precedence) - middleware3 = ResponseCachingMiddleware( - cache, included_tools=["tool1", "tool2"], excluded_tools=["tool2"] + assert ( + middleware1._matches_tool_cache_settings( + context=MiddlewareContext( + method="tools/call", + message=mcp.types.CallToolRequestParams(name=tool_name), + ) + ) + is result ) - assert middleware3._should_cache_tool("tool1") is True - assert middleware3._should_cache_tool("tool2") is False def test_cache_key_generation(self): """Test cache key generation.""" @@ -274,56 +486,70 @@ class TestResponseCachingMiddleware: assert len(key) == 64 assert all(c in "0123456789abcdef" for c in key) - async def test_cache_miss_and_hit(self, mock_context, mock_call_next): + async def test_cache_miss_and_hit( + self, + ): """Test cache miss and hit scenarios.""" - cache = InMemoryCache() - middleware = ResponseCachingMiddleware(cache) + middleware = ResponseCachingMiddleware() + + mock_call_next = AsyncMock( + return_value=ToolResult( + content=[{"type": "text", "text": "test result"}], + structured_content={"result": "success", "value": 123}, + ) + ) + + mock_context = MagicMock( + spec=MiddlewareContext[mcp.types.CallToolRequestParams] + ) + mock_context.message = mcp.types.CallToolRequestParams( + name="test_tool", arguments={"param1": "value1", "param2": 42} + ) + mock_context.method = "tools/call" # First call - cache miss - result1 = await middleware.on_call_tool(mock_context, mock_call_next) - assert middleware._stats.misses == 1 - assert middleware._stats.hits == 0 + result1 = await middleware.on_call_tool( + context=mock_context, call_next=mock_call_next + ) + assert middleware._stats.get_misses("tools/call") == 1 + assert middleware._stats.get_hits("tools/call") == 0 # Second call - cache hit mock_call_next.reset_mock() - result2 = await middleware.on_call_tool(mock_context, mock_call_next) + result2 = await middleware.on_call_tool( + context=mock_context, call_next=mock_call_next + ) assert result1.content == result2.content assert not mock_call_next.called # Should not call downstream - assert middleware._stats.hits == 1 - assert middleware._stats.misses == 1 + assert middleware._stats.get_hits("tools/call") == 1 + assert middleware._stats.get_misses("tools/call") == 1 class TestResponseCachingMiddlewareIntegration: """Integration tests with real FastMCP server.""" - @pytest.fixture - async def disk_cache(self): - with tempfile.TemporaryDirectory() as temp_dir: - yield DiskCache(path=temp_dir) - - @pytest.fixture - async def in_memory_cache(self): - return InMemoryCache() - @pytest.fixture(params=["memory", "disk"]) async def caching_server( self, tracking_calculator: TrackingCalculator, request, - disk_cache, - in_memory_cache, ): """Create a FastMCP server for caching tests.""" mcp = FastMCP("CachingTestServer") - cache = disk_cache if request.param == "disk" else in_memory_cache - - response_caching_middleware = ResponseCachingMiddleware(cache_backend=cache) + with tempfile.TemporaryDirectory() as temp_dir: + response_caching_middleware = ResponseCachingMiddleware( + cache_backend=DiskCache(path=temp_dir) + if request.param == "disk" + else InMemoryCache() + ) mcp.add_middleware(middleware=response_caching_middleware) - tracking_calculator.add_tools(mcp) + tracking_calculator.add_tools(fastmcp=mcp) + tracking_calculator.add_resources(fastmcp=mcp) + tracking_calculator.add_prompts(fastmcp=mcp) return mcp @@ -331,148 +557,119 @@ class TestResponseCachingMiddlewareIntegration: def non_caching_server(self, tracking_calculator: TrackingCalculator): """Create a FastMCP server for non-caching tests.""" mcp = FastMCP("NonCachingTestServer") - tracking_calculator.add_tools(mcp) + tracking_calculator.add_tools(fastmcp=mcp) return mcp - async def test_caching_works_with_real_server( + 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) == 3 + + # 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) == 3 + + assert pre_tool_list == post_tool_list + + async def test_call_tool( self, caching_server: FastMCP, tracking_calculator: TrackingCalculator, - crazy_model: CrazyModel, ): """Test that caching works with a real FastMCP server.""" - tracking_calculator.add_tools(caching_server) + tracking_calculator.add_tools(fastmcp=caching_server) - async with Client(caching_server) as client: - call_tool_result = await client.call_tool("add", {"a": 5, "b": 3}) + async with Client[FastMCPTransport](caching_server) as client: + call_tool_result_one: CallToolResult = await client.call_tool( + "add", {"a": 5, "b": 3} + ) assert tracking_calculator.add_calls == 1 - assert extract_content_for_snapshot(call_tool_result) == snapshot( - { - "content": [ - {"type": "text", "text": "8", "annotations": None, "meta": None} - ], - "structured_content": {"result": 8}, - } + call_tool_result_two: CallToolResult = await client.call_tool( + "add", {"a": 5, "b": 3} ) + assert call_tool_result_one == call_tool_result_two - call_tool_result = await client.call_tool("add", {"a": 5, "b": 3}) - assert tracking_calculator.add_calls == 1 - assert extract_content_for_snapshot(call_tool_result) == snapshot( - { - "content": [ - {"type": "text", "text": "8", "annotations": None, "meta": None} - ], - "structured_content": {"result": 8}, - } - ) - - call_tool_result = await client.call_tool("crazy", {"a": crazy_model}) - assert tracking_calculator.crazy_calls == 1 - assert extract_content_for_snapshot(call_tool_result) == snapshot( - { - "content": [ - { - "type": "text", - "text": '{"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]}}', - "annotations": None, - "meta": None, - } - ], - "structured_content": { - "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]}, - }, - } - ) - - call_tool_result = await client.call_tool("crazy", {"a": crazy_model}) - assert tracking_calculator.crazy_calls == 1 - assert extract_content_for_snapshot(call_tool_result) == snapshot( - { - "content": [ - { - "type": "text", - "text": '{"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]}}', - "annotations": None, - "meta": None, - } - ], - "structured_content": { - "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]}, - }, - } - ) - - async def test_different_arguments_create_different_entries( + async def test_list_resources( self, caching_server: FastMCP, tracking_calculator: TrackingCalculator ): - """Test that different arguments create different cache entries.""" + """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() - async with Client(caching_server) as client: - result1 = await client.call_tool("add", {"a": 5, "b": 10}) - assert tracking_calculator.add_calls == 1 - result2 = await client.call_tool("add", {"a": 1, "b": 5}) - assert tracking_calculator.add_calls == 2 + assert len(pre_resource_list) == 3 - # Results should be different - assert result1.structured_content["result"] == 15 - assert result2.structured_content["result"] == 6 + tracking_calculator.add_resources(fastmcp=caching_server) - async def test_tool_filtering_integration( - self, non_caching_server: FastMCP, tracking_calculator: TrackingCalculator + 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 tool filtering in integration.""" - partial_caching_server = non_caching_server + """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" - partial_caching_server.add_middleware( - ResponseCachingMiddleware( - cache_backend=InMemoryCache(), - included_tools=["add"], # Only cache this tool + 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} ) - ) - async with Client(partial_caching_server) as client: - # This should be cached - await client.call_tool("add", {"a": 5, "b": 10}) - await client.call_tool("add", {"a": 5, "b": 10}) - assert tracking_calculator.add_calls == 1 + 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." + ) - # This should not be cached - await client.call_tool("multiply", {"a": 1, "b": 5}) - await client.call_tool("multiply", {"a": 1, "b": 5}) - assert tracking_calculator.multiply_calls == 2 + tracking_calculator.add_prompts(fastmcp=caching_server) - async def test_cache_stats_tracking(self, non_caching_server: FastMCP): - """Test that cache statistics are properly tracked.""" - middleware = ResponseCachingMiddleware(cache_backend=InMemoryCache()) - non_caching_server.add_middleware(middleware) + post_prompt = await client.get_prompt( + name="how_to_calculate", arguments={"a": 5, "b": 3} + ) - async with Client(non_caching_server) as client: - # First call - cache miss - await client.call_tool("add", {"a": 5, "b": 10}) - assert middleware._stats.misses == 1 - assert middleware._stats.hits == 0 - - # Second call - cache hit - await client.call_tool("add", {"a": 5, "b": 10}) - assert middleware._stats.misses == 1 - assert middleware._stats.hits == 1 + assert pre_prompt == post_prompt class TestCacheStats: @@ -480,6 +677,14 @@ class TestCacheStats: def test_stats_initialization(self): """Test cache stats initialization.""" - stats = CacheStats(hits=5, misses=10) - assert stats.hits == 5 - assert stats.misses == 10 + stats = CacheStats( + collections={ + "tools/call": CacheMethodStats(hits=5, misses=10), + "tools/list": CacheMethodStats(hits=0, misses=0), + } + ) + + assert stats.get_hits("tools/call") == 5 + assert stats.get_misses("tools/call") == 10 + assert stats.get_hits("tools/list") == 0 + assert stats.get_misses("tools/list") == 0 diff --git a/uv.lock b/uv.lock index 36b3b1181..fb632a6c4 100644 --- a/uv.lock +++ b/uv.lock @@ -69,6 +69,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a0/59/76ab57e3fe74484f48a53f8e337171b4a2349e506eabe136d7e01d059086/backports_asyncio_runner-1.2.0-py3-none-any.whl", hash = "sha256:0da0a936a8aeb554eccb426dc55af3ba63bcdc69fa1a600b5bb305413a4477b5", size = 12313, upload-time = "2025-07-02T02:27:14.263Z" }, ] +[[package]] +name = "cachetools" +version = "6.2.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/9d/61/e4fad8155db4a04bfb4734c7c8ff0882f078f24294d42798b3568eb63bff/cachetools-6.2.0.tar.gz", hash = "sha256:38b328c0889450f05f5e120f56ab68c8abaf424e1275522b138ffc93253f7e32", size = 30988, upload-time = "2025-08-25T18:57:30.924Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/6c/56/3124f61d37a7a4e7cc96afc5492c78ba0cb551151e530b54669ddd1436ef/cachetools-6.2.0-py3-none-any.whl", hash = "sha256:1c76a8960c0041fcc21097e357f882197c79da0dbff766e7317890a65d7d8ba6", size = 11276, upload-time = "2025-08-25T18:57:29.684Z" }, +] + [[package]] name = "certifi" version = "2025.8.3" @@ -530,7 +539,6 @@ source = { editable = "." } dependencies = [ { name = "authlib" }, { name = "cyclopts" }, - { name = "diskcache" }, { name = "exceptiongroup" }, { name = "httpx" }, { name = "mcp" }, @@ -543,6 +551,10 @@ dependencies = [ ] [package.optional-dependencies] +caching = [ + { name = "cachetools" }, + { name = "diskcache" }, +] openai = [ { name = "openai" }, ] @@ -554,7 +566,7 @@ websockets = [ dev = [ { name = "dirty-equals" }, { name = "fastapi" }, - { name = "fastmcp", extra = ["openai"] }, + { name = "fastmcp", extra = ["caching", "openai"] }, { name = "inline-snapshot", extra = ["dirty-equals"] }, { name = "ipython", version = "8.37.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "ipython", version = "9.4.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, @@ -565,6 +577,7 @@ dev = [ { name = "pyperclip" }, { name = "pytest" }, { name = "pytest-asyncio" }, + { name = "pytest-benchmark" }, { name = "pytest-cov" }, { name = "pytest-env" }, { name = "pytest-flakefinder" }, @@ -579,8 +592,9 @@ dev = [ [package.metadata] requires-dist = [ { name = "authlib", specifier = ">=1.5.2" }, + { name = "cachetools", marker = "extra == 'caching'", specifier = ">=6.2.0" }, { name = "cyclopts", specifier = ">=3.0.0" }, - { name = "diskcache", specifier = ">=5.6.3" }, + { name = "diskcache", marker = "extra == 'caching'", specifier = ">=5.6.3" }, { name = "exceptiongroup", specifier = ">=1.2.2" }, { name = "httpx", specifier = ">=0.28.1" }, { name = "mcp", specifier = ">=1.12.4,<2.0.0" }, @@ -593,12 +607,13 @@ requires-dist = [ { name = "rich", specifier = ">=13.9.4" }, { name = "websockets", marker = "extra == 'websockets'", specifier = ">=15.0.1" }, ] -provides-extras = ["openai", "websockets"] +provides-extras = ["caching", "openai", "websockets"] [package.metadata.requires-dev] dev = [ { name = "dirty-equals", specifier = ">=0.9.0" }, { name = "fastapi", specifier = ">=0.115.12" }, + { name = "fastmcp", extras = ["caching"] }, { name = "fastmcp", extras = ["openai"] }, { name = "inline-snapshot", extras = ["dirty-equals"], specifier = ">=0.27.2" }, { name = "ipython", specifier = ">=8.12.3" }, @@ -609,6 +624,7 @@ dev = [ { name = "pyperclip", specifier = ">=1.9.0" }, { name = "pytest", specifier = ">=8.3.3" }, { name = "pytest-asyncio", specifier = ">=0.23.5" }, + { name = "pytest-benchmark", specifier = ">=5.1.0" }, { name = "pytest-cov", specifier = ">=6.1.1" }, { name = "pytest-env", specifier = ">=1.1.5" }, { name = "pytest-flakefinder" }, @@ -1290,6 +1306,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/8e/37/efad0257dc6e593a18957422533ff0f87ede7c9c6ea010a2177d738fb82f/pure_eval-0.2.3-py3-none-any.whl", hash = "sha256:1db8e35b67b3d218d818ae653e27f06c3aa420901fa7b081ca98cbedc874e0d0", size = 11842, upload-time = "2024-07-21T12:58:20.04Z" }, ] +[[package]] +name = "py-cpuinfo" +version = "9.0.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/37/a8/d832f7293ebb21690860d2e01d8115e5ff6f2ae8bbdc953f0eb0fa4bd2c7/py-cpuinfo-9.0.0.tar.gz", hash = "sha256:3cdbbf3fac90dc6f118bfd64384f309edeadd902d7c8fb17f02ffa1fc3f49690", size = 104716, upload-time = "2022-10-25T20:38:06.303Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e0/a9/023730ba63db1e494a271cb018dcd361bd2c917ba7004c3e49d5daf795a2/py_cpuinfo-9.0.0-py3-none-any.whl", hash = "sha256:859625bc251f64e21f077d099d4162689c762b5d6a4c3c97553d56241c9674d5", size = 22335, upload-time = "2022-10-25T20:38:27.636Z" }, +] + [[package]] name = "pycparser" version = "2.22" @@ -1540,6 +1565,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c7/9d/bf86eddabf8c6c9cb1ea9a869d6873b46f105a5d292d3a6f7071f5b07935/pytest_asyncio-1.1.0-py3-none-any.whl", hash = "sha256:5fe2d69607b0bd75c656d1211f969cadba035030156745ee09e7d71740e58ecf", size = 15157, upload-time = "2025-07-16T04:29:24.929Z" }, ] +[[package]] +name = "pytest-benchmark" +version = "5.1.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "py-cpuinfo" }, + { name = "pytest" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/39/d0/a8bd08d641b393db3be3819b03e2d9bb8760ca8479080a26a5f6e540e99c/pytest-benchmark-5.1.0.tar.gz", hash = "sha256:9ea661cdc292e8231f7cd4c10b0319e56a2118e2c09d9f50e1b3d150d2aca105", size = 337810, upload-time = "2024-10-30T11:51:48.521Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9e/d6/b41653199ea09d5969d4e385df9bbfd9a100f28ca7e824ce7c0a016e3053/pytest_benchmark-5.1.0-py3-none-any.whl", hash = "sha256:922de2dfa3033c227c96da942d1878191afa135a29485fb942e85dff1c592c89", size = 44259, upload-time = "2024-10-30T11:51:45.94Z" }, +] + [[package]] name = "pytest-cov" version = "6.2.1"