mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-25 15:04:18 +02:00
d41ff5bcd8A task tool that returns InputRequiredResult now ends its leg (freeing the worker) and stores the ask as durable state; tasks/update enqueues a fresh Docket execution (the next leg) with accumulated request_state/input_responses injected via ctx. No worker ever blocks on input, so a parked task no longer holds up shutdown. Imperative ctx.elicit() inside a task is removed and raises with guard-pattern guidance. Co-Authored-By: Claude <noreply@anthropic.com>
371 lines
13 KiB
Python
371 lines
13 KiB
Python
"""End-to-end tests for the SEP-2663 `TasksExtension` server adapter.
|
|
|
|
Covers the decide-and-task interceptor (forbidden/optional/required modes and the
|
|
-32003 missing-capability error), the tasks/get|update|cancel handlers, status
|
|
mapping, inlined completed results, argument-coercion parity, TTL, and capability
|
|
advertisement. Server-side tasks are driven in-process via `task_helpers` because
|
|
there is no client task-submission API until Phase 4.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from contextlib import AsyncExitStack
|
|
from types import SimpleNamespace
|
|
from typing import cast
|
|
|
|
import pytest
|
|
from fastmcp_tasks.models import (
|
|
MISSING_REQUIRED_CLIENT_CAPABILITY,
|
|
CreateTaskResult,
|
|
)
|
|
from mcp.server.context import ServerRequestContext
|
|
from mcp.server.session import ServerSession
|
|
from mcp.shared.exceptions import MCPError
|
|
|
|
from fastmcp import FastMCP
|
|
from fastmcp.client import Client
|
|
from fastmcp.exceptions import ToolError
|
|
from fastmcp.server.dependencies import bind_request_context
|
|
from fastmcp.tools.base import ToolResult
|
|
from fastmcp.utilities.tasks import TASKS_EXTENSION_ID, TaskConfig
|
|
from fastmcp_tasks import TasksExtension
|
|
from tests.tasks.task_helpers import (
|
|
_opted_in_request,
|
|
auth_scope,
|
|
call_tool_without_optin,
|
|
get_task,
|
|
make_access_token,
|
|
opt_in_meta,
|
|
run_task,
|
|
running_task_server,
|
|
submit_task,
|
|
wait_for_task,
|
|
)
|
|
|
|
|
|
def _tasks_server() -> FastMCP:
|
|
mcp = FastMCP("tasks")
|
|
mcp.add_extension(TasksExtension())
|
|
|
|
@mcp.tool(task=True)
|
|
async def square(n: int) -> int:
|
|
return n * n
|
|
|
|
@mcp.tool(task=TaskConfig(mode="required"))
|
|
async def must_task(n: int) -> int:
|
|
return n + 1
|
|
|
|
@mcp.tool
|
|
async def plain(n: int) -> int:
|
|
return n - 1
|
|
|
|
@mcp.tool(task=True)
|
|
async def boom() -> int:
|
|
raise ToolError("kaboom")
|
|
|
|
return mcp
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Capability advertisement
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
async def test_capability_advertised_to_modern_client():
|
|
mcp = FastMCP("t")
|
|
mcp.add_extension(TasksExtension())
|
|
|
|
@mcp.tool(task=True)
|
|
async def t(n: int) -> int:
|
|
return n
|
|
|
|
async with Client(mcp, mode="auto") as client:
|
|
extensions = client.server_capabilities.extensions or {}
|
|
assert extensions.get(TASKS_EXTENSION_ID) == {}
|
|
|
|
|
|
async def test_capability_absent_without_extension():
|
|
mcp = FastMCP("t")
|
|
|
|
@mcp.tool
|
|
async def t(n: int) -> int:
|
|
return n
|
|
|
|
async with Client(mcp, mode="auto") as client:
|
|
extensions = client.server_capabilities.extensions or {}
|
|
assert TASKS_EXTENSION_ID not in extensions
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Decide-and-task interceptor
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
async def test_optional_tool_tasks_when_opted_in():
|
|
mcp = _tasks_server()
|
|
async with running_task_server(mcp):
|
|
created = await submit_task(mcp, "square", {"n": 5})
|
|
assert isinstance(created, CreateTaskResult)
|
|
assert created.status == "working"
|
|
final = await wait_for_task(mcp, created.task_id)
|
|
assert final.status == "completed"
|
|
assert final.result is not None
|
|
assert final.result["structuredContent"]["result"] == 25
|
|
|
|
|
|
async def test_optional_tool_runs_sync_without_opt_in():
|
|
mcp = _tasks_server()
|
|
async with running_task_server(mcp):
|
|
result = await call_tool_without_optin(mcp, "square", {"n": 5})
|
|
assert not isinstance(result, CreateTaskResult)
|
|
assert result.structured_content == {"result": 25}
|
|
|
|
|
|
async def test_forbidden_tool_never_tasks_even_with_opt_in():
|
|
mcp = _tasks_server()
|
|
async with running_task_server(mcp):
|
|
# `plain` is mode=forbidden; opting in must not task it.
|
|
result = await submit_task_expecting_sync(mcp, "plain", {"n": 5})
|
|
assert result.structured_content == {"result": 4}
|
|
|
|
|
|
async def submit_task_expecting_sync(mcp, name, args):
|
|
with auth_scope(None), _opted_in_request(name, args, None):
|
|
return await mcp.call_tool(name, args)
|
|
|
|
|
|
async def test_required_tool_tasks_when_opted_in():
|
|
mcp = _tasks_server()
|
|
async with running_task_server(mcp):
|
|
created = await submit_task(mcp, "must_task", {"n": 10})
|
|
final = await wait_for_task(mcp, created.task_id)
|
|
assert final.status == "completed"
|
|
assert final.result is not None
|
|
assert final.result["structuredContent"]["result"] == 11
|
|
|
|
|
|
async def test_required_tool_without_opt_in_raises_missing_capability():
|
|
mcp = _tasks_server()
|
|
async with running_task_server(mcp):
|
|
with pytest.raises(MCPError) as exc_info:
|
|
await call_tool_without_optin(mcp, "must_task", {"n": 1})
|
|
error = exc_info.value.error
|
|
assert error.code == MISSING_REQUIRED_CLIENT_CAPABILITY
|
|
assert error.data == {
|
|
"requiredCapabilities": {"extensions": {TASKS_EXTENSION_ID: {}}}
|
|
}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Task id and status
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
async def test_task_ids_are_server_generated_and_distinct():
|
|
mcp = _tasks_server()
|
|
async with running_task_server(mcp):
|
|
a = await submit_task(mcp, "square", {"n": 1})
|
|
b = await submit_task(mcp, "square", {"n": 2})
|
|
assert a.task_id != b.task_id
|
|
assert len(a.task_id) >= 20
|
|
|
|
|
|
async def test_get_unknown_task_raises_not_found():
|
|
mcp = _tasks_server()
|
|
async with running_task_server(mcp):
|
|
with pytest.raises(MCPError, match="not found"):
|
|
await get_task(mcp, "does-not-exist")
|
|
|
|
|
|
async def test_failed_task_surfaces_error_not_completed():
|
|
mcp = _tasks_server()
|
|
async with running_task_server(mcp):
|
|
created = await submit_task(mcp, "boom", {})
|
|
final = await wait_for_task(mcp, created.task_id)
|
|
assert final.status == "failed"
|
|
assert final.error is not None
|
|
assert "kaboom" in final.error["message"]
|
|
assert final.result is None
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Argument coercion parity
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
async def test_task_arguments_are_coerced_like_sync_path():
|
|
mcp = _tasks_server()
|
|
async with running_task_server(mcp):
|
|
# "6" coerces to int 6 exactly as the synchronous path would.
|
|
final = await run_task(mcp, "square", {"n": "6"})
|
|
assert final.result is not None
|
|
assert final.result["structuredContent"]["result"] == 36
|
|
|
|
|
|
async def test_invalid_task_arguments_reject_at_submission():
|
|
mcp = _tasks_server()
|
|
async with running_task_server(mcp):
|
|
with pytest.raises(Exception):
|
|
await submit_task(mcp, "square", {"n": "not-a-number"})
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# TTL / poll interval
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
async def test_create_and_get_carry_ttl_and_poll_interval():
|
|
mcp = _tasks_server()
|
|
async with running_task_server(mcp):
|
|
created = await submit_task(mcp, "square", {"n": 3})
|
|
assert created.ttl_ms is not None and created.ttl_ms > 0
|
|
assert created.poll_interval_ms == 5000
|
|
got = await get_task(mcp, created.task_id)
|
|
assert got.ttl_ms == created.ttl_ms
|
|
assert got.poll_interval_ms == 5000
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Cancellation
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
async def test_cancel_transitions_task_to_cancelled():
|
|
mcp = FastMCP("t")
|
|
mcp.add_extension(TasksExtension())
|
|
release = asyncio.Event()
|
|
|
|
@mcp.tool(task=True)
|
|
async def slow() -> str:
|
|
await release.wait()
|
|
return "done"
|
|
|
|
async with running_task_server(mcp):
|
|
created = await submit_task(mcp, "slow", {})
|
|
ack = await cancel_and_release(mcp, created.task_id, release)
|
|
assert ack is not None
|
|
final = await wait_for_task(
|
|
mcp, created.task_id, target_states=frozenset({"cancelled", "completed"})
|
|
)
|
|
assert final.status in {"cancelled", "completed"}
|
|
|
|
|
|
async def cancel_and_release(mcp, task_id, release):
|
|
from tests.tasks.task_helpers import cancel_task
|
|
|
|
ack = await cancel_task(mcp, task_id)
|
|
release.set()
|
|
return ack
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Serve-time guard
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
async def test_task_tool_without_extension_fails_at_serve_time():
|
|
mcp = FastMCP("t")
|
|
|
|
@mcp.tool(task=True)
|
|
async def t(n: int) -> int:
|
|
return n
|
|
|
|
with pytest.raises(RuntimeError, match="tasks extension"):
|
|
async with mcp._lifespan_manager():
|
|
pass
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Auth-scoped isolation
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
async def test_tasks_isolated_across_auth_scopes():
|
|
mcp = _tasks_server()
|
|
alice = make_access_token("alice")
|
|
bob = make_access_token("bob")
|
|
async with running_task_server(mcp):
|
|
created = await submit_task(mcp, "square", {"n": 4}, access_token=alice)
|
|
# Alice sees her task.
|
|
mine = await get_task(mcp, created.task_id, access_token=alice)
|
|
assert mine.task_id == created.task_id
|
|
# Bob cannot: a cross-scope id is indistinguishable from missing.
|
|
with pytest.raises(MCPError, match="not found"):
|
|
await get_task(mcp, created.task_id, access_token=bob)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Protocol-era gating of the tasking decision
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
async def test_legacy_era_opt_in_is_ignored():
|
|
"""A handshake-era request cannot be tasked, even with the _meta opt-in.
|
|
|
|
The SDK strips `capabilities.extensions` from pre-2026 handshakes, so a
|
|
legacy client can never have negotiated the tasks extension — a stray
|
|
per-request opt-in on a legacy connection is treated as absent and an
|
|
`optional` tool runs synchronously.
|
|
"""
|
|
mcp = _tasks_server()
|
|
async with running_task_server(mcp):
|
|
srctx = ServerRequestContext(
|
|
session=cast(ServerSession, SimpleNamespace()),
|
|
lifespan_context={},
|
|
protocol_version="2025-06-18",
|
|
method="tools/call",
|
|
params={"name": "square", "arguments": {"n": 3}, "_meta": opt_in_meta()},
|
|
)
|
|
with bind_request_context(srctx):
|
|
result = await mcp.call_tool("square", {"n": 3})
|
|
assert isinstance(result, ToolResult)
|
|
|
|
|
|
async def test_legacy_era_required_tool_raises_missing_capability():
|
|
"""`required` tools refuse legacy-era calls with -32003 even when opted in."""
|
|
mcp = _tasks_server()
|
|
async with running_task_server(mcp):
|
|
srctx = ServerRequestContext(
|
|
session=cast(ServerSession, SimpleNamespace()),
|
|
lifespan_context={},
|
|
protocol_version="2025-06-18",
|
|
method="tools/call",
|
|
params={
|
|
"name": "must_task",
|
|
"arguments": {"n": 3},
|
|
"_meta": opt_in_meta(),
|
|
},
|
|
)
|
|
with bind_request_context(srctx):
|
|
with pytest.raises(MCPError) as exc_info:
|
|
await mcp.call_tool("must_task", {"n": 3})
|
|
assert exc_info.value.error.code == -32003
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Worker-hook lifecycle across multiple servers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
async def test_worker_hooks_survive_sibling_server_shutdown():
|
|
"""One server's shutdown must not strand another server's workers.
|
|
|
|
The worker-side hooks core exposes are process-global; two sibling servers
|
|
each running a TasksExtension refcount them, so the hooks clear only when
|
|
the last extension lifespan exits.
|
|
"""
|
|
from fastmcp.server import dependencies as core_dependencies
|
|
|
|
server_a = _tasks_server()
|
|
server_b = _tasks_server()
|
|
|
|
async with AsyncExitStack() as stack_b:
|
|
await stack_b.enter_async_context(server_b._lifespan_manager())
|
|
async with AsyncExitStack() as stack_a:
|
|
await stack_a.enter_async_context(server_a._lifespan_manager())
|
|
assert core_dependencies._background_context_factory is not None
|
|
# Server A has shut down; server B's workers still need the hooks.
|
|
assert core_dependencies._background_context_factory is not None
|
|
# The last extension exited; hooks are cleared.
|
|
assert core_dependencies._background_context_factory is None
|