diff --git a/tests/server/middleware/test_caching.py b/tests/server/middleware/test_caching.py index a88efbd3f..b45961bcf 100644 --- a/tests/server/middleware/test_caching.py +++ b/tests/server/middleware/test_caching.py @@ -26,7 +26,6 @@ from pydantic import AnyUrl, BaseModel from fastmcp import Context, FastMCP from fastmcp.client.client import CallToolResult, Client -from fastmcp.client.elicitation import ElicitResult from fastmcp.client.transports import FastMCPTransport from fastmcp.prompts.base import Message, Prompt from fastmcp.prompts.function_prompt import FunctionPrompt @@ -903,94 +902,3 @@ 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: - params = mcp_types.ElicitRequestFormParams( - message="Which quarter?", - requested_schema={ - "type": "object", - "properties": {"q": {"type": "string"}}, - "required": ["q"], - }, - ) - request = mcp_types.ElicitRequest(method="elicitation/create", params=params) - return mcp_types.InputRequiredResult( - result_type="input_required", input_requests={"q": request} - ) - - @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): - return ElicitResult(action="accept", content=response_type(q="Q3")) - - def _client(self) -> Client: - return Client(self._server(), mode="auto", elicitation_handler=self._handler) - - async def test_tool_guard_completes_under_caching(self): - async with self._client() 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 self._client() 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 self._client() 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. - """ - async with self._client() 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" diff --git a/tests/server/middleware/test_caching_guards.py b/tests/server/middleware/test_caching_guards.py new file mode 100644 index 000000000..e552dc49e --- /dev/null +++ b/tests/server/middleware/test_caching_guards.py @@ -0,0 +1,123 @@ +"""Response caching around multi-round-trip asks (SEP-2322). + +A guard component answers a call by *returning* an `InputRequiredResult` — a +request for client input rather than a final answer. Two things follow for +`ResponseCachingMiddleware`, and they apply equally to tools, prompts, and +resources: + +- An ask must never be stored. It carries no content of its own, so caching one + writes an empty result, and every later caller is served that emptiness + instead of being asked the question. +- A continuation leg 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 hand this leg a prior flow's final answer, and + writing would hand a later fresh caller *this* flow's answer, skipping the + questions altogether. +""" + +import mcp_types + +from fastmcp import Context, FastMCP +from fastmcp.client.client import Client +from fastmcp.client.elicitation import ElicitResult +from fastmcp.server.middleware.caching import ResponseCachingMiddleware + + +def _ask() -> mcp_types.InputRequiredResult: + """The single-question ask every guard in this module returns.""" + params = mcp_types.ElicitRequestFormParams( + message="Which quarter?", + requested_schema={ + "type": "object", + "properties": {"q": {"type": "string"}}, + "required": ["q"], + }, + ) + request = mcp_types.ElicitRequest(method="elicitation/create", params=params) + return mcp_types.InputRequiredResult( + result_type="input_required", + input_requests={"q": request}, + ) + + +def _answer(responses: mcp_types.InputResponses) -> str: + """The accepted value for the question `_ask` poses.""" + result = responses["q"] + assert isinstance(result, mcp_types.ElicitResult) + assert result.content is not None + return str(result.content["q"]) + + +async def _handler(message, response_type, params, ctx): + """An elicitation handler that always answers "Q3".""" + return ElicitResult(action="accept", content=response_type(q="Q3")) + + +def cached_guard_server() -> FastMCP: + """A caching server whose tool, prompt, and resource are all guards.""" + 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 _ask() + return f"Summary for {_answer(ctx.input_responses)}" + + @mcp.prompt + async def summarize(ctx: Context) -> str | mcp_types.InputRequiredResult: + if ctx.input_responses is None: + return _ask() + return f"Summary for {_answer(ctx.input_responses)}" + + @mcp.resource("report://x") + async def report(ctx: Context) -> str | mcp_types.InputRequiredResult: + if ctx.input_responses is None: + return _ask() + return f"Report for {_answer(ctx.input_responses)}" + + return mcp + + +def guard_client() -> Client: + """A client that answers each round automatically.""" + return Client(cached_guard_server(), mode="auto", elicitation_handler=_handler) + + +class TestGuardsCompleteUnderCaching: + """Each component type drives its loop to a real answer with caching on.""" + + async def test_tool(self): + async with guard_client() as client: + result = await client.call_tool("summarize_tool", {}) + + assert result.data == "Summary for Q3" + + async def test_prompt(self): + async with guard_client() as client: + result = await client.get_prompt("summarize") + + assert result.messages[0].content.text == "Summary for Q3" + + async def test_resource(self): + async with guard_client() as client: + result = await client.read_resource("report://x") + + assert result[0].text == "Report for Q3" + + +class TestAsksAreNotCached: + """A stored ask would poison every later caller.""" + + async def test_second_fresh_flow_is_asked_again(self): + """A second fresh flow must be asked the same question. + + Serving it a cached final answer would skip the component's own + per-round logic — it would receive an answer it never supplied input for. + """ + async with guard_client() 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"