mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-15 01:59:10 +02:00
123 lines
4.4 KiB
Python
123 lines
4.4 KiB
Python
"""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"
|