mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-09 15:19:10 +02:00
53741dc9c7A resumed leg that runs longer than its pointer's wall-clock TTL stranded _lookup_task on the base leg (false completion / not found). Each poll now refreshes the routing keys' TTL (sliding expiration), so an actively-polled task keeps them alive regardless of execution duration, and the resumed-leg write uses the same buffered TTL as creation. Separately, remote-worker server resolution now respects the requested tool version, so two versions of the same mounted tool name resolve to their own child server.
562 lines
21 KiB
Python
562 lines
21 KiB
Python
"""SEP-2663 task execution through mounted servers (tools-only).
|
|
|
|
Verifies that background tasks work when a tool lives on a mounted child server:
|
|
the parent (which registers the tasks extension and owns the Docket) runs the
|
|
tool as a task, the worker resolves back to the child server, dependencies
|
|
resolve, and mode enforcement / metadata survive mounting. SEP-2663 is
|
|
tools-only, so the SEP-1686 prompt/resource mount cases are gone.
|
|
|
|
Two architectural notes vs. SEP-1686:
|
|
- The `tools/call` interceptor composes at the *registering* (parent/root)
|
|
server's dispatch and short-circuits before delegating into a mounted child,
|
|
so for a tasked call only the root's middleware wraps submission (the tool
|
|
body runs later in the worker). Child/grandchild middleware do not wrap a
|
|
tasked submission.
|
|
- Worker server resolution is single-level: a tool reached through nested mounts
|
|
resolves to the outermost mounted child (the mount point the call arrived
|
|
through), which still reaches deeper components via its own mounts.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from typing import cast
|
|
|
|
import mcp_types as mt
|
|
import pytest
|
|
from docket import Docket
|
|
from fastmcp_tasks.dependencies import CurrentDocket
|
|
from mcp_types import Tool as MCPTool
|
|
from mcp_types import ToolExecution
|
|
|
|
from fastmcp import Context, FastMCP
|
|
from fastmcp.server.dependencies import CurrentFastMCP
|
|
from fastmcp.server.middleware import CallNext, Middleware, MiddlewareContext
|
|
from fastmcp.server.providers.proxy import ClientFactoryT, ProxyTool
|
|
from fastmcp.tools.base import ToolResult
|
|
from fastmcp.utilities.tasks import TaskConfig
|
|
from fastmcp_tasks import TasksExtension
|
|
from tests.tasks.task_helpers import (
|
|
call_tool_without_optin,
|
|
running_task_server,
|
|
submit_task,
|
|
wait_for_task,
|
|
)
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def reset_docket_memory_server():
|
|
"""Reset the shared memory:// Docket server between tests for isolation."""
|
|
if hasattr(Docket, "_memory_server"):
|
|
delattr(Docket, "_memory_server")
|
|
yield
|
|
if hasattr(Docket, "_memory_server"):
|
|
delattr(Docket, "_memory_server")
|
|
|
|
|
|
@pytest.fixture
|
|
def child_server() -> FastMCP:
|
|
mcp = FastMCP("child-server")
|
|
|
|
@mcp.tool(task=True)
|
|
async def multiply(a: int, b: int) -> int:
|
|
return a * b
|
|
|
|
@mcp.tool(task=False)
|
|
async def sync_child_tool(message: str) -> str:
|
|
return f"child sync: {message}"
|
|
|
|
return mcp
|
|
|
|
|
|
@pytest.fixture
|
|
def parent_server(child_server: FastMCP) -> FastMCP:
|
|
parent = FastMCP("parent-server")
|
|
parent.add_extension(TasksExtension())
|
|
|
|
@parent.tool(task=True)
|
|
async def parent_tool(value: int) -> int:
|
|
return value * 10
|
|
|
|
parent.mount(child_server, namespace="child")
|
|
return parent
|
|
|
|
|
|
class TestMountedToolTasks:
|
|
async def test_mounted_tool_task_returns_correct_result(self, parent_server):
|
|
async with running_task_server(parent_server):
|
|
created = await submit_task(
|
|
parent_server, "child_multiply", {"a": 8, "b": 9}
|
|
)
|
|
assert created.status == "working"
|
|
final = await wait_for_task(parent_server, created.task_id)
|
|
assert final.status == "completed"
|
|
assert final.result is not None
|
|
assert final.result["structuredContent"]["result"] == 72
|
|
|
|
async def test_mounted_and_parent_tasks_both_work(self, parent_server):
|
|
async with running_task_server(parent_server):
|
|
parent_created = await submit_task(
|
|
parent_server, "parent_tool", {"value": 5}
|
|
)
|
|
child_created = await submit_task(
|
|
parent_server, "child_multiply", {"a": 2, "b": 3}
|
|
)
|
|
parent_final = await wait_for_task(parent_server, parent_created.task_id)
|
|
child_final = await wait_for_task(parent_server, child_created.task_id)
|
|
assert parent_final.result is not None
|
|
assert parent_final.result["structuredContent"]["result"] == 50
|
|
assert child_final.result is not None
|
|
assert child_final.result["structuredContent"]["result"] == 6
|
|
|
|
async def test_sync_only_mounted_tool_runs_synchronously(self, parent_server):
|
|
"""A task=False mounted tool runs sync even when the client opts in."""
|
|
async with running_task_server(parent_server):
|
|
# Opting in on a forbidden tool must not task it.
|
|
from tests.tasks.task_helpers import _opted_in_request
|
|
|
|
with _opted_in_request("child_sync_child_tool", {"message": "hi"}, None):
|
|
result = await parent_server.call_tool(
|
|
"child_sync_child_tool", {"message": "hi"}
|
|
)
|
|
assert not hasattr(result, "task_id")
|
|
assert "child sync: hi" in result.content[0].text
|
|
|
|
|
|
class TestRemoteWorkerServerResolution:
|
|
"""A separate worker process re-resolves the owning child from the root.
|
|
|
|
The in-process submission map is unreachable across processes, so the worker
|
|
recovers the mounted child server from the snapshotted tool name instead of
|
|
falling back to the root (which would break child-specific state/config).
|
|
"""
|
|
|
|
async def test_resolve_owning_server_recovers_mounted_child(self, parent_server):
|
|
import weakref
|
|
|
|
from fastmcp_tasks.context import (
|
|
TaskContextSnapshot,
|
|
_resolve_owning_server,
|
|
)
|
|
|
|
from fastmcp.server.dependencies import _current_server
|
|
|
|
child = await parent_server.get_tool("child_multiply")
|
|
|
|
token = _current_server.set(weakref.ref(parent_server))
|
|
try:
|
|
snapshot = TaskContextSnapshot(owning_tool_name="child_multiply")
|
|
resolved = await _resolve_owning_server(snapshot)
|
|
assert resolved is child._server
|
|
|
|
# A parent-owned (unmounted) tool resolves to None so the caller
|
|
# falls back to the root, and a missing name is likewise None.
|
|
assert (
|
|
await _resolve_owning_server(
|
|
TaskContextSnapshot(owning_tool_name="parent_tool")
|
|
)
|
|
is None
|
|
)
|
|
assert (
|
|
await _resolve_owning_server(
|
|
TaskContextSnapshot(owning_tool_name="does_not_exist")
|
|
)
|
|
is None
|
|
)
|
|
assert await _resolve_owning_server(TaskContextSnapshot()) is None
|
|
finally:
|
|
_current_server.reset(token)
|
|
|
|
async def test_resolve_owning_server_respects_version(self):
|
|
"""Two versions of a mounted tool name resolve to their own child server."""
|
|
import weakref
|
|
|
|
from fastmcp_tasks.context import (
|
|
TaskContextSnapshot,
|
|
_resolve_owning_server,
|
|
)
|
|
|
|
from fastmcp.server.dependencies import _current_server
|
|
|
|
child_v1 = FastMCP("child-v1")
|
|
|
|
@child_v1.tool(name="calc", version="1.0", task=True)
|
|
async def calc_v1() -> str:
|
|
return "v1"
|
|
|
|
child_v2 = FastMCP("child-v2")
|
|
|
|
@child_v2.tool(name="calc", version="2.0", task=True)
|
|
async def calc_v2() -> str:
|
|
return "v2"
|
|
|
|
parent = FastMCP("parent-versions")
|
|
parent.add_extension(TasksExtension())
|
|
parent.mount(child_v1)
|
|
parent.mount(child_v2)
|
|
|
|
token = _current_server.set(weakref.ref(parent))
|
|
try:
|
|
resolved_v1 = await _resolve_owning_server(
|
|
TaskContextSnapshot(owning_tool_name="calc", owning_tool_version="1.0")
|
|
)
|
|
resolved_v2 = await _resolve_owning_server(
|
|
TaskContextSnapshot(owning_tool_name="calc", owning_tool_version="2.0")
|
|
)
|
|
assert resolved_v1 is child_v1
|
|
assert resolved_v2 is child_v2
|
|
finally:
|
|
_current_server.reset(token)
|
|
|
|
|
|
class TestMountedToolTasksNoPrefix:
|
|
async def test_mounted_tool_without_prefix_works(self, child_server):
|
|
parent = FastMCP("parent-no-prefix")
|
|
parent.add_extension(TasksExtension())
|
|
parent.mount(child_server) # no prefix
|
|
async with running_task_server(parent):
|
|
final = await wait_for_task(
|
|
parent,
|
|
(await submit_task(parent, "multiply", {"a": 5, "b": 6})).task_id,
|
|
)
|
|
assert final.result is not None
|
|
assert final.result["structuredContent"]["result"] == 30
|
|
|
|
|
|
class TestMountedTaskDependencies:
|
|
async def test_mounted_task_receives_docket_dependency(self):
|
|
child = FastMCP("dep-child")
|
|
|
|
@child.tool(task=True)
|
|
async def tool_with_docket(docket: Docket = CurrentDocket()) -> str:
|
|
return f"docket available: {docket is not None}"
|
|
|
|
parent = FastMCP("dep-parent")
|
|
parent.add_extension(TasksExtension())
|
|
parent.mount(child, namespace="child")
|
|
|
|
async with running_task_server(parent):
|
|
final = await wait_for_task(
|
|
parent,
|
|
(await submit_task(parent, "child_tool_with_docket", {})).task_id,
|
|
)
|
|
assert final.result is not None
|
|
assert "docket available: True" in final.result["content"][0]["text"]
|
|
|
|
|
|
class TestMountedTaskServerContext:
|
|
async def test_current_fastmcp_resolves_to_child_server(self):
|
|
child = FastMCP("child")
|
|
|
|
@child.tool(task=True)
|
|
async def whoami(server: FastMCP = CurrentFastMCP()) -> str:
|
|
return f"server name: {server.name}"
|
|
|
|
parent = FastMCP("parent")
|
|
parent.add_extension(TasksExtension())
|
|
parent.mount(child, namespace="child")
|
|
|
|
async with running_task_server(parent):
|
|
final = await wait_for_task(
|
|
parent, (await submit_task(parent, "child_whoami", {})).task_id
|
|
)
|
|
assert final.result is not None
|
|
assert "server name: child" in final.result["content"][0]["text"]
|
|
|
|
async def test_context_fastmcp_resolves_to_child_server(self):
|
|
child = FastMCP("child")
|
|
|
|
@child.tool(task=True)
|
|
async def whoami_ctx(ctx: Context) -> str:
|
|
return f"context server: {ctx.fastmcp.name}"
|
|
|
|
parent = FastMCP("parent")
|
|
parent.add_extension(TasksExtension())
|
|
parent.mount(child, namespace="child")
|
|
|
|
async with running_task_server(parent):
|
|
final = await wait_for_task(
|
|
parent, (await submit_task(parent, "child_whoami_ctx", {})).task_id
|
|
)
|
|
assert final.result is not None
|
|
assert "context server: child" in final.result["content"][0]["text"]
|
|
|
|
|
|
class TestMultipleMounts:
|
|
async def test_tasks_work_with_multiple_mounts(self):
|
|
child1 = FastMCP("child1")
|
|
child2 = FastMCP("child2")
|
|
|
|
@child1.tool(task=True)
|
|
async def add(a: int, b: int) -> int:
|
|
return a + b
|
|
|
|
@child2.tool(task=True)
|
|
async def subtract(a: int, b: int) -> int:
|
|
return a - b
|
|
|
|
parent = FastMCP("multi-parent")
|
|
parent.add_extension(TasksExtension())
|
|
parent.mount(child1, namespace="math1")
|
|
parent.mount(child2, namespace="math2")
|
|
|
|
async with running_task_server(parent):
|
|
r1 = await wait_for_task(
|
|
parent,
|
|
(await submit_task(parent, "math1_add", {"a": 10, "b": 5})).task_id,
|
|
)
|
|
r2 = await wait_for_task(
|
|
parent,
|
|
(
|
|
await submit_task(parent, "math2_subtract", {"a": 10, "b": 5})
|
|
).task_id,
|
|
)
|
|
assert r1.result is not None
|
|
assert r1.result["structuredContent"]["result"] == 15
|
|
assert r2.result is not None
|
|
assert r2.result["structuredContent"]["result"] == 5
|
|
|
|
async def test_same_function_names_do_not_collide(self):
|
|
child1 = FastMCP("child1")
|
|
child2 = FastMCP("child2")
|
|
|
|
@child1.tool(task=True)
|
|
async def process(value: int) -> int:
|
|
return value * 2
|
|
|
|
@child2.tool(task=True)
|
|
async def process(value: int) -> int: # noqa: F811
|
|
return value * 3
|
|
|
|
parent = FastMCP("parent")
|
|
parent.add_extension(TasksExtension())
|
|
parent.mount(child1, namespace="c1")
|
|
parent.mount(child2, namespace="c2")
|
|
|
|
async with running_task_server(parent):
|
|
r1 = await wait_for_task(
|
|
parent,
|
|
(await submit_task(parent, "c1_process", {"value": 10})).task_id,
|
|
)
|
|
r2 = await wait_for_task(
|
|
parent,
|
|
(await submit_task(parent, "c2_process", {"value": 10})).task_id,
|
|
)
|
|
assert r1.result is not None
|
|
assert r1.result["structuredContent"]["result"] == 20
|
|
assert r2.result is not None
|
|
assert r2.result["structuredContent"]["result"] == 30
|
|
|
|
async def test_nested_mount_prefix_accumulation(self):
|
|
grandchild = FastMCP("gc")
|
|
child = FastMCP("child")
|
|
parent = FastMCP("parent")
|
|
parent.add_extension(TasksExtension())
|
|
|
|
@grandchild.tool(task=True)
|
|
async def deep_tool() -> str:
|
|
return "deep"
|
|
|
|
child.mount(grandchild, namespace="gc")
|
|
parent.mount(child, namespace="child")
|
|
|
|
async with running_task_server(parent):
|
|
final = await wait_for_task(
|
|
parent,
|
|
(await submit_task(parent, "child_gc_deep_tool", {})).task_id,
|
|
)
|
|
assert final.result is not None
|
|
assert final.result["structuredContent"]["result"] == "deep"
|
|
|
|
|
|
class TestMountedTaskMetadata:
|
|
async def test_mounted_tool_list_preserves_task_support_metadata(self):
|
|
child = FastMCP("child")
|
|
|
|
@child.tool(task=True)
|
|
async def foo() -> dict[str, bool]:
|
|
return {"ok": True}
|
|
|
|
parent = FastMCP("parent")
|
|
parent.mount(child)
|
|
|
|
child_tool = next(t for t in await child.list_tools() if t.name == "foo")
|
|
parent_tool = next(t for t in await parent.list_tools() if t.name == "foo")
|
|
|
|
child_mcp = child_tool.to_mcp_tool(name=child_tool.name)
|
|
parent_mcp = parent_tool.to_mcp_tool(name=parent_tool.name)
|
|
assert child_mcp.execution is not None
|
|
assert parent_mcp.execution is not None
|
|
assert child_mcp.execution.task_support == "optional"
|
|
assert parent_mcp.execution.task_support == "optional"
|
|
|
|
async def test_proxy_tool_preserves_execution_metadata(self):
|
|
mcp_tool = MCPTool(
|
|
name="remote_task_tool",
|
|
description="A remote tool that supports tasks",
|
|
input_schema={"type": "object", "properties": {}},
|
|
execution=ToolExecution(task_support="optional"),
|
|
)
|
|
proxy = ProxyTool.from_mcp_tool(cast(ClientFactoryT, lambda: None), mcp_tool)
|
|
result = proxy.to_mcp_tool(name=proxy.name)
|
|
assert result.execution is not None
|
|
assert result.execution.task_support == "optional"
|
|
|
|
|
|
class TestMountedTaskConfigModes:
|
|
@pytest.fixture
|
|
def parent_with_modes(self) -> FastMCP:
|
|
child = FastMCP("child-modes")
|
|
|
|
@child.tool(task=TaskConfig(mode="optional"))
|
|
async def optional_tool() -> str:
|
|
return "optional result"
|
|
|
|
@child.tool(task=TaskConfig(mode="required"))
|
|
async def required_tool() -> str:
|
|
return "required result"
|
|
|
|
@child.tool(task=TaskConfig(mode="forbidden"))
|
|
async def forbidden_tool() -> str:
|
|
return "forbidden result"
|
|
|
|
parent = FastMCP("parent-modes")
|
|
parent.add_extension(TasksExtension())
|
|
parent.mount(child, namespace="child")
|
|
return parent
|
|
|
|
async def test_optional_mode_sync_through_mount(self, parent_with_modes):
|
|
async with running_task_server(parent_with_modes):
|
|
result = await call_tool_without_optin(
|
|
parent_with_modes, "child_optional_tool", {}
|
|
)
|
|
assert "optional result" in result.content[0].text
|
|
|
|
async def test_optional_mode_task_through_mount(self, parent_with_modes):
|
|
async with running_task_server(parent_with_modes):
|
|
final = await wait_for_task(
|
|
parent_with_modes,
|
|
(
|
|
await submit_task(parent_with_modes, "child_optional_tool", {})
|
|
).task_id,
|
|
)
|
|
assert final.result is not None
|
|
assert final.result["structuredContent"]["result"] == "optional result"
|
|
|
|
async def test_required_mode_with_task_through_mount(self, parent_with_modes):
|
|
async with running_task_server(parent_with_modes):
|
|
final = await wait_for_task(
|
|
parent_with_modes,
|
|
(
|
|
await submit_task(parent_with_modes, "child_required_tool", {})
|
|
).task_id,
|
|
)
|
|
assert final.result is not None
|
|
assert final.result["structuredContent"]["result"] == "required result"
|
|
|
|
async def test_required_mode_without_task_through_mount(self, parent_with_modes):
|
|
from fastmcp_tasks.models import MISSING_REQUIRED_CLIENT_CAPABILITY
|
|
from mcp.shared.exceptions import MCPError
|
|
|
|
async with running_task_server(parent_with_modes):
|
|
with pytest.raises(MCPError) as exc_info:
|
|
await call_tool_without_optin(
|
|
parent_with_modes, "child_required_tool", {}
|
|
)
|
|
assert exc_info.value.error.code == MISSING_REQUIRED_CLIENT_CAPABILITY
|
|
|
|
async def test_forbidden_mode_sync_through_mount(self, parent_with_modes):
|
|
async with running_task_server(parent_with_modes):
|
|
result = await call_tool_without_optin(
|
|
parent_with_modes, "child_forbidden_tool", {}
|
|
)
|
|
assert "forbidden result" in result.content[0].text
|
|
|
|
|
|
class ToolTracingMiddleware(Middleware):
|
|
def __init__(self, name: str, calls: list[str]):
|
|
super().__init__()
|
|
self._name = name
|
|
self._calls = calls
|
|
|
|
async def on_call_tool(
|
|
self,
|
|
context: MiddlewareContext[mt.CallToolRequestParams],
|
|
call_next: CallNext[mt.CallToolRequestParams, ToolResult],
|
|
) -> ToolResult:
|
|
self._calls.append(f"{self._name}:before")
|
|
result = await call_next(context)
|
|
self._calls.append(f"{self._name}:after")
|
|
return result
|
|
|
|
|
|
class TestMiddlewareWithMountedTasks:
|
|
async def test_root_middleware_wraps_task_submission(self):
|
|
"""For a tasked call, the root's middleware wraps submission.
|
|
|
|
The interceptor composes at the registering (parent) server and
|
|
short-circuits before delegating into the mounted child, so child
|
|
middleware does not wrap a tasked submission; the tool body runs later
|
|
in the worker.
|
|
"""
|
|
calls: list[str] = []
|
|
|
|
grandchild = FastMCP("Grandchild")
|
|
|
|
@grandchild.tool(task=True)
|
|
async def compute(x: int) -> int:
|
|
calls.append("grandchild:tool")
|
|
return x * 2
|
|
|
|
grandchild.add_middleware(ToolTracingMiddleware("grandchild", calls))
|
|
child = FastMCP("Child")
|
|
child.mount(grandchild, namespace="gc")
|
|
child.add_middleware(ToolTracingMiddleware("child", calls))
|
|
parent = FastMCP("Parent")
|
|
parent.add_extension(TasksExtension())
|
|
parent.mount(child, namespace="c")
|
|
parent.add_middleware(ToolTracingMiddleware("parent", calls))
|
|
|
|
async with running_task_server(parent):
|
|
created = await submit_task(parent, "c_gc_compute", {"x": 5})
|
|
final = await wait_for_task(parent, created.task_id)
|
|
assert final.result is not None
|
|
assert final.result["structuredContent"]["result"] == 10
|
|
|
|
assert calls == ["parent:before", "parent:after", "grandchild:tool"]
|
|
|
|
|
|
class TestMountedDocketOwnership:
|
|
async def test_mounted_child_does_not_own_docket(self, parent_server, child_server):
|
|
"""The parent owns the Docket; the mounted child does not."""
|
|
async with running_task_server(parent_server):
|
|
assert parent_server.docket is not None
|
|
assert child_server.docket is None
|
|
|
|
|
|
class TestSlowMountedTaskCancellation:
|
|
async def test_cancel_mounted_task(self):
|
|
child = FastMCP("child")
|
|
release = asyncio.Event()
|
|
|
|
@child.tool(task=True)
|
|
async def slow() -> str:
|
|
await release.wait()
|
|
return "done"
|
|
|
|
parent = FastMCP("parent")
|
|
parent.add_extension(TasksExtension())
|
|
parent.mount(child, namespace="child")
|
|
|
|
from tests.tasks.task_helpers import cancel_task
|
|
|
|
async with running_task_server(parent):
|
|
created = await submit_task(parent, "child_slow", {})
|
|
await cancel_task(parent, created.task_id)
|
|
release.set()
|
|
final = await wait_for_task(
|
|
parent,
|
|
created.task_id,
|
|
target_states=frozenset({"cancelled", "completed"}),
|
|
)
|
|
assert final.status in {"cancelled", "completed"}
|