diff --git a/fastmcp_slim/fastmcp/server/middleware/caching.py b/fastmcp_slim/fastmcp/server/middleware/caching.py index cbf48b849..d470f1b90 100644 --- a/fastmcp_slim/fastmcp/server/middleware/caching.py +++ b/fastmcp_slim/fastmcp/server/middleware/caching.py @@ -18,8 +18,18 @@ from key_value.aio.wrappers.statistics.wrapper import ( from pydantic import Field from typing_extensions import NotRequired, Self, TypeVar, override -from fastmcp.prompts.base import Message, Prompt, PromptResult -from fastmcp.resources.base import Resource, ResourceContent, ResourceResult +from fastmcp.prompts.base import ( + InputRequiredPromptResult, + Message, + Prompt, + PromptResult, +) +from fastmcp.resources.base import ( + InputRequiredResourceResult, + 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 InputRequiredToolResult, Tool, ToolResult @@ -28,6 +38,28 @@ from fastmcp.utilities.types import FastMCPBaseModel logger: Logger = get_logger(name=__name__) + +def _is_continuation_leg(context: MiddlewareContext[Any]) -> bool: + """Whether this request is answering a previous round's ask (SEP-2322). + + A continuation must bypass the cache entirely. Cache keys are built from the + component's identity and arguments alone, so a continuation shares its key + with a fresh call: reading could serve a prior flow's final answer to this + leg, and writing would serve THIS flow's final answer to a later fresh call, + which would then never be asked the questions at all. + + Either signal marks a continuation. A state-only round (one that carried + `request_state` without asking anything) retries with `input_responses` + still `None`. + """ + fastmcp_ctx = context.fastmcp_context + if fastmcp_ctx is None: + return False + return ( + fastmcp_ctx.input_responses is not None or fastmcp_ctx.request_state is not None + ) + + # Constants ONE_HOUR_IN_SECONDS = 3600 FIVE_MINUTES_IN_SECONDS = 300 @@ -413,19 +445,7 @@ class ResponseCachingMiddleware(Middleware): ) is False or not self._matches_tool_cache_settings(tool_name=tool_name): return await call_next(context) - # A multi-round continuation leg (SEP-2322) must bypass the cache - # entirely: the cache key is built from the tool name and arguments - # only, so a continuation shares its key with a fresh call. Reading - # could serve a prior flow's final answer to this leg; writing would - # serve THIS flow's final answer to a later fresh call — which would - # then never be asked the tool's questions. Either signal marks a - # continuation: a state-only round (request_state with no questions) - # retries with input_responses=None but still carries request_state. - fastmcp_ctx = context.fastmcp_context - if fastmcp_ctx is not None and ( - fastmcp_ctx.input_responses is not None - or fastmcp_ctx.request_state is not None - ): + if _is_continuation_leg(context): return await call_next(context) cache_key: str = _make_call_tool_cache_key( @@ -477,6 +497,9 @@ class ResponseCachingMiddleware(Middleware): if self._read_resource_settings.get("enabled") is False: return await call_next(context) + if _is_continuation_leg(context): + return await call_next(context) + cache_key: str = _make_read_resource_cache_key( msg=context.message, auth_key=_get_auth_partition_key() ) @@ -486,6 +509,14 @@ class ResponseCachingMiddleware(Middleware): return cached_value.unwrap() value: ResourceResult = await call_next(context) + + # Never cache a multi-round-trip ask (SEP-2322). An + # InputRequiredResourceResult is a request for client input on this leg, + # not a stable answer, and it carries no contents — wrapping it would + # cache an empty read and the client would never see the question. + if isinstance(value, InputRequiredResourceResult): + return value + cached_value = CacheableResourceResult.wrap(value) await self._read_resource_cache.put( @@ -507,6 +538,9 @@ class ResponseCachingMiddleware(Middleware): if self._get_prompt_settings.get("enabled") is False: return await call_next(context) + if _is_continuation_leg(context): + return await call_next(context) + cache_key: str = _make_get_prompt_cache_key( msg=context.message, auth_key=_get_auth_partition_key() ) @@ -515,6 +549,14 @@ class ResponseCachingMiddleware(Middleware): return cached_value.unwrap() value: PromptResult = await call_next(context) + + # Never cache a multi-round-trip ask (SEP-2322). An + # InputRequiredPromptResult is a request for client input on this leg, + # not a stable answer, and it carries no messages — wrapping it would + # cache an empty prompt and the client would never see the question. + if isinstance(value, InputRequiredPromptResult): + return value + cached_value = CacheablePromptResult.wrap(value) await self._get_prompt_cache.put( diff --git a/tests/server/middleware/test_caching.py b/tests/server/middleware/test_caching.py index b45961bcf..dd1782a7e 100644 --- a/tests/server/middleware/test_caching.py +++ b/tests/server/middleware/test_caching.py @@ -902,3 +902,110 @@ class TestAuthAwareCaching: assert {p.name for p in prompts} == {"public_prompt"} finally: auth_context_var.reset(tok) + + +class TestCachingWithInputRequiredResults: + """A multi-round-trip ask must survive `ResponseCachingMiddleware`. + + An ask carries no content of its own, so caching one would store an empty + result and the client would never see the question. Continuation legs must + also bypass the cache, since they share a cache key with a fresh call. + """ + + @staticmethod + def _answer(responses: mcp_types.InputResponses) -> str: + """The accepted value for the single question these guards ask.""" + result = responses["q"] + assert isinstance(result, mcp_types.ElicitResult) + assert result.content is not None + return str(result.content["q"]) + + @staticmethod + def _ask() -> mcp_types.InputRequiredResult: + return mcp_types.InputRequiredResult( + result_type="input_required", + input_requests={ + "q": mcp_types.ElicitRequest( + method="elicitation/create", + params=mcp_types.ElicitRequestFormParams( + message="Which quarter?", + requested_schema={ + "type": "object", + "properties": {"q": {"type": "string"}}, + "required": ["q"], + }, + ), + ) + }, + ) + + @classmethod + def _server(cls) -> FastMCP: + mcp = FastMCP("cached-guards") + mcp.add_middleware(ResponseCachingMiddleware()) + + @mcp.tool + async def summarize_tool(ctx: Context) -> str | mcp_types.InputRequiredResult: + if ctx.input_responses is None: + return cls._ask() + return f"Summary for {cls._answer(ctx.input_responses)}" + + @mcp.prompt + async def summarize(ctx: Context) -> str | mcp_types.InputRequiredResult: + if ctx.input_responses is None: + return cls._ask() + return f"Summary for {cls._answer(ctx.input_responses)}" + + @mcp.resource("report://x") + async def report(ctx: Context) -> str | mcp_types.InputRequiredResult: + if ctx.input_responses is None: + return cls._ask() + return f"Report for {cls._answer(ctx.input_responses)}" + + return mcp + + @staticmethod + async def _handler(message, response_type, params, ctx): + from fastmcp.client.elicitation import ElicitResult + + return ElicitResult(action="accept", content=response_type(q="Q3")) + + async def test_tool_guard_completes_under_caching(self): + async with Client( + self._server(), mode="auto", elicitation_handler=self._handler + ) as client: + result = await client.call_tool("summarize_tool", {}) + + assert result.data == "Summary for Q3" + + async def test_prompt_guard_completes_under_caching(self): + async with Client( + self._server(), mode="auto", elicitation_handler=self._handler + ) as client: + result = await client.get_prompt("summarize") + + assert result.messages[0].content.text == "Summary for Q3" + + async def test_resource_guard_completes_under_caching(self): + async with Client( + self._server(), mode="auto", elicitation_handler=self._handler + ) as client: + result = await client.read_resource("report://x") + + assert result[0].text == "Report for Q3" + + async def test_guard_still_asks_on_a_second_fresh_call(self): + """The ask must not be cached away for the next caller. + + A second fresh flow has to be asked the same question; serving it a + cached final answer would skip the component's own per-round logic. + """ + mcp = self._server() + async with Client( + mcp, mode="auto", elicitation_handler=self._handler + ) as client: + first = await client.get_prompt("summarize") + second = await client.get_prompt("summarize") + + assert first.messages[0].content.text == "Summary for Q3" + assert second.messages[0].content.text == "Summary for Q3"