From d346dc66e71f8a71808effca7010a0b968011180 Mon Sep 17 00:00:00 2001 From: "claude[bot]" <41898282+claude[bot]@users.noreply.github.com> Date: Wed, 29 Oct 2025 15:24:33 +0000 Subject: [PATCH] fix: preserve key field in ResponseCachingMiddleware for prefixed tools/resources/prompts The ResponseCachingMiddleware was losing prefix information when caching tools, resources, and prompts from mounted/imported servers. The root cause was that the private _key attribute wasn't being serialized by Pydantic. Changes: - Add model_serializer to FastMCPComponent to include _key in serialization - Update model_validate to restore _key from deserialized data - Add key parameter when creating cached Tool/Resource/Prompt objects - Add comprehensive test for mounted server prefix preservation Fixes #2300 Co-authored-by: William Easton --- src/fastmcp/server/middleware/caching.py | 3 + src/fastmcp/utilities/components.py | 43 +++++++++++++- tests/server/middleware/test_caching.py | 76 ++++++++++++++++++++++++ 3 files changed, 121 insertions(+), 1 deletion(-) diff --git a/src/fastmcp/server/middleware/caching.py b/src/fastmcp/server/middleware/caching.py index 52540248e..58392eb26 100644 --- a/src/fastmcp/server/middleware/caching.py +++ b/src/fastmcp/server/middleware/caching.py @@ -242,6 +242,7 @@ class ResponseCachingMiddleware(Middleware): cachable_tools: list[Tool] = [ Tool( name=tool.name, + key=tool.key, title=tool.title, description=tool.description, parameters=tool.parameters, @@ -282,6 +283,7 @@ class ResponseCachingMiddleware(Middleware): cachable_resources: list[Resource] = [ Resource( name=resource.name, + key=resource.key, title=resource.title, description=resource.description, tags=resource.tags, @@ -322,6 +324,7 @@ class ResponseCachingMiddleware(Middleware): cachable_prompts: list[Prompt] = [ Prompt( name=prompt.name, + key=prompt.key, title=prompt.title, description=prompt.description, tags=prompt.tags, diff --git a/src/fastmcp/utilities/components.py b/src/fastmcp/utilities/components.py index 91e0777b0..e31506d90 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 +from pydantic import BeforeValidator, Field, PrivateAttr, model_serializer from typing_extensions import Self, TypeVar import fastmcp @@ -119,6 +119,47 @@ 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 diff --git a/tests/server/middleware/test_caching.py b/tests/server/middleware/test_caching.py index 696933530..26dfcc3ac 100644 --- a/tests/server/middleware/test_caching.py +++ b/tests/server/middleware/test_caching.py @@ -505,3 +505,79 @@ class TestResponseCachingMiddlewareIntegration: ), ) ) + + async def test_mounted_server_prefixes_preserved(self): + """Test that caching preserves prefixes from mounted servers.""" + # Create child servers with tools, resources, and prompts + child = FastMCP("child") + calculator = TrackingCalculator() + calculator.add_tools(fastmcp=child) + calculator.add_resources(fastmcp=child) + calculator.add_prompts(fastmcp=child) + + # Create parent with caching middleware + parent = FastMCP("parent") + parent.add_middleware(ResponseCachingMiddleware()) + await parent.import_server(child, prefix="child") + + async with Client[FastMCPTransport](transport=parent) as client: + # First call - populates cache + tools1 = await client.list_tools() + tool_names1 = [tool.name for tool in tools1] + + # Second call - from cache (this is where the bug would occur) + tools2 = await client.list_tools() + tool_names2 = [tool.name for tool in tools2] + + # All tools should have the prefix in both calls + for name in tool_names1: + assert name.startswith("child_"), ( + f"Tool {name} missing prefix (first call)" + ) + for name in tool_names2: + assert name.startswith("child_"), ( + f"Tool {name} missing prefix (cached call)" + ) + + # Both calls should return the same tools + assert tool_names1 == tool_names2 + + # Verify tool can be called with prefixed name + result = await client.call_tool("child_add", {"a": 5, "b": 3}) + assert not result.is_error + + # Test resources + resources1 = await client.list_resources() + resource_names1 = [resource.name for resource in resources1] + + resources2 = await client.list_resources() + resource_names2 = [resource.name for resource in resources2] + + for name in resource_names1: + assert name.startswith("child_"), ( + f"Resource {name} missing prefix (first call)" + ) + for name in resource_names2: + assert name.startswith("child_"), ( + f"Resource {name} missing prefix (cached call)" + ) + + assert resource_names1 == resource_names2 + + # Test prompts + prompts1 = await client.list_prompts() + prompt_names1 = [prompt.name for prompt in prompts1] + + prompts2 = await client.list_prompts() + prompt_names2 = [prompt.name for prompt in prompts2] + + for name in prompt_names1: + assert name.startswith("child_"), ( + f"Prompt {name} missing prefix (first call)" + ) + for name in prompt_names2: + assert name.startswith("child_"), ( + f"Prompt {name} missing prefix (cached call)" + ) + + assert prompt_names1 == prompt_names2