From 5831c4bb60919618cef47d238dc884d1955b34a2 Mon Sep 17 00:00:00 2001 From: William Easton Date: Thu, 16 Oct 2025 21:56:55 -0500 Subject: [PATCH] PR Clean-up --- docs/servers/middleware.mdx | 26 +- pyproject.toml | 3 +- src/fastmcp/server/middleware/caching.py | 568 ++++++++--------------- tests/server/middleware/test_caching.py | 257 +++++----- 4 files changed, 305 insertions(+), 549 deletions(-) diff --git a/docs/servers/middleware.mdx b/docs/servers/middleware.mdx index 4919462b5..b38a0a649 100644 --- a/docs/servers/middleware.mdx +++ b/docs/servers/middleware.mdx @@ -461,21 +461,22 @@ from fastmcp.server.middleware.caching import ResponseCachingMiddleware mcp.add_middleware(ResponseCachingMiddleware()) ``` -Out of the box, it caches call/list tool, resources, and prompts. Sending a notification of a tool/resource/prompt change will invalidate the cache for the affected method. +Out of the box, it caches call/list tool, resources, and prompts to an in-memory cache. Sending a notification of a tool/resource/prompt change will invalidate the cache for the affected method. List calls are stored under global keys, if you share a key_value backend across servers, keep this in mind and consider using the PrefixCollectionsWrapper in py-key-value-aio to namespace collections by server. -Alternatively, it can be configured to only cache specific methods, for example, only caching list tools and only caching calls to `tool1`: +Each method can be configured individually, for example, caching list tools for 30 seconds, skipping caching for tools other than `tool1` and not caching and requests to read resources: ```python -from fastmcp.server.middleware.caching import ResponseCachingMiddleware, MethodSettings, CallToolSettings, ListToolsSettings +from fastmcp.server.middleware.caching import ResponseCachingMiddleware, CallToolSettings, ListToolsSettings, ReadResourceSettings mcp.add_middleware(ResponseCachingMiddleware( - method_settings=MethodSettings( - call_tool=CallToolSettings( - included_tools=["tool1"], - ), - list_tools=ListToolsSettings( - ttl=30, - ) + list_tools_settings=ListToolsSettings( + ttl=30, + ), + call_tool_settings=CallToolSettings( + included_tools=["tool1"], + ), + read_resource_settings=ReadResourceSettings( + enabled=False ) )) ``` @@ -483,10 +484,11 @@ mcp.add_middleware(ResponseCachingMiddleware( It can also be configured to cache to disk: ```python -from fastmcp.server.middleware.caching import ResponseCachingMiddleware, DiskCache +from fastmcp.server.middleware.caching import ResponseCachingMiddleware +from key_value.aio.stores.disk import DiskStore mcp.add_middleware(ResponseCachingMiddleware( - cache_backend=DiskCache(path="cache"), + cache_storage=DiskStore(directory="cache"), )) ``` diff --git a/pyproject.toml b/pyproject.toml index bd6563745..e6c93d38f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -15,7 +15,7 @@ dependencies = [ "pydantic[email]>=2.11.7", "pyperclip>=1.9.0", "openapi-core>=0.19.5", - "py-key-value-aio[disk,memory]>=0.2.2", + "py-key-value-aio[disk,memory]>=0.2.2,<0.3.0", "websockets>=15.0.1", ] @@ -70,7 +70,6 @@ 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 fb61d4f96..9bc64b1fa 100644 --- a/src/fastmcp/server/middleware/caching.py +++ b/src/fastmcp/server/middleware/caching.py @@ -1,20 +1,20 @@ """A middleware for response caching.""" -import json from collections.abc import Sequence -from typing import Any, TypedDict, TypeVar, cast +from logging import Logger +from typing import Any, TypedDict import mcp.types import pydantic_core from key_value.aio.adapters.pydantic import PydanticAdapter from key_value.aio.protocols.key_value import AsyncKeyValue from key_value.aio.stores.memory import MemoryStore +from key_value.aio.wrappers.limit_size import LimitSizeWrapper from key_value.aio.wrappers.statistics import StatisticsWrapper from key_value.aio.wrappers.statistics.wrapper import ( KVStoreCollectionStatistics, ) from mcp.server.lowlevel.helper_types import ReadResourceContents -from mcp.types import PromptMessage from pydantic import BaseModel, Field from typing_extensions import NotRequired, Self, override @@ -24,130 +24,40 @@ from fastmcp.server.middleware.middleware import CallNext, Middleware, Middlewar from fastmcp.tools.tool import Tool, ToolResult from fastmcp.utilities.logging import get_logger -logger = get_logger(__name__) +logger: Logger = get_logger(name=__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 CachableToolResult(BaseModel, ToolResult): - structured_content: dict[str, Any] | None - content: list[mcp.types.ContentBlock] +class CachableReadResourceContents(BaseModel): + """A wrapper for ReadResourceContents that can be cached.""" - @classmethod - def from_tool_result(cls, tool_result: ToolResult) -> Self: - return cls( - structured_content=tool_result.structured_content, - content=tool_result.content, - ) + content: str | bytes + mime_type: str | None = None def get_size(self) -> int: - return _get_size_of_tool_result(self) - - -class CachablePrompt(Prompt): - @override - async def render( - self, - arguments: dict[str, Any] | None = None, - ) -> list[PromptMessage]: - """Render the prompt with arguments.""" - raise NotImplementedError( - "Prompt.render() is not implemented on cached prompts" - ) + return len(self.model_dump_json()) @classmethod - def from_list_prompts(cls, prompts: list[Prompt]) -> list[Self]: - cachable_prompts: list[Self] = [] - for prompt in prompts: - cachable_prompts.append( - cls( - name=prompt.name, - title=prompt.title, - description=prompt.description, - arguments=prompt.arguments, - meta=prompt.meta, - tags=prompt.tags, - enabled=prompt.enabled, - ) - ) - return cachable_prompts - - -class CachablePromptResult(mcp.types.GetPromptResult): ... - - -class CachableResource(Resource): - @override - async def read(self) -> str | bytes: - """Read the resource content.""" - raise NotImplementedError( - "Resource.read() is not implemented on cached resources" - ) + def get_sizes(cls, values: Sequence[Self]) -> int: + return sum([item.get_size() for item in values]) @classmethod - def from_list_resources(cls, resources: list[Resource]) -> list[Self]: - cachable_resources: list[Self] = [] - for resource in resources: - cachable_resources.append( - cls( - name=resource.name, - description=resource.description, - uri=resource.uri, - mime_type=resource.mime_type, - annotations=resource.annotations, - meta=resource.meta, - tags=resource.tags, - enabled=resource.enabled, - ) - ) - return cachable_resources + def wrap(cls, values: Sequence[ReadResourceContents]) -> list[Self]: + return [cls(content=item.content, mime_type=item.mime_type) for item in values] - -class CachableTool(Tool): @classmethod - def from_list_tools(cls, tools: list[Tool]) -> list[Self]: - cachable_tools: list[Self] = [] - for tool in tools: - cachable_tools.append( - cls( - name=tool.name, - description=tool.description, - parameters=tool.parameters, - output_schema=tool.output_schema, - annotations=tool.annotations, - serializer=tool.serializer, - meta=tool.meta, - tags=tool.tags, - enabled=tool.enabled, - ) - ) - return cachable_tools - - -class CachableReadResourceContents(BaseModel, ReadResourceContents): ... - - -class CachableToolList(BaseModel): - cachable_tools: list[CachableTool] - - -class CachableResourceList(BaseModel): - cachable_resources: list[CachableResource] - - -class CachablePromptList(BaseModel): - cachable_prompts: list[CachablePrompt] - - -class CachableReadResourceContentsList(BaseModel): - cachable_read_resource_contents: list[CachableReadResourceContents] + def unwrap(cls, values: Sequence[Self]) -> list[ReadResourceContents]: + return [ + ReadResourceContents(content=item.content, mime_type=item.mime_type) + for item in values + ] class SharedMethodSettings(TypedDict): @@ -184,53 +94,6 @@ class GetPromptSettings(SharedMethodSettings): """Configuration options for Prompt-related caching.""" -class MethodSettings(TypedDict): - """Configuration options for mcp "methods" in the response caching middleware.""" - - list_tools: NotRequired[ListToolsSettings] - call_tool: NotRequired[CallToolSettings] - - list_resources: NotRequired[ListResourcesSettings] - read_resource: NotRequired[ReadResourceSettings] - - list_prompts: NotRequired[ListPromptsSettings] - get_prompt: NotRequired[GetPromptSettings] - - -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 ResponseCachingStatistics(BaseModel): list_tools: KVStoreCollectionStatistics | None = Field(default=None) list_resources: KVStoreCollectionStatistics | None = Field(default=None) @@ -255,66 +118,93 @@ class ResponseCachingMiddleware(Middleware): def __init__( self, - cache_store: AsyncKeyValue | None = None, - method_settings: MethodSettings | None = None, - default_ttl: int = ONE_HOUR_IN_SECONDS, - max_item_size: int | None = None, + cache_storage: AsyncKeyValue | None = None, + list_tools_settings: ListToolsSettings | None = None, + list_resources_settings: ListResourcesSettings | None = None, + list_prompts_settings: ListPromptsSettings | None = None, + read_resource_settings: ReadResourceSettings | None = None, + get_prompt_settings: GetPromptSettings | None = None, + call_tool_settings: CallToolSettings | None = None, + max_item_size: int = ONE_MB_IN_BYTES, ): """Initialize the response caching middleware. Args: - cache_backend: The cache backend to use. If None, an in-memory cache is used. - 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. + cache_storage: The cache backend to use. If None, an in-memory cache is used. + list_tools_settings: The settings for the list tools method. If None, the default settings are used (5 minute TTL). + list_resources_settings: The settings for the list resources method. If None, the default settings are used (5 minute TTL). + list_prompts_settings: The settings for the list prompts method. If None, the default settings are used (5 minute TTL). + read_resource_settings: The settings for the read resource method. If None, the default settings are used (1 hour TTL). + get_prompt_settings: The settings for the get prompt method. If None, the default settings are used (1 hour TTL). + call_tool_settings: The settings for the call tool method. If None, the default settings are used (1 hour TTL). + max_item_size: The maximum size of items eligible for caching. Defaults to 1MB. """ - self._default_ttl: int = default_ttl - self._backend: AsyncKeyValue = cache_store or MemoryStore() - self._stats: StatisticsWrapper = StatisticsWrapper(store=self._backend) - self._max_item_size: int | None = max_item_size + self._backend: AsyncKeyValue = cache_storage or MemoryStore() - self.method_settings: MethodSettings = ( - method_settings or DEFAULT_METHOD_SETTINGS + # When the size limit is exceeded, the put will silently fail + self._size_limiter: LimitSizeWrapper = LimitSizeWrapper( + key_value=self._backend, max_size=max_item_size, raise_on_too_large=False + ) + self._stats: StatisticsWrapper = StatisticsWrapper(key_value=self._size_limiter) + + self._list_tools_settings: ListToolsSettings = ( + list_tools_settings or ListToolsSettings() + ) + self._list_resources_settings: ListResourcesSettings = ( + list_resources_settings or ListResourcesSettings() + ) + self._list_prompts_settings: ListPromptsSettings = ( + list_prompts_settings or ListPromptsSettings() ) - self._list_tools_cache: PydanticAdapter[CachableToolList] = PydanticAdapter( + self._read_resource_settings: ReadResourceSettings = ( + read_resource_settings or ReadResourceSettings() + ) + self._get_prompt_settings: GetPromptSettings = ( + get_prompt_settings or GetPromptSettings() + ) + self._call_tool_settings: CallToolSettings = ( + call_tool_settings or CallToolSettings() + ) + + self._list_tools_cache: PydanticAdapter[list[Tool]] = PydanticAdapter( key_value=self._stats, - pydantic_model=CachableToolList, + pydantic_model=list[Tool], default_collection="tools/list", ) - self._list_resources_cache: PydanticAdapter[CachableResourceList] = ( - PydanticAdapter( - key_value=self._stats, - pydantic_model=CachableResourceList, - default_collection="resources/list", - ) + self._list_resources_cache: PydanticAdapter[list[Resource]] = PydanticAdapter( + key_value=self._stats, + pydantic_model=list[Resource], + default_collection="resources/list", ) - self._list_prompts_cache: PydanticAdapter[CachablePromptList] = PydanticAdapter( + self._list_prompts_cache: PydanticAdapter[list[Prompt]] = PydanticAdapter( key_value=self._stats, - pydantic_model=CachablePromptList, + pydantic_model=list[Prompt], default_collection="prompts/list", ) - self._read_resource_cache: PydanticAdapter[CachableReadResourceContentsList] = ( + self._read_resource_cache: PydanticAdapter[ + list[CachableReadResourceContents] + ] = PydanticAdapter( + key_value=self._stats, + pydantic_model=list[CachableReadResourceContents], + default_collection="resources/read", + ) + + self._get_prompt_cache: PydanticAdapter[mcp.types.GetPromptResult] = ( PydanticAdapter( key_value=self._stats, - pydantic_model=CachableReadResourceContentsList, - default_collection="resources/read", + pydantic_model=mcp.types.GetPromptResult, + default_collection="prompts/get", ) ) - self._get_prompt_cache: PydanticAdapter[CachablePromptResult] = PydanticAdapter( + self._call_tool_cache: PydanticAdapter[ToolResult] = PydanticAdapter( key_value=self._stats, - pydantic_model=CachablePromptResult, - default_collection="prompts/get", - ) - - self._call_tool_cache: PydanticAdapter[CachableToolResult] = PydanticAdapter( - key_value=self._stats, - pydantic_model=CachableToolResult, + pydantic_model=ToolResult, default_collection="tools/call", ) @@ -322,21 +212,38 @@ class ResponseCachingMiddleware(Middleware): 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) + call_next: CallNext[mcp.types.ListToolsRequest, Sequence[Tool]], + ) -> Sequence[Tool]: + """List tools from the cache, if caching is enabled, and the result is in the cache. Otherwise, + otherwise call the next middleware and store the result in the cache if caching is enabled.""" + if self._list_tools_settings.get("enabled") is False: + return await call_next(context) if cached_value := await self._list_tools_cache.get(key=GLOBAL_KEY): - return cached_value.cachable_tools + return cached_value - value: list[Tool] = await call_next(context=context) + tools: Sequence[Tool] = await call_next(context=context) - cachable_tools: list[CachableTool] = CachableTool.from_list_tools(tools=value) + # Turn any subclass of Tool into a Tool + cachable_tools: list[Tool] = [ + 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, + enabled=tool.enabled, + ) + for tool in tools + ] await self._list_tools_cache.put( key=GLOBAL_KEY, - value=CachableToolList(cachable_tools=cachable_tools), + value=cachable_tools, + ttl=self._list_tools_settings.get("ttl", FIVE_MINUTES_IN_SECONDS), ) return cachable_tools @@ -345,25 +252,38 @@ class ResponseCachingMiddleware(Middleware): async def on_list_resources( self, context: MiddlewareContext[mcp.types.ListResourcesRequest], - call_next: CallNext[mcp.types.ListResourcesRequest, list[Resource]], - ) -> list[Resource]: + call_next: CallNext[mcp.types.ListResourcesRequest, Sequence[Resource]], + ) -> Sequence[Resource]: """List resources from the cache, if caching is enabled, and the result is in the cache. Otherwise, otherwise call the next middleware and store the result in the cache if caching is enabled.""" - if self._should_bypass_caching(context=context): + if self._list_resources_settings.get("enabled") is False: return await call_next(context) if cached_value := await self._list_resources_cache.get(key=GLOBAL_KEY): - return cached_value.cachable_resources + return cached_value - value: list[Resource] = await call_next(context=context) + resources: Sequence[Resource] = await call_next(context=context) - cachable_resources: list[CachableResource] = ( - CachableResource.from_list_resources(resources=value) - ) + # Turn any subclass of Resource into a Resource + cachable_resources: list[Resource] = [ + Resource( + name=resource.name, + title=resource.title, + description=resource.description, + tags=resource.tags, + meta=resource.meta, + mime_type=resource.mime_type, + annotations=resource.annotations, + enabled=resource.enabled, + uri=resource.uri, + ) + for resource in resources + ] await self._list_resources_cache.put( key=GLOBAL_KEY, - value=CachableResourceList(cachable_resources=cachable_resources), + value=cachable_resources, + ttl=self._list_resources_settings.get("ttl", FIVE_MINUTES_IN_SECONDS), ) return cachable_resources @@ -372,25 +292,36 @@ class ResponseCachingMiddleware(Middleware): async def on_list_prompts( self, context: MiddlewareContext[mcp.types.ListPromptsRequest], - call_next: CallNext[mcp.types.ListPromptsRequest, list[Prompt]], - ) -> list[Prompt]: + call_next: CallNext[mcp.types.ListPromptsRequest, Sequence[Prompt]], + ) -> Sequence[Prompt]: """List prompts from the cache, if caching is enabled, and the result is in the cache. Otherwise, otherwise call the next middleware and store the result in the cache if caching is enabled.""" - if self._should_bypass_caching(context=context): + if self._list_prompts_settings.get("enabled") is False: return await call_next(context) if cached_value := await self._list_prompts_cache.get(key=GLOBAL_KEY): - return cached_value.cachable_prompts + return cached_value - value: list[Prompt] = await call_next(context=context) + prompts: Sequence[Prompt] = await call_next(context=context) - cachable_prompts: list[CachablePrompt] = CachablePrompt.from_list_prompts( - prompts=value - ) + # Turn any subclass of Prompt into a Prompt + cachable_prompts: list[Prompt] = [ + Prompt( + name=prompt.name, + title=prompt.title, + description=prompt.description, + tags=prompt.tags, + meta=prompt.meta, + enabled=prompt.enabled, + arguments=prompt.arguments, + ) + for prompt in prompts + ] await self._list_prompts_cache.put( key=GLOBAL_KEY, - value=CachablePromptList(cachable_prompts=cachable_prompts), + value=cachable_prompts, + ttl=self._list_prompts_settings.get("ttl", FIVE_MINUTES_IN_SECONDS), ) return cachable_prompts @@ -400,69 +331,60 @@ class ResponseCachingMiddleware(Middleware): self, context: MiddlewareContext[mcp.types.CallToolRequestParams], call_next: CallNext[mcp.types.CallToolRequestParams, ToolResult], - ) -> Any: + ) -> ToolResult: """Call a tool from the cache, if caching is enabled, and the result is in the cache. Otherwise, otherwise call the next middleware and store the result in the cache if caching is enabled.""" - if self._should_bypass_caching(context=context): + tool_name = context.message.name + + if self._call_tool_settings.get( + "enabled" + ) is False or not self._matches_tool_cache_settings(tool_name=tool_name): return await call_next(context=context) - if not self._matches_tool_cache_settings(context=context): - return await call_next(context=context) + cache_key: str = f"{tool_name}:{_get_arguments_str(context.message.arguments)}" - if cached_value := await self._call_tool_cache.get( - key=_make_call_tool_cache_key(msg=context.message) - ): + if cached_value := await self._call_tool_cache.get(key=cache_key): return cached_value tool_result: ToolResult = await call_next(context=context) - cachable_value: CachableToolResult = CachableToolResult.from_tool_result( - tool_result=tool_result - ) - - if self._max_item_size and cachable_value.get_size() > self._max_item_size: - return tool_result - await self._call_tool_cache.put( - key=_make_call_tool_cache_key(msg=context.message), - value=cachable_value, + key=cache_key, + value=tool_result, + ttl=self._call_tool_settings.get("ttl", ONE_HOUR_IN_SECONDS), ) - return cachable_value + return tool_result @override async def on_read_resource( self, context: MiddlewareContext[mcp.types.ReadResourceRequestParams], call_next: CallNext[ - mcp.types.ReadResourceRequestParams, list[ReadResourceContents] + mcp.types.ReadResourceRequestParams, Sequence[ReadResourceContents] ], - ) -> list[ReadResourceContents]: + ) -> Sequence[ReadResourceContents]: """Read a resource from the cache, if caching is enabled, and the result is in the cache. Otherwise, otherwise call the next middleware and store the result in the cache if caching is enabled.""" - if self._should_bypass_caching(context=context): + if self._read_resource_settings.get("enabled") is False: return await call_next(context=context) - if cached_value := await self._read_resource_cache.get( - key=_make_read_resource_cache_key(msg=context.message) - ): - return cached_value.cachable_read_resource_contents + cache_key: str = str(context.message.uri) + cached_value: list[CachableReadResourceContents] | None - value: list[ReadResourceContents] = await call_next(context=context) + if cached_value := await self._read_resource_cache.get(key=cache_key): + return CachableReadResourceContents.unwrap(values=cached_value) - cachable_read_resource_contents: list[CachableReadResourceContents] = [ - CachableReadResourceContents(content=item.content, mime_type=item.mime_type) - for item in value - ] + value: Sequence[ReadResourceContents] = await call_next(context=context) + cached_value = CachableReadResourceContents.wrap(values=value) await self._read_resource_cache.put( - key=_make_read_resource_cache_key(msg=context.message), - value=CachableReadResourceContentsList( - cachable_read_resource_contents=cachable_read_resource_contents - ), + key=cache_key, + value=cached_value, + ttl=self._read_resource_settings.get("ttl", ONE_HOUR_IN_SECONDS), ) - return cachable_read_resource_contents + return CachableReadResourceContents.unwrap(values=cached_value) @override async def on_get_prompt( @@ -474,117 +396,39 @@ class ResponseCachingMiddleware(Middleware): ) -> mcp.types.GetPromptResult: """Get a prompt from the cache, if caching is enabled, and the result is in the cache. Otherwise, otherwise call the next middleware and store the result in the cache if caching is enabled.""" - if self._should_bypass_caching(context=context): - return await call_next(context) + if self._get_prompt_settings.get("enabled") is False: + return await call_next(context=context) - if cached_value := await self._get_prompt_cache.get( - key=_make_get_prompt_cache_key(msg=context.message) - ): + cache_key: str = f"{context.message.name}:{_get_arguments_str(arguments=context.message.arguments)}" + + if cached_value := await self._get_prompt_cache.get(key=cache_key): return cached_value value: mcp.types.GetPromptResult = await call_next(context=context) - cachable_value: CachablePromptResult = CachablePromptResult( - messages=value.messages, description=value.description, _meta=value.meta - ) - await self._get_prompt_cache.put( - key=_make_get_prompt_cache_key(msg=context.message), - value=cachable_value, + key=cache_key, + value=value, + ttl=self._get_prompt_settings.get("ttl", ONE_HOUR_IN_SECONDS), ) - return cachable_value + return value - @override - async def on_notification( - self, - context: MiddlewareContext[mcp.types.Notification[Any, Any]], - call_next: CallNext[mcp.types.Notification[Any, Any], Any], - ) -> Any: - """Handle a notification from the server. If the notification is a tool/resource/prompt list changed - notification, delete the cache for the affected method.""" - if isinstance(context.message, mcp.types.ToolListChangedNotification): - _ = await self._list_tools_cache.delete(key=GLOBAL_KEY) - elif isinstance(context.message, mcp.types.ResourceListChangedNotification): - _ = await self._list_resources_cache.delete(key=GLOBAL_KEY) - elif isinstance(context.message, mcp.types.PromptListChangedNotification): - _ = await self._list_prompts_cache.delete(key=GLOBAL_KEY) - else: - pass - - return await call_next(context=context) - - def _matches_tool_cache_settings( - self, context: MiddlewareContext[mcp.types.CallToolRequestParams] - ) -> bool: + def _matches_tool_cache_settings(self, tool_name: str) -> bool: """Check if the tool matches the cache settings for tool calls.""" - 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 included_tools := self._call_tool_settings.get("included_tools"): if tool_name not in included_tools: return False - if excluded_tools := tool_call_cache_settings.get("excluded_tools"): + if excluded_tools := self._call_tool_settings.get("excluded_tools"): if tool_name in excluded_tools: return False return True - def _get_cache_settings( - self, - context: MiddlewareContext[Any], - settings_type: type[MethodSettingsType] = SharedMethodSettings, - ) -> MethodSettingsType | None: - """Get the cache settings for a method.""" - - 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: - """Get the cache TTL for a method.""" - - 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: - """Check if the method should bypass caching.""" - - if not (cache_settings := self._get_cache_settings(context=context)): - return True - - if cache_settings.get("enabled") is False: - return True - - return False - def statistics(self) -> ResponseCachingStatistics: + """Get the statistics for the cache.""" return ResponseCachingStatistics( list_tools=self._stats.statistics.collections.get("tools/list"), list_resources=self._stats.statistics.collections.get("resources/list"), @@ -595,24 +439,6 @@ class ResponseCachingMiddleware(Middleware): ) -def _make_call_tool_cache_key(msg: mcp.types.CallToolRequestParams) -> str: - """Make a cache key for a tool call by hashing the tool name and its arguments.""" - - return f"{msg.name}:{_get_arguments_str(msg.arguments)}" - - -def _make_read_resource_cache_key(msg: mcp.types.ReadResourceRequestParams) -> str: - """Make a cache key for a resource read by hashing the resource URI.""" - - return f"{msg.uri}" - - -def _make_get_prompt_cache_key(msg: mcp.types.GetPromptRequestParams) -> str: - """Make a cache key for a prompt get by hashing the prompt name and its arguments.""" - - return f"{msg.name}:{_get_arguments_str(msg.arguments)}" - - def _get_arguments_str(arguments: dict[str, Any] | None) -> str: """Get a string representation of the arguments.""" @@ -624,27 +450,3 @@ def _get_arguments_str(arguments: dict[str, Any] | None) -> str: except TypeError: return repr(arguments) - - -def _get_size_of_content_blocks( - value: mcp.types.ContentBlock | Sequence[mcp.types.ContentBlock], -) -> int: - """Get the size of a series of content blocks by summing the size of the JSON representation of each block.""" - - if isinstance(value, mcp.types.ContentBlock): - value = [value] - - return sum([len(item.model_dump_json()) for item in value]) - - -def _get_size_of_tool_result(value: ToolResult) -> int: - """Get the size of a tool result by summing the size of the content blocks and the size of the structured content.""" - - content_size = _get_size_of_content_blocks(value.content) - structured_content_size = len( - json.dumps( - value.structured_content, sort_keys=True, separators=(",", ":") - ).encode("utf-8") - ) - - return content_size + structured_content_size diff --git a/tests/server/middleware/test_caching.py b/tests/server/middleware/test_caching.py index 7701f2735..696933530 100644 --- a/tests/server/middleware/test_caching.py +++ b/tests/server/middleware/test_caching.py @@ -1,8 +1,6 @@ """Tests for response caching middleware.""" import tempfile -from collections.abc import Sequence -from typing import Any from unittest.mock import AsyncMock, MagicMock import mcp.types @@ -10,19 +8,24 @@ import pytest from inline_snapshot import snapshot from key_value.aio.stores.disk import DiskStore 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 PromptMessage, TextContent, TextResourceContents from pydantic import AnyUrl, BaseModel -from fastmcp import FastMCP +from fastmcp import Context, FastMCP from fastmcp.client.client import CallToolResult, Client from fastmcp.client.transports import FastMCPTransport from fastmcp.prompts.prompt import FunctionPrompt, Prompt from fastmcp.resources.resource import Resource from fastmcp.server.middleware.caching import ( CallToolSettings, - MethodSettings, ResponseCachingMiddleware, + ResponseCachingStatistics, ) from fastmcp.server.middleware.middleware import CallNext, MiddlewareContext from fastmcp.tools.tool import Tool, ToolResult @@ -81,41 +84,6 @@ class CrazyModel(BaseModel): i: dict[str, list[int]] -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] - | Sequence[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() -> CrazyModel: return CrazyModel( @@ -171,41 +139,51 @@ class TrackingCalculator: def get_crazy_calls(self) -> int: return self.crazy_calls + async def update_tool_list(self, context: Context): + await context.send_tool_list_changed() + 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( + _ = 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( + _ = 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( + _ = 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( + _ = fastmcp.add_resource( resource=Resource.from_function( fn=self.get_add_calls, uri="resource://add_calls", name=f"{prefix}add_calls", ) ) - fastmcp.add_resource( + _ = fastmcp.add_resource( resource=Resource.from_function( fn=self.get_multiply_calls, uri="resource://multiply_calls", name=f"{prefix}multiply_calls", ) ) - fastmcp.add_resource( + _ = fastmcp.add_resource( resource=Resource.from_function( fn=self.get_crazy_calls, uri="resource://crazy_calls", @@ -235,7 +213,7 @@ def mock_call_next() -> CallNext[mcp.types.CallToolRequestParams, ToolResult]: """Create a mock call_next function.""" return AsyncMock( return_value=ToolResult( - content=[{"type": "text", "text": "test result"}], + content=[TextContent(type="text", text="test result")], structured_content={"result": "success", "value": 123}, ) ) @@ -245,7 +223,7 @@ def mock_call_next() -> CallNext[mcp.types.CallToolRequestParams, ToolResult]: def sample_tool_result() -> ToolResult: """Create a sample tool result for testing.""" return ToolResult( - content=[{"type": "text", "text": "cached result"}], + content=[TextContent(type="text", text="cached result")], structured_content={"cached": True, "data": "test"}, ) @@ -255,22 +233,13 @@ class TestResponseCachingMiddleware: def test_initialization(self): """Test middleware initialization.""" - middleware = ResponseCachingMiddleware( - method_settings=MethodSettings( - call_tool=CallToolSettings( - included_tools=["tool1"], - excluded_tools=["tool2"], - ) + assert ResponseCachingMiddleware( + call_tool_settings=CallToolSettings( + included_tools=["tool1"], + excluded_tools=["tool2"], ), - default_ttl=1800, ) - assert middleware.method_settings == snapshot( - {"call_tool": {"included_tools": ["tool1"], "excluded_tools": ["tool2"]}} - ) - assert middleware._default_ttl == 1800 - assert middleware._max_item_size is None - @pytest.mark.parametrize( ("tool_name", "included_tools", "excluded_tools", "result"), [ @@ -300,90 +269,11 @@ class TestResponseCachingMiddleware: """Test tool filtering logic.""" middleware1 = ResponseCachingMiddleware( - method_settings=MethodSettings( - call_tool=CallToolSettings( - included_tools=included_tools, excluded_tools=excluded_tools - ) + call_tool_settings=CallToolSettings( + included_tools=included_tools, excluded_tools=excluded_tools ), ) - assert ( - middleware1._matches_tool_cache_settings( - context=MiddlewareContext( - method="tools/call", - message=mcp.types.CallToolRequestParams(name=tool_name), - ) - ) - is result - ) - - def test_method_settings(self): - """Test method TTL.""" - middleware = ResponseCachingMiddleware( - method_settings={ - "list_tools": {"ttl": 100}, - "call_tool": {"enabled": False}, - }, - default_ttl=1000, - ) - - tool_list_settings = middleware._get_cache_settings( - context=MiddlewareContext(method="tools/list", message=MagicMock()) - ) - assert tool_list_settings == {"ttl": 100} - - call_tool_settings = middleware._get_cache_settings( - context=MiddlewareContext(method="tools/call", message=MagicMock()) - ) - assert call_tool_settings == {"enabled": False} - - other_methods = [ - "resources/list", - "prompts/list", - "resources/read", - "prompts/get", - ] - for method in other_methods: - cache_settings = middleware._get_cache_settings( - context=MiddlewareContext(method=method, message=MagicMock()) - ) - assert cache_settings is None - - should_bypass = middleware._should_bypass_caching( - context=MiddlewareContext(method=method, message=MagicMock()) - ) - assert should_bypass - - def test_cache_key_generation(self): - """Test cache key generation.""" - from fastmcp.server.middleware.caching import ( - _make_call_tool_cache_key, - _make_get_prompt_cache_key, - _make_read_resource_cache_key, - ) - - msg = mcp.types.CallToolRequestParams( - name="test_tool", arguments={"param1": "value1", "param2": 42} - ) - - key = _make_call_tool_cache_key(msg) - - assert key == snapshot('test_tool:{"param1":"value1","param2":42}') - - msg = mcp.types.ReadResourceRequestParams( - uri=AnyUrl("https://test_uri"), - ) - - key = _make_read_resource_cache_key(msg) - - assert key == snapshot("https://test_uri/") - - msg = mcp.types.GetPromptRequestParams( - name="test_prompt", arguments={"param1": "value1"} - ) - - key = _make_get_prompt_cache_key(msg) - - assert key == snapshot('test_prompt:{"param1":"value1"}') + assert middleware1._matches_tool_cache_settings(tool_name=tool_name) is result class TestResponseCachingMiddlewareIntegration: @@ -393,7 +283,7 @@ class TestResponseCachingMiddlewareIntegration: async def caching_server( self, tracking_calculator: TrackingCalculator, - request, + request: pytest.FixtureRequest, ): """Create a FastMCP server for caching tests.""" mcp = FastMCP("CachingTestServer") @@ -401,8 +291,7 @@ class TestResponseCachingMiddlewareIntegration: with tempfile.TemporaryDirectory() as temp_dir: disk_store = DiskStore(directory=temp_dir) response_caching_middleware = ResponseCachingMiddleware( - cache_store=disk_store if request.param == "disk" else MemoryStore(), - max_item_size=100000, # 100kb + cache_storage=disk_store if request.param == "disk" else MemoryStore(), ) mcp.add_middleware(middleware=response_caching_middleware) @@ -429,15 +318,15 @@ class TestResponseCachingMiddlewareIntegration: async with Client(caching_server) as client: pre_tool_list: list[mcp.types.Tool] = await client.list_tools() - assert len(pre_tool_list) == 4 + assert len(pre_tool_list) == 5 # Add a tool and make sure it's missing from the list tool response - caching_server.add_tool( + _ = 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) == 4 + assert len(post_tool_list) == 5 assert pre_tool_list == post_tool_list @@ -449,7 +338,7 @@ class TestResponseCachingMiddlewareIntegration: """Test that caching works with a real FastMCP server.""" tracking_calculator.add_tools(fastmcp=caching_server) - async with Client[FastMCPTransport](caching_server) as client: + async with Client[FastMCPTransport](transport=caching_server) as client: call_tool_result_one: CallToolResult = await client.call_tool( "add", {"a": 5, "b": 3} ) @@ -468,7 +357,7 @@ class TestResponseCachingMiddlewareIntegration: """Test that caching works with a real FastMCP server.""" tracking_calculator.add_tools(fastmcp=caching_server) - async with Client[FastMCPTransport](caching_server) as client: + async with Client[FastMCPTransport](transport=caching_server) as client: call_tool_result_one: CallToolResult = await client.call_tool( "very_large_response", {} ) @@ -480,6 +369,27 @@ class TestResponseCachingMiddlewareIntegration: 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 ): @@ -552,3 +462,46 @@ class TestResponseCachingMiddlewareIntegration: ) assert pre_prompt == post_prompt + + async def test_statistics( + self, + caching_server: FastMCP, + ): + """Test that statistics are collected correctly.""" + caching_middleware = caching_server.middleware[0] + assert isinstance(caching_middleware, ResponseCachingMiddleware) + + async with Client[FastMCPTransport](transport=caching_server) as client: + statistics = caching_middleware.statistics() + assert statistics == snapshot(ResponseCachingStatistics()) + + _ = await client.call_tool("add", {"a": 5, "b": 3}) + + statistics = caching_middleware.statistics() + assert statistics == snapshot( + ResponseCachingStatistics( + list_tools=KVStoreCollectionStatistics( + get=GetStatistics(count=2, hit=1, miss=1), + put=PutStatistics(count=1), + ), + call_tool=KVStoreCollectionStatistics( + get=GetStatistics(count=1, miss=1), put=PutStatistics(count=1) + ), + ) + ) + + _ = await client.call_tool("add", {"a": 5, "b": 3}) + + statistics = caching_middleware.statistics() + assert statistics == snapshot( + ResponseCachingStatistics( + list_tools=KVStoreCollectionStatistics( + get=GetStatistics(count=2, hit=1, miss=1), + put=PutStatistics(count=1), + ), + call_tool=KVStoreCollectionStatistics( + get=GetStatistics(count=2, hit=1, miss=1), + put=PutStatistics(count=1), + ), + ) + )