mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-09 15:19:10 +02:00
c3ad5e9ecbThree review fixes. A Docket worker may reuse an asyncio context across tasks, so snapshot restore now always resets auth and headers to the current task's state — an anonymous task following an authenticated one no longer inherits the prior caller's identity. A stalled in-task elicitation handler is now bounded by the call's remaining timeout, like polling and sleeps. And call_tool_task takes a version= to task a specific component version rather than the highest.
233 lines
8.2 KiB
Python
233 lines
8.2 KiB
Python
"""The transparent client task flow over a real in-memory connection.
|
|
|
|
A real `Client(server, mode="auto")` calls a `task=True` tool; the server runs it
|
|
as a task and answers `tools/call` with a `CreateTaskResult`; the client's
|
|
auto-registered tasks extension resolves it by polling `tasks/get` to completion.
|
|
The caller of `call_tool` sees only the tool's real result — never that the call
|
|
was tasked. This is the whole point of the client half.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from dataclasses import dataclass
|
|
|
|
import mcp_types
|
|
import pytest
|
|
from mcp.shared.exceptions import MCPError
|
|
|
|
from fastmcp import Context, FastMCP
|
|
from fastmcp.client import Client
|
|
from fastmcp.exceptions import ToolError
|
|
from fastmcp_tasks import TasksExtension, call_tool_task
|
|
|
|
|
|
@pytest.fixture
|
|
def task_server() -> FastMCP:
|
|
mcp = FastMCP("transparent-tasks")
|
|
mcp.add_extension(TasksExtension())
|
|
|
|
@mcp.tool(task=True)
|
|
async def multiply(a: int, b: int) -> int:
|
|
await asyncio.sleep(0.01)
|
|
return a * b
|
|
|
|
@mcp.tool(task=True)
|
|
async def boom() -> str:
|
|
raise ValueError("kaboom")
|
|
|
|
@mcp.tool(task=True)
|
|
async def slow() -> str:
|
|
await asyncio.sleep(5)
|
|
return "done"
|
|
|
|
return mcp
|
|
|
|
|
|
async def test_call_tool_timeout_bounds_total_task_drive(task_server: FastMCP):
|
|
"""A per-call timeout bounds the whole tasked drive, not just one poll.
|
|
|
|
The tool runs far longer than the timeout while each individual poll answers
|
|
instantly; the transparent path must still abort once total execution passes
|
|
the deadline, matching the synchronous `tools/call` timeout contract.
|
|
"""
|
|
async with Client(task_server, mode="auto") as client:
|
|
with pytest.raises((TimeoutError, MCPError)):
|
|
await client.call_tool("slow", {}, timeout=0.3)
|
|
|
|
|
|
async def test_call_tool_transparently_completes_a_task(task_server: FastMCP):
|
|
"""call_tool returns the tool's real result; the caller never sees a task."""
|
|
async with Client(task_server, mode="auto") as client:
|
|
result = await client.call_tool("multiply", {"a": 6, "b": 7})
|
|
|
|
assert result.data == 42
|
|
|
|
|
|
async def test_call_tool_mcp_returns_completed_result(task_server: FastMCP):
|
|
"""call_tool_mcp resolves the tasked call into an ordinary CallToolResult."""
|
|
async with Client(task_server, mode="auto") as client:
|
|
result = await client.call_tool_mcp("multiply", {"a": 3, "b": 4})
|
|
|
|
assert result.structured_content == {"result": 12}
|
|
assert not result.is_error
|
|
|
|
|
|
async def test_failed_task_raises_tool_error(task_server: FastMCP):
|
|
"""A task whose tool raises surfaces as a ToolError through call_tool."""
|
|
async with Client(task_server, mode="auto") as client:
|
|
with pytest.raises(ToolError, match="kaboom"):
|
|
await client.call_tool("boom", {})
|
|
|
|
|
|
async def test_call_tool_task_forwards_requested_version():
|
|
"""`call_tool_task(..., version=...)` tasks the requested version, not the highest."""
|
|
mcp = FastMCP("versioned-task-client")
|
|
mcp.add_extension(TasksExtension())
|
|
|
|
@mcp.tool(name="pick", version="1.0", task=True)
|
|
async def pick_v1() -> str:
|
|
return "v1"
|
|
|
|
@mcp.tool(name="pick", version="2.0", task=True)
|
|
async def pick_v2() -> str:
|
|
return "v2"
|
|
|
|
async with Client(mcp, mode="auto") as client:
|
|
task = await call_tool_task(client, "pick", version="1.0")
|
|
result = await task.result()
|
|
|
|
assert result.data == "v1"
|
|
|
|
|
|
async def test_raw_create_task_result_is_exposed(task_server: FastMCP):
|
|
"""The raw claimed CreateTaskResult is reachable via the session/handle path."""
|
|
async with Client(task_server, mode="auto") as client:
|
|
task = await call_tool_task(client, "multiply", {"a": 2, "b": 5})
|
|
# The raw claimed shape is exposed on the handle.
|
|
assert task.create_result.result_type == "task"
|
|
assert task.create_result.status == "working"
|
|
assert isinstance(task.task_id, str) and task.task_id
|
|
|
|
result = await task.result()
|
|
assert result.data == 10
|
|
|
|
|
|
async def test_legacy_client_never_tasks(task_server: FastMCP):
|
|
"""A legacy-era client never negotiates the capability, so nothing is tasked.
|
|
|
|
The optional-mode tool simply runs synchronously and returns its result
|
|
directly (no CreateTaskResult on the wire).
|
|
"""
|
|
async with Client(task_server, mode="legacy") as client:
|
|
result = await client.call_tool("multiply", {"a": 8, "b": 9})
|
|
|
|
assert result.data == 72
|
|
|
|
|
|
# --- In-task input over the wire -------------------------------------------
|
|
|
|
|
|
@dataclass
|
|
class DinnerPrefs:
|
|
cuisine: str
|
|
vegetarian: bool
|
|
|
|
|
|
def _elicit_request(message: str) -> mcp_types.ElicitRequest:
|
|
return mcp_types.ElicitRequest(
|
|
params=mcp_types.ElicitRequestFormParams(
|
|
message=message,
|
|
requested_schema={
|
|
"type": "object",
|
|
"properties": {
|
|
"cuisine": {"type": "string"},
|
|
"vegetarian": {"type": "boolean"},
|
|
},
|
|
"required": ["cuisine", "vegetarian"],
|
|
},
|
|
)
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def guard_server() -> FastMCP:
|
|
mcp = FastMCP("guard-tasks")
|
|
mcp.add_extension(TasksExtension())
|
|
|
|
@mcp.tool(task=True)
|
|
async def plan_dinner(
|
|
ctx: Context,
|
|
) -> str | mcp_types.InputRequiredResult:
|
|
responses = ctx.input_responses
|
|
if responses is None:
|
|
return mcp_types.InputRequiredResult(
|
|
result_type="input_required",
|
|
input_requests={"prefs": _elicit_request("What's for dinner?")},
|
|
)
|
|
answer = responses["prefs"]
|
|
assert isinstance(answer, mcp_types.ElicitResult)
|
|
assert answer.content is not None
|
|
veg = "vegetarian " if answer.content["vegetarian"] else ""
|
|
return f"Tonight: a {veg}{answer.content['cuisine']} dinner!"
|
|
|
|
return mcp
|
|
|
|
|
|
async def test_in_task_input_answered_transparently(guard_server: FastMCP):
|
|
"""A guard task that asks for input is answered via the elicitation handler."""
|
|
|
|
async def handle_elicitation(message, response_type, params, context):
|
|
return DinnerPrefs(cuisine="Thai", vegetarian=True)
|
|
|
|
client = Client(guard_server, mode="auto", elicitation_handler=handle_elicitation)
|
|
async with client:
|
|
result = await client.call_tool("plan_dinner", {})
|
|
|
|
assert result.data == "Tonight: a vegetarian Thai dinner!"
|
|
|
|
|
|
async def test_in_task_input_without_handler_errors(guard_server: FastMCP):
|
|
"""A guard task with no elicitation handler surfaces a clear error."""
|
|
async with Client(guard_server, mode="auto") as client:
|
|
with pytest.raises(ToolError, match="no elicitation handler"):
|
|
await client.call_tool("plan_dinner", {})
|
|
|
|
|
|
async def test_call_tool_timeout_bounds_a_stalled_elicitation(guard_server: FastMCP):
|
|
"""A stalled elicitation handler cannot outlast the call's timeout.
|
|
|
|
The deadline covers the whole drive, elicitation callbacks included: a
|
|
handler that hangs must abort the tasked call once `timeout=N` elapses,
|
|
matching the synchronous path rather than blocking forever inside the
|
|
callback.
|
|
"""
|
|
|
|
async def slow_elicitation(message, response_type, params, context):
|
|
await asyncio.sleep(5)
|
|
return DinnerPrefs(cuisine="Thai", vegetarian=True)
|
|
|
|
client = Client(guard_server, mode="auto", elicitation_handler=slow_elicitation)
|
|
async with client:
|
|
with pytest.raises((TimeoutError, ToolError, MCPError)):
|
|
await client.call_tool("plan_dinner", {}, timeout=0.3)
|
|
|
|
|
|
async def test_in_task_input_answered_by_handler_set_after_construction(
|
|
guard_server: FastMCP,
|
|
):
|
|
"""An elicitation handler set via set_elicitation_callback reaches in-task input.
|
|
|
|
The tasks client extension is built at construction; set_elicitation_callback
|
|
must rebuild it so a later-configured handler still answers a task's input.
|
|
"""
|
|
|
|
async def handle_elicitation(message, response_type, params, context):
|
|
return DinnerPrefs(cuisine="Thai", vegetarian=True)
|
|
|
|
client = Client(guard_server, mode="auto")
|
|
client.set_elicitation_callback(handle_elicitation)
|
|
async with client:
|
|
result = await client.call_tool("plan_dinner", {})
|
|
|
|
assert result.data == "Tonight: a vegetarian Thai dinner!"
|