"""A middleware for response caching.""" import hashlib from collections.abc import Sequence 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 pydantic import Field from typing_extensions import NotRequired, Self, override from fastmcp.prompts.base import Message, Prompt, PromptResult from fastmcp.resources.base import Resource, ResourceContent, ResourceResult from fastmcp.server.dependencies import get_access_token from fastmcp.server.middleware.middleware import CallNext, Middleware, MiddlewareContext from fastmcp.tools.base import Tool, ToolResult from fastmcp.utilities.logging import get_logger from fastmcp.utilities.types import FastMCPBaseModel logger: Logger = get_logger(name=__name__) # Constants ONE_HOUR_IN_SECONDS = 3600 FIVE_MINUTES_IN_SECONDS = 300 ONE_MB_IN_BYTES = 1024 * 1024 ANONYMOUS_AUTH_KEY = "__anonymous__" class CachableResourceContent(FastMCPBaseModel): """A wrapper for ResourceContent that can be cached.""" content: str | bytes mime_type: str | None = None meta: dict[str, Any] | None = None class CachableResourceResult(FastMCPBaseModel): """A wrapper for ResourceResult that can be cached.""" contents: list[CachableResourceContent] meta: dict[str, Any] | None = None def get_size(self) -> int: return len(self.model_dump_json()) @classmethod def wrap(cls, value: ResourceResult) -> Self: return cls( contents=[ CachableResourceContent( content=item.content, mime_type=item.mime_type, meta=item.meta ) for item in value.contents ], meta=value.meta, ) def unwrap(self) -> ResourceResult: return ResourceResult( contents=[ ResourceContent( content=item.content, mime_type=item.mime_type, meta=item.meta ) for item in self.contents ], meta=self.meta, ) class CachableToolResult(FastMCPBaseModel): content: list[mcp.types.ContentBlock] structured_content: dict[str, Any] | None meta: dict[str, Any] | None is_error: bool = False @classmethod def wrap(cls, value: ToolResult) -> Self: return cls( content=value.content, structured_content=value.structured_content, meta=value.meta, is_error=value.is_error, ) def unwrap(self) -> ToolResult: return ToolResult( content=self.content, structured_content=self.structured_content, meta=self.meta, is_error=self.is_error, ) class CachableMessage(FastMCPBaseModel): """A wrapper for Message that can be cached.""" role: str content: ( mcp.types.TextContent | mcp.types.ImageContent | mcp.types.AudioContent | mcp.types.EmbeddedResource ) class CachablePromptResult(FastMCPBaseModel): """A wrapper for PromptResult that can be cached.""" messages: list[CachableMessage] description: str | None = None meta: dict[str, Any] | None = None def get_size(self) -> int: return len(self.model_dump_json()) @classmethod def wrap(cls, value: PromptResult) -> Self: return cls( messages=[ CachableMessage(role=m.role, content=m.content) for m in value.messages ], description=value.description, meta=value.meta, ) def unwrap(self) -> PromptResult: return PromptResult( messages=[ Message(content=m.content, role=m.role) # type: ignore[arg-type] # ty:ignore[invalid-argument-type] for m in self.messages ], description=self.description, meta=self.meta, ) class SharedMethodSettings(TypedDict): """Shared config for a cache method.""" ttl: NotRequired[int] enabled: NotRequired[bool] class ListToolsSettings(SharedMethodSettings): """Configuration options for Tool-related caching.""" class ListResourcesSettings(SharedMethodSettings): """Configuration options for Resource-related caching.""" class ListPromptsSettings(SharedMethodSettings): """Configuration options for Prompt-related caching.""" class CallToolSettings(SharedMethodSettings): """Configuration options for Tool-related caching.""" included_tools: NotRequired[list[str]] excluded_tools: NotRequired[list[str]] class ReadResourceSettings(SharedMethodSettings): """Configuration options for Resource-related caching.""" class GetPromptSettings(SharedMethodSettings): """Configuration options for Prompt-related caching.""" class ResponseCachingStatistics(FastMCPBaseModel): list_tools: KVStoreCollectionStatistics | None = Field(default=None) list_resources: KVStoreCollectionStatistics | None = Field(default=None) list_prompts: KVStoreCollectionStatistics | None = Field(default=None) read_resource: KVStoreCollectionStatistics | None = Field(default=None) get_prompt: KVStoreCollectionStatistics | None = Field(default=None) call_tool: KVStoreCollectionStatistics | None = Field(default=None) class ResponseCachingMiddleware(Middleware): """The response caching middleware offers a simple way to cache responses to mcp methods. The Middleware supports cache invalidation via notifications from the server. The Middleware implements TTL-based caching but cache implementations may offer additional features like LRU eviction, size limits, and more. When items are retrieved from the cache they will no longer be the original objects, but rather no-op objects this means that response caching may not be compatible with other middleware that expects original subclasses. Notes: - Caches `tools/call`, `resources/read`, `prompts/get`, `tools/list`, `resources/list`, and `prompts/list` requests. - Cache keys are derived from the method name, arguments, and the caller's access token. Entries are partitioned per-token so that responses filtered by per-component authorization (e.g. `auth=require_scopes(...)`) cannot leak across users with different permissions. Unauthenticated callers (including STDIO) share a single anonymous partition. """ def __init__( self, 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_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._backend: AsyncKeyValue = cache_storage or MemoryStore() # 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._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=list[Tool], default_collection="tools/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[list[Prompt]] = PydanticAdapter( key_value=self._stats, pydantic_model=list[Prompt], default_collection="prompts/list", ) self._read_resource_cache: PydanticAdapter[CachableResourceResult] = ( PydanticAdapter( key_value=self._stats, pydantic_model=CachableResourceResult, default_collection="resources/read", ) ) self._get_prompt_cache: PydanticAdapter[CachablePromptResult] = 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, default_collection="tools/call", ) @override async def on_list_tools( self, context: MiddlewareContext[mcp.types.ListToolsRequest], 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) cache_key: str = _get_auth_partition_key() if cached_value := await self._list_tools_cache.get(key=cache_key): return cached_value tools: Sequence[Tool] = await call_next(context) # 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, ) for tool in tools ] await self._list_tools_cache.put( key=cache_key, value=cachable_tools, ttl=self._list_tools_settings.get("ttl", FIVE_MINUTES_IN_SECONDS), ) return cachable_tools @override async def on_list_resources( self, context: MiddlewareContext[mcp.types.ListResourcesRequest], 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._list_resources_settings.get("enabled") is False: return await call_next(context) cache_key: str = _get_auth_partition_key() if cached_value := await self._list_resources_cache.get(key=cache_key): return cached_value resources: Sequence[Resource] = await call_next(context) # 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, uri=resource.uri, ) for resource in resources ] await self._list_resources_cache.put( key=cache_key, value=cachable_resources, ttl=self._list_resources_settings.get("ttl", FIVE_MINUTES_IN_SECONDS), ) return cachable_resources @override async def on_list_prompts( self, context: MiddlewareContext[mcp.types.ListPromptsRequest], 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._list_prompts_settings.get("enabled") is False: return await call_next(context) cache_key: str = _get_auth_partition_key() if cached_value := await self._list_prompts_cache.get(key=cache_key): return cached_value prompts: Sequence[Prompt] = await call_next(context) # 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, arguments=prompt.arguments, ) for prompt in prompts ] await self._list_prompts_cache.put( key=cache_key, value=cachable_prompts, ttl=self._list_prompts_settings.get("ttl", FIVE_MINUTES_IN_SECONDS), ) return cachable_prompts @override async def on_call_tool( self, context: MiddlewareContext[mcp.types.CallToolRequestParams], call_next: CallNext[mcp.types.CallToolRequestParams, ToolResult], ) -> 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.""" 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) cache_key: str = _make_call_tool_cache_key( msg=context.message, auth_key=_get_auth_partition_key() ) if cached_value := await self._call_tool_cache.get(key=cache_key): return cached_value.unwrap() tool_result: ToolResult = await call_next(context) cachable_tool_result: CachableToolResult = CachableToolResult.wrap( value=tool_result ) await self._call_tool_cache.put( key=cache_key, value=cachable_tool_result, ttl=self._call_tool_settings.get("ttl", ONE_HOUR_IN_SECONDS), ) return cachable_tool_result.unwrap() @override async def on_read_resource( self, context: MiddlewareContext[mcp.types.ReadResourceRequestParams], call_next: CallNext[mcp.types.ReadResourceRequestParams, ResourceResult], ) -> ResourceResult: """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._read_resource_settings.get("enabled") is False: return await call_next(context) cache_key: str = _make_read_resource_cache_key( msg=context.message, auth_key=_get_auth_partition_key() ) cached_value: CachableResourceResult | None if cached_value := await self._read_resource_cache.get(key=cache_key): return cached_value.unwrap() value: ResourceResult = await call_next(context) cached_value = CachableResourceResult.wrap(value) await self._read_resource_cache.put( key=cache_key, value=cached_value, ttl=self._read_resource_settings.get("ttl", ONE_HOUR_IN_SECONDS), ) return cached_value.unwrap() @override async def on_get_prompt( self, context: MiddlewareContext[mcp.types.GetPromptRequestParams], call_next: CallNext[mcp.types.GetPromptRequestParams, PromptResult], ) -> PromptResult: """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._get_prompt_settings.get("enabled") is False: return await call_next(context) cache_key: str = _make_get_prompt_cache_key( msg=context.message, auth_key=_get_auth_partition_key() ) if cached_value := await self._get_prompt_cache.get(key=cache_key): return cached_value.unwrap() value: PromptResult = await call_next(context) cached_value = CachablePromptResult.wrap(value) await self._get_prompt_cache.put( key=cache_key, value=cached_value, ttl=self._get_prompt_settings.get("ttl", ONE_HOUR_IN_SECONDS), ) return cached_value.unwrap() def _matches_tool_cache_settings(self, tool_name: str) -> bool: """Check if the tool matches the cache settings for tool calls.""" if included_tools := self._call_tool_settings.get("included_tools"): if tool_name not in included_tools: return False if excluded_tools := self._call_tool_settings.get("excluded_tools"): if tool_name in excluded_tools: return False return True 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"), list_prompts=self._stats.statistics.collections.get("prompts/list"), read_resource=self._stats.statistics.collections.get("resources/read"), get_prompt=self._stats.statistics.collections.get("prompts/get"), call_tool=self._stats.statistics.collections.get("tools/call"), ) def _get_arguments_str(arguments: dict[str, Any] | None) -> str: """Get a string representation of the arguments.""" if arguments is None: return "null" try: return pydantic_core.to_json(value=arguments, fallback=str).decode() except TypeError: return repr(arguments) def _hash_cache_key(value: str) -> str: """Build a fixed-length SHA-256 cache key from request-derived input.""" return hashlib.sha256(value.encode()).hexdigest() def _get_auth_partition_key() -> str: """Return a stable, hashed identifier for the current access token. Cache entries are partitioned by access token so that responses filtered by per-component authorization (e.g. `auth=require_scopes(...)`) are not leaked across users with different permissions. Unauthenticated callers (including STDIO) share a single anonymous partition. """ token = get_access_token() if token is None: return ANONYMOUS_AUTH_KEY return _hash_cache_key(token.token) def _make_call_tool_cache_key( msg: mcp.types.CallToolRequestParams, auth_key: str = ANONYMOUS_AUTH_KEY ) -> str: """Make a cache key for a tool call using a stable hash of name and arguments.""" return _hash_cache_key(f"{auth_key}:{msg.name}:{_get_arguments_str(msg.arguments)}") def _make_read_resource_cache_key( msg: mcp.types.ReadResourceRequestParams, auth_key: str = ANONYMOUS_AUTH_KEY ) -> str: """Make a cache key for a resource read using a stable hash of URI.""" return _hash_cache_key(f"{auth_key}:{msg.uri}") def _make_get_prompt_cache_key( msg: mcp.types.GetPromptRequestParams, auth_key: str = ANONYMOUS_AUTH_KEY ) -> str: """Make a cache key for a prompt get using a stable hash of name and arguments.""" return _hash_cache_key(f"{auth_key}:{msg.name}:{_get_arguments_str(msg.arguments)}")