mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-09 07:09:11 +02:00
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:
parent
7ca58583fd
commit
4ef7d16419
2 changed files with 164 additions and 15 deletions
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue