From c4b53ce141a3d456ff2218a82b95ca33d85462b2 Mon Sep 17 00:00:00 2001 From: "claude[bot]" <41898282+claude[bot]@users.noreply.github.com> Date: Sun, 16 Nov 2025 21:04:32 +0000 Subject: [PATCH] Use custom cacheable models in middleware to preserve key field Instead of adding custom serializers to FastMCPComponent (which affects equality globally), use dedicated cacheable models in the middleware. This approach: - Isolates caching concerns to the middleware - Follows existing patterns (CachableReadResourceContents, CachableToolResult) - Preserves key field through serialization without side effects - Maintains architectural separation Co-authored-by: William Easton --- src/fastmcp/server/middleware/caching.py | 214 ++++++++++++++++------- src/fastmcp/utilities/components.py | 43 +---- 2 files changed, 155 insertions(+), 102 deletions(-) diff --git a/src/fastmcp/server/middleware/caching.py b/src/fastmcp/server/middleware/caching.py index 58392eb26..ba3f31127 100644 --- a/src/fastmcp/server/middleware/caching.py +++ b/src/fastmcp/server/middleware/caching.py @@ -74,6 +74,132 @@ class CachableToolResult(BaseModel): ) +class CachableTool(BaseModel): + """A wrapper for Tool that can be cached, preserving the key field.""" + + name: str + key: str + title: str | None + description: str | None + parameters: dict[str, Any] + output_schema: dict[str, Any] | None + annotations: Any | None + meta: dict[str, Any] | None + tags: set[str] + enabled: bool + + @classmethod + def wrap(cls, tool: Tool) -> Self: + return cls( + name=tool.name, + key=tool.key, + 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, + ) + + def unwrap(self) -> Tool: + return Tool( + name=self.name, + key=self.key, + title=self.title, + description=self.description, + parameters=self.parameters, + output_schema=self.output_schema, + annotations=self.annotations, + meta=self.meta, + tags=self.tags, + enabled=self.enabled, + ) + + +class CachableResource(BaseModel): + """A wrapper for Resource that can be cached, preserving the key field.""" + + name: str + key: str + title: str | None + description: str | None + uri: str + mime_type: str + annotations: Any | None + meta: dict[str, Any] | None + tags: set[str] + enabled: bool + + @classmethod + def wrap(cls, resource: Resource) -> Self: + return cls( + name=resource.name, + key=resource.key, + title=resource.title, + description=resource.description, + uri=str(resource.uri), + mime_type=resource.mime_type, + annotations=resource.annotations, + meta=resource.meta, + tags=resource.tags, + enabled=resource.enabled, + ) + + def unwrap(self) -> Resource: + return Resource( + name=self.name, + key=self.key, + title=self.title, + description=self.description, + uri=self.uri, + mime_type=self.mime_type, + annotations=self.annotations, + meta=self.meta, + tags=self.tags, + enabled=self.enabled, + ) + + +class CachablePrompt(BaseModel): + """A wrapper for Prompt that can be cached, preserving the key field.""" + + name: str + key: str + title: str | None + description: str | None + arguments: list[Any] | None + meta: dict[str, Any] | None + tags: set[str] + enabled: bool + + @classmethod + def wrap(cls, prompt: Prompt) -> Self: + return cls( + name=prompt.name, + key=prompt.key, + title=prompt.title, + description=prompt.description, + arguments=prompt.arguments, + meta=prompt.meta, + tags=prompt.tags, + enabled=prompt.enabled, + ) + + def unwrap(self) -> Prompt: + return Prompt( + name=self.name, + key=self.key, + title=self.title, + description=self.description, + arguments=self.arguments, + meta=self.meta, + tags=self.tags, + enabled=self.enabled, + ) + + class SharedMethodSettings(TypedDict): """Shared config for a cache method.""" @@ -182,22 +308,26 @@ class ResponseCachingMiddleware(Middleware): call_tool_settings or CallToolSettings() ) - self._list_tools_cache: PydanticAdapter[list[Tool]] = PydanticAdapter( + self._list_tools_cache: PydanticAdapter[list[CachableTool]] = PydanticAdapter( key_value=self._stats, - pydantic_model=list[Tool], + pydantic_model=list[CachableTool], 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_resources_cache: PydanticAdapter[list[CachableResource]] = ( + PydanticAdapter( + key_value=self._stats, + pydantic_model=list[CachableResource], + 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._list_prompts_cache: PydanticAdapter[list[CachablePrompt]] = ( + PydanticAdapter( + key_value=self._stats, + pydantic_model=list[CachablePrompt], + default_collection="prompts/list", + ) ) self._read_resource_cache: PydanticAdapter[ @@ -234,26 +364,12 @@ class ResponseCachingMiddleware(Middleware): return await call_next(context) if cached_value := await self._list_tools_cache.get(key=GLOBAL_KEY): - return cached_value + return [item.unwrap() for item in cached_value] tools: Sequence[Tool] = await call_next(context=context) - # Turn any subclass of Tool into a Tool - cachable_tools: list[Tool] = [ - Tool( - name=tool.name, - key=tool.key, - 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 - ] + # Wrap tools in cacheable models + cachable_tools: list[CachableTool] = [CachableTool.wrap(tool) for tool in tools] await self._list_tools_cache.put( key=GLOBAL_KEY, @@ -261,7 +377,7 @@ class ResponseCachingMiddleware(Middleware): ttl=self._list_tools_settings.get("ttl", FIVE_MINUTES_IN_SECONDS), ) - return cachable_tools + return [item.unwrap() for item in cachable_tools] @override async def on_list_resources( @@ -275,25 +391,13 @@ class ResponseCachingMiddleware(Middleware): return await call_next(context) if cached_value := await self._list_resources_cache.get(key=GLOBAL_KEY): - return cached_value + return [item.unwrap() for item in cached_value] resources: Sequence[Resource] = await call_next(context=context) - # Turn any subclass of Resource into a Resource - cachable_resources: list[Resource] = [ - Resource( - name=resource.name, - key=resource.key, - 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 + # Wrap resources in cacheable models + cachable_resources: list[CachableResource] = [ + CachableResource.wrap(resource) for resource in resources ] await self._list_resources_cache.put( @@ -302,7 +406,7 @@ class ResponseCachingMiddleware(Middleware): ttl=self._list_resources_settings.get("ttl", FIVE_MINUTES_IN_SECONDS), ) - return cachable_resources + return [item.unwrap() for item in cachable_resources] @override async def on_list_prompts( @@ -316,23 +420,13 @@ class ResponseCachingMiddleware(Middleware): return await call_next(context) if cached_value := await self._list_prompts_cache.get(key=GLOBAL_KEY): - return cached_value + return [item.unwrap() for item in cached_value] prompts: Sequence[Prompt] = await call_next(context=context) - # Turn any subclass of Prompt into a Prompt - cachable_prompts: list[Prompt] = [ - Prompt( - name=prompt.name, - key=prompt.key, - title=prompt.title, - description=prompt.description, - tags=prompt.tags, - meta=prompt.meta, - enabled=prompt.enabled, - arguments=prompt.arguments, - ) - for prompt in prompts + # Wrap prompts in cacheable models + cachable_prompts: list[CachablePrompt] = [ + CachablePrompt.wrap(prompt) for prompt in prompts ] await self._list_prompts_cache.put( @@ -341,7 +435,7 @@ class ResponseCachingMiddleware(Middleware): ttl=self._list_prompts_settings.get("ttl", FIVE_MINUTES_IN_SECONDS), ) - return cachable_prompts + return [item.unwrap() for item in cachable_prompts] @override async def on_call_tool( diff --git a/src/fastmcp/utilities/components.py b/src/fastmcp/utilities/components.py index e31506d90..91e0777b0 100644 --- a/src/fastmcp/utilities/components.py +++ b/src/fastmcp/utilities/components.py @@ -4,7 +4,7 @@ from collections.abc import Sequence from typing import Annotated, Any, TypedDict from mcp.types import Icon -from pydantic import BeforeValidator, Field, PrivateAttr, model_serializer +from pydantic import BeforeValidator, Field, PrivateAttr from typing_extensions import Self, TypeVar import fastmcp @@ -119,47 +119,6 @@ class FastMCPComponent(FastMCPBaseModel): copy._key = key return copy - @model_serializer(mode="wrap") - def _serialize_model(self, serializer: Any, info: Any) -> dict[str, Any]: - """Custom serializer to include the key field.""" - data = serializer(self) - # Include _key in serialization if it's set - if self._key is not None: - data["key"] = self._key - return data - - @classmethod - def model_validate( - cls, - obj: Any, - *, - strict: bool | None = None, - from_attributes: bool | None = None, - context: dict[str, Any] | None = None, - ) -> Self: - """Validate and create a model instance, handling the key attribute.""" - if isinstance(obj, dict): - # Extract key from dict if present (don't mutate original dict) - key = obj.get("key") - if key is not None: - # Create a copy without the key field - obj = {k: v for k, v in obj.items() if k != "key"} - instance = super().model_validate( - obj, - strict=strict, - from_attributes=from_attributes, - context=context, - ) - if key is not None: - instance._key = key - return instance - return super().model_validate( - obj, - strict=strict, - from_attributes=from_attributes, - context=context, - ) - def __eq__(self, other: object) -> bool: if type(self) is not type(other): return False