Keep multi-round-trip asks out of the response cache

Prompt and resource asks carry no content, so caching one stored an empty
result and the client never saw the question. Bypass the cache on
continuation legs and return asks unwrapped, as tool calls already did.
This commit is contained in:
Jeremiah Lowin 2026-07-26 16:06:17 -04:00
commit 4ef7d16419
No known key found for this signature in database
2 changed files with 164 additions and 15 deletions

View file

@ -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(

View file

@ -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"