diff --git a/src/fastmcp/server/middleware/caching.py b/src/fastmcp/server/middleware/caching.py index 35565de36..fb61d4f96 100644 --- a/src/fastmcp/server/middleware/caching.py +++ b/src/fastmcp/server/middleware/caching.py @@ -1,11 +1,11 @@ """A middleware for response caching.""" -import hashlib import json from collections.abc import Sequence from typing import Any, TypedDict, TypeVar, cast 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 @@ -598,22 +598,19 @@ 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.""" - raw = f"{msg.name}:{_get_arguments_str(msg.arguments)}" - return hashlib.sha256(raw.encode("utf-8")).hexdigest() + 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.""" - raw = f"{msg.uri}" - return hashlib.sha256(raw.encode("utf-8")).hexdigest() + 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.""" - raw = f"{msg.name}:{_get_arguments_str(msg.arguments)}" - return hashlib.sha256(raw.encode("utf-8")).hexdigest() + return f"{msg.name}:{_get_arguments_str(msg.arguments)}" def _get_arguments_str(arguments: dict[str, Any] | None) -> str: @@ -623,7 +620,7 @@ def _get_arguments_str(arguments: dict[str, Any] | None) -> str: return "null" try: - return json.dumps(arguments, sort_keys=True, separators=(",", ":")) + return pydantic_core.to_json(value=arguments, fallback=str).decode() except TypeError: return repr(arguments) @@ -651,21 +648,3 @@ def _get_size_of_tool_result(value: ToolResult) -> int: ) return content_size + structured_content_size - - -# def get_size_of_one_value(value: BaseModel | ToolResult | ReadResourceContents) -> int: -# """Get the size of an mcp type.""" - -# if isinstance(value, ToolResult): -# return get_size_of_tool_result(value) -# if isinstance(value, ReadResourceContents): -# return len(value.content) -# return len(value.model_dump_json()) - - -# def get_size_of_value(value: CachableListTypes) -> int: -# """Get the size of a cache entry.""" -# if isinstance(value, (BaseModel | ToolResult | ReadResourceContents)): -# return get_size_of_one_value(value) - -# return sum(get_size_of_one_value(item) for item in value) diff --git a/tests/server/middleware/test_caching.py b/tests/server/middleware/test_caching.py index fd7701ae3..7701f2735 100644 --- a/tests/server/middleware/test_caching.py +++ b/tests/server/middleware/test_caching.py @@ -367,11 +367,7 @@ class TestResponseCachingMiddleware: key = _make_call_tool_cache_key(msg) - # Should be a SHA256 hash - assert len(key) == 64 - assert key == snapshot( - "7fa3d5c7967a202457eeca0731709fd87ec98546ddaee829ee86ca54b1858c59" - ) + assert key == snapshot('test_tool:{"param1":"value1","param2":42}') msg = mcp.types.ReadResourceRequestParams( uri=AnyUrl("https://test_uri"), @@ -379,10 +375,7 @@ class TestResponseCachingMiddleware: key = _make_read_resource_cache_key(msg) - assert len(key) == 64 - assert key == snapshot( - "e34cc47c03ed1ad54f02501d95ecc463b65646568961c97ca4b730cb274e9d42" - ) + assert key == snapshot("https://test_uri/") msg = mcp.types.GetPromptRequestParams( name="test_prompt", arguments={"param1": "value1"} @@ -390,10 +383,7 @@ class TestResponseCachingMiddleware: key = _make_get_prompt_cache_key(msg) - assert len(key) == 64 - assert key == snapshot( - "6306ff84fd3ff247a4bd91271e9d727d7f051bba53fb2e3bf80958988c4baf57" - ) + assert key == snapshot('test_prompt:{"param1":"value1"}') class TestResponseCachingMiddlewareIntegration: