ceos/fastmcp
1
0
Fork 0
mirror of https://github.com/PrefectHQ/fastmcp.git synced 2026-08-24 22:44:17 +02:00
Code Issues Projects Releases Packages Wiki Activity
fastmcp/tests/server/tasks/test_task_mount.py
Chris Guidry d41bcb2c9e
 Unify background task context forwarding, fix concurrent dependency bugs (#3710) 
* Unify background task context forwarding and fix concurrent dependency bugs

We've been getting a steady trickle of edge-case reports around background tasks
and contextual dependencies over the last few months (#3654, #3656, #3569). Each
one pointed at a different symptom, but they all traced back to the same area:
the way context is negotiated between the "frontend" server and Docket workers
was grown piecemeal, with each new piece of context (access tokens, HTTP headers,
origin request IDs) getting its own Redis key, its own restore function, and its
own ContextVar. This made it hard to reason about what state was available where,
and the shared-instance Dependency pattern made concurrent tasks stomp on each
other's cleanup state.

This takes a step back and reworks the whole thing as a single unified system:

- Dependency subclasses (_CurrentContext, Progress, _CurrentAccessToken, etc.)
  are now stateless factories — __aenter__ returns a fresh per-invocation
  object, so concurrent tasks never share mutable state. Fixes #3654, #3656.

- The three individual context-snapshot Redis keys (access_token, http_headers,
  origin_request_id) are collapsed into a single TaskContextSnapshot stored as
  one JSON key per task. The three _restore_task_* functions and two ContextVars
  they populated are gone.

- Sync functions like get_http_request() and get_access_token() now find the
  snapshot transparently in background tasks via a 3-tier sync fallback:
  ContextVar (set by _CurrentContext for functions with deps) → in-memory dict
  (same-process workers) → sync Redis GET (out-of-process workers). No function
  wrapping needed.

- The _wrap_for_task_http_headers hack is deleted. FunctionTool registers its
  raw function with Docket so Docket sees and resolves ALL dependencies,
  including Docket-native ones like Retry and Timeout.

- ProxyTool.from_mcp_tool() now propagates execution.taskSupport metadata from
  remote tools. Fixes #3569.

- Removed redundant _current_docket/_current_worker ContextVar management from
  Context.__aenter__/__aexit__ (they're only set in the lifespan now).

Closes #3654
Closes #3656
Closes #3569

🤖 Generated with Claude Code

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* Address code review feedback

- _OptionalCurrentContext: guard __aexit__ against cleaning up contexts it
  didn't create (check is_background_task before delegating)
- Narrow except clauses in snapshot loading (OSError, JSONDecodeError, etc.
  instead of bare Exception)
- Fix docstrings on register_with_docket for resources/prompts/templates
- Simplify Progress: read ExecutionProgress directly from current_execution
  instead of creating and manually entering a DocketProgress wrapper

🤖 Generated with Claude Code

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* Use pop-on-access transfer buffer instead of bounded LRU cache for snapshots

The in-memory snapshot dict is a transfer mechanism, not a cache. Entries go
in at submission and come out at the worker's first access. Using pop instead
of get means the dict only holds entries during the brief submission-to-execution
window, bounded by task concurrency (~10) rather than a 10,000-entry LRU limit.

🤖 Generated with Claude Code

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* Drop in-memory transfer buffer, use sync Redis for all backends

Instead of maintaining an in-memory dict to bridge the async/sync gap, use
a sync Redis client directly. For memory:// backends (fakeredis), shares the
same FakeServer instance via docket._redis.get_memory_server() so data written
by the async Docket client is visible to sync reads. For real Redis, creates a
standard sync connection. No in-process state to manage at all.

🤖 Generated with Claude Code

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* Move snapshot operations to TaskContextSnapshot methods

capture(), from_json(), to_json(), save() are now classmethod/instance methods
on the dataclass instead of free functions. Deduplicates JSON parsing that was
copy-pasted between the async and sync load paths.

🤖 Generated with Claude Code

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* Trim implementation details from register_with_docket docstrings

🤖 Generated with Claude Code

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* Clarify docket lookup comment in submit_to_docket

🤖 Generated with Claude Code

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* Restore docket/worker ContextVar bridge in Context.__aenter__

Servers that own the Docket (the parent) re-set _current_docket/_current_worker
from their instance attributes when entering a Context. Mounted children skip
this (their _docket is None), so they inherit the parent's value. This is needed
for ASGI deployments where ContextVars set during the lifespan don't propagate
to request handlers.

🤖 Generated with Claude Code

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* Key snapshot cache by task_id to prevent cross-task context leakage

Docket workers may reuse the same asyncio context for sequential tasks.
The ContextVar cache now stores (task_id, snapshot) tuples so stale entries
from previous tasks are automatically ignored.

🤖 Generated with Claude Code

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-04-03 10:48:27 -04:00

1083 lines
39 KiB
Python
Raw Blame History

"""
Tests for MCP SEP-1686 task protocol support through mounted servers.
Verifies that tasks work seamlessly when calling tools/prompts/resources
on mounted child servers through a parent server.
"""
import asyncio
import mcp.types as mt
import pytest
from docket import Docket
from mcp.types import Tool as MCPTool
from mcp.types import ToolExecution
from fastmcp import FastMCP
from fastmcp.client import Client
from fastmcp.prompts.base import PromptResult
from fastmcp.resources.base import ResourceResult
from fastmcp.server.dependencies import CurrentDocket, CurrentFastMCP
from fastmcp.server.middleware import CallNext, Middleware, MiddlewareContext
from fastmcp.server.providers.proxy import ProxyTool
from fastmcp.server.tasks import TaskConfig
from fastmcp.tools.base import ToolResult
@pytest.fixture(autouse=True)
def reset_docket_memory_server():
"""Reset the shared Docket memory server between tests.
Docket uses a class-level FakeServer instance for memory:// URLs which
persists between tests, causing test isolation issues. This fixture
clears that shared state before each test.
"""
# Clear the shared FakeServer before each test
if hasattr(Docket, "_memory_server"):
delattr(Docket, "_memory_server")
yield
# Clean up after test as well
if hasattr(Docket, "_memory_server"):
delattr(Docket, "_memory_server")
@pytest.fixture
def child_server():
"""Create a child server with task-enabled components."""
mcp = FastMCP("child-server")
@mcp.tool(task=True)
async def multiply(a: int, b: int) -> int:
"""Multiply two numbers."""
return a * b
@mcp.tool(task=True)
async def slow_child_tool(duration: float = 0.1) -> str:
"""A child tool that takes time to execute."""
await asyncio.sleep(duration)
return "child completed"
@mcp.tool(task=False)
async def sync_child_tool(message: str) -> str:
"""Child tool that only supports synchronous execution."""
return f"child sync: {message}"
@mcp.prompt(task=True)
async def child_prompt(topic: str) -> str:
"""A child prompt that can execute as a task."""
return f"Here is information about {topic} from the child server."
@mcp.resource("child://data.txt", task=True)
async def child_resource() -> str:
"""A child resource that can be read as a task."""
return "Data from child server"
@mcp.resource("child://item/{item_id}.json", task=True)
async def child_item_resource(item_id: str) -> str:
"""A child resource template that can execute as a task."""
return f'{{"itemId": "{item_id}", "source": "child"}}'
return mcp
@pytest.fixture
def parent_server(child_server):
"""Create a parent server with the child mounted."""
parent = FastMCP("parent-server")
@parent.tool(task=True)
async def parent_tool(value: int) -> int:
"""A tool on the parent server."""
return value * 10
# Mount child with prefix
parent.mount(child_server, namespace="child")
return parent
@pytest.fixture
def parent_server_no_prefix(child_server):
"""Create a parent server with child mounted without prefix."""
parent = FastMCP("parent-no-prefix")
parent.mount(child_server) # No prefix
return parent
class TestMountedToolTasks:
"""Test task execution for mounted tools."""
async def test_mounted_tool_task_returns_task_object(self, parent_server):
"""Mounted tool called with task=True returns a task object."""
async with Client(parent_server) as client:
# Tool name is prefixed: child_multiply
task = await client.call_tool("child_multiply", {"a": 6, "b": 7}, task=True)
assert task is not None
assert hasattr(task, "task_id")
assert isinstance(task.task_id, str)
assert len(task.task_id) > 0
async def test_mounted_tool_task_executes_in_background(self, parent_server):
"""Mounted tool task executes in background."""
async with Client(parent_server) as client:
task = await client.call_tool("child_multiply", {"a": 3, "b": 4}, task=True)
# Should execute in background
assert not task.returned_immediately
async def test_mounted_tool_task_returns_correct_result(
self, parent_server: FastMCP
):
"""Mounted tool task returns correct result."""
async with Client(parent_server) as client:
task = await client.call_tool("child_multiply", {"a": 8, "b": 9}, task=True)
result = await task.result()
assert result.data == 72
async def test_mounted_tool_task_status(self, parent_server):
"""Can poll task status for mounted tool."""
async with Client(parent_server) as client:
task = await client.call_tool(
"child_slow_child_tool", {"duration": 0.5}, task=True
)
# Check status while running
status = await task.status()
assert status.status in ["working", "completed"]
# Wait for completion
await task.wait(timeout=2.0)
# Check status after completion
status = await task.status()
assert status.status == "completed"
@pytest.mark.timeout(10)
async def test_mounted_tool_task_cancellation(self, parent_server):
"""Can cancel a mounted tool task."""
async with Client(parent_server) as client:
task = await client.call_tool(
"child_slow_child_tool", {"duration": 10.0}, task=True
)
# Let it start
await asyncio.sleep(0.1)
# Cancel the task
await task.cancel()
# Check status
status = await task.status()
assert status.status == "cancelled"
async def test_graceful_degradation_sync_mounted_tool(self, parent_server):
"""Sync-only mounted tool returns error with task=True."""
async with Client(parent_server) as client:
task = await client.call_tool(
"child_sync_child_tool", {"message": "hello"}, task=True
)
# Should return immediately with an error
assert task.returned_immediately
result = await task.result()
assert result.is_error
async def test_parent_and_mounted_tools_both_work(self, parent_server):
"""Both parent and mounted tools work as tasks."""
async with Client(parent_server) as client:
# Parent tool
parent_task = await client.call_tool("parent_tool", {"value": 5}, task=True)
# Mounted tool
child_task = await client.call_tool(
"child_multiply", {"a": 2, "b": 3}, task=True
)
parent_result = await parent_task.result()
child_result = await child_task.result()
assert parent_result.data == 50
assert child_result.data == 6
class TestMountedToolTasksNoPrefix:
"""Test task execution for mounted tools without prefix."""
async def test_mounted_tool_without_prefix_task_works(
self, parent_server_no_prefix
):
"""Mounted tool without prefix works as task."""
async with Client(parent_server_no_prefix) as client:
# No prefix, so tool keeps original name
task = await client.call_tool("multiply", {"a": 5, "b": 6}, task=True)
assert not task.returned_immediately
result = await task.result()
assert result.data == 30
class TestMountedPromptTasks:
"""Test task execution for mounted prompts."""
async def test_mounted_prompt_task_returns_task_object(self, parent_server):
"""Mounted prompt called with task=True returns a task object."""
async with Client(parent_server) as client:
# Prompt name is prefixed: child_child_prompt
task = await client.get_prompt(
"child_child_prompt", {"topic": "FastMCP"}, task=True
)
assert task is not None
assert hasattr(task, "task_id")
assert isinstance(task.task_id, str)
async def test_mounted_prompt_task_executes_in_background(self, parent_server):
"""Mounted prompt task executes in background."""
async with Client(parent_server) as client:
task = await client.get_prompt(
"child_child_prompt", {"topic": "testing"}, task=True
)
assert not task.returned_immediately
async def test_mounted_prompt_task_returns_correct_result(
self, parent_server: FastMCP
):
"""Mounted prompt task returns correct result."""
async with Client(parent_server) as client:
task = await client.get_prompt(
"child_child_prompt", {"topic": "MCP protocol"}, task=True
)
result = await task.result()
assert "MCP protocol" in result.messages[0].content.text
assert "child server" in result.messages[0].content.text
class TestMountedResourceTasks:
"""Test task execution for mounted resources."""
async def test_mounted_resource_task_returns_task_object(self, parent_server):
"""Mounted resource read with task=True returns a task object."""
async with Client(parent_server) as client:
# Resource URI is prefixed: child://child/data.txt
task = await client.read_resource("child://child/data.txt", task=True)
assert task is not None
assert hasattr(task, "task_id")
assert isinstance(task.task_id, str)
async def test_mounted_resource_task_executes_in_background(self, parent_server):
"""Mounted resource task executes in background."""
async with Client(parent_server) as client:
task = await client.read_resource("child://child/data.txt", task=True)
assert not task.returned_immediately
async def test_mounted_resource_task_returns_correct_result(self, parent_server):
"""Mounted resource task returns correct result."""
async with Client(parent_server) as client:
task = await client.read_resource("child://child/data.txt", task=True)
result = await task.result()
assert len(result) > 0
assert "Data from child server" in result[0].text
async def test_mounted_resource_template_task(self, parent_server):
"""Mounted resource template with task=True works."""
async with Client(parent_server) as client:
task = await client.read_resource("child://child/item/99.json", task=True)
assert not task.returned_immediately
result = await task.result()
assert '"itemId": "99"' in result[0].text
assert '"source": "child"' in result[0].text
class TestMountedTaskDependencies:
"""Test that dependencies work correctly in mounted task execution."""
async def test_mounted_task_receives_docket_dependency(self):
"""Mounted tool task receives CurrentDocket dependency."""
child = FastMCP("dep-child")
received_docket = []
@child.tool(task=True)
async def tool_with_docket(docket: CurrentDocket = CurrentDocket()) -> str: # type: ignore[invalid-type-form] # ty:ignore[invalid-type-form]
received_docket.append(docket)
return f"docket available: {docket is not None}"
parent = FastMCP("dep-parent")
parent.mount(child, namespace="child")
async with Client(parent) as client:
task = await client.call_tool("child_tool_with_docket", {}, task=True)
result = await task.result()
assert "docket available: True" in str(result)
assert len(received_docket) == 1
assert received_docket[0] is not None
async def test_mounted_task_receives_server_dependency(self):
"""Mounted tool task receives CurrentFastMCP dependency."""
child = FastMCP("server-dep-child")
received_server = []
@child.tool(task=True)
async def tool_with_server(server: CurrentFastMCP = CurrentFastMCP()) -> str: # type: ignore[invalid-type-form] # ty:ignore[invalid-type-form]
received_server.append(server)
return f"server name: {server.name}"
parent = FastMCP("server-dep-parent")
parent.mount(child, namespace="child")
async with Client(parent) as client:
task = await client.call_tool("child_tool_with_server", {}, task=True)
await task.result()
assert len(received_server) == 1
assert received_server[0].name == "server-dep-child"
class TestMountedTaskServerContext:
"""Test that background tasks on mounted servers resolve to the child server (#3571)."""
async def test_current_fastmcp_resolves_to_child_server(self):
"""CurrentFastMCP() inside a mounted background task returns the child server."""
child = FastMCP("child")
received_server: list[FastMCP] = []
@child.tool(task=True)
async def whoami(server: CurrentFastMCP = CurrentFastMCP()) -> str: # type: ignore[invalid-type-form] # ty:ignore[invalid-type-form]
received_server.append(server)
return f"server name: {server.name}"
parent = FastMCP("parent")
parent.mount(child, namespace="child")
async with Client(parent) as client:
task = await client.call_tool("child_whoami", {}, task=True)
result = await task.result()
assert len(received_server) == 1
assert received_server[0].name == "child"
assert "server name: child" in str(result)
async def test_context_fastmcp_resolves_to_child_server(self):
"""ctx.fastmcp inside a mounted background task returns the child server."""
from fastmcp import Context
child = FastMCP("child")
received_server: list[FastMCP] = []
@child.tool(task=True)
async def whoami_ctx(ctx: Context) -> str:
received_server.append(ctx.fastmcp)
return f"context server: {ctx.fastmcp.name}"
parent = FastMCP("parent")
parent.mount(child, namespace="child")
async with Client(parent) as client:
task = await client.call_tool("child_whoami_ctx", {}, task=True)
result = await task.result()
assert len(received_server) == 1
assert received_server[0].name == "child"
assert "context server: child" in str(result)
async def test_nested_mount_resolves_to_innermost_server(self):
"""Doubly-nested mounts resolve to the innermost child server."""
grandchild = FastMCP("grandchild")
received_server: list[FastMCP] = []
@grandchild.tool(task=True)
async def deep_whoami(server: CurrentFastMCP = CurrentFastMCP()) -> str: # type: ignore[invalid-type-form] # ty:ignore[invalid-type-form]
received_server.append(server)
return f"server name: {server.name}"
child = FastMCP("child")
child.mount(grandchild, namespace="gc")
parent = FastMCP("parent")
parent.mount(child, namespace="child")
async with Client(parent) as client:
task = await client.call_tool("child_gc_deep_whoami", {}, task=True)
result = await task.result()
assert len(received_server) == 1
assert received_server[0].name == "grandchild"
assert "server name: grandchild" in str(result)
class TestMultipleMounts:
"""Test tasks with multiple mounted servers."""
async def test_tasks_work_with_multiple_mounts(self):
"""Tasks work correctly with multiple mounted servers."""
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.mount(child1, namespace="math1")
parent.mount(child2, namespace="math2")
async with Client(parent) as client:
task1 = await client.call_tool("math1_add", {"a": 10, "b": 5}, task=True)
task2 = await client.call_tool(
"math2_subtract", {"a": 10, "b": 5}, task=True
)
result1 = await task1.result()
result2 = await task2.result()
assert result1.data == 15
assert result2.data == 5
class TestMountedFunctionNameCollisions:
"""Test task execution when mounted servers have identically-named functions."""
async def test_multiple_mounts_with_same_function_names(self):
"""Two mounted servers with identically-named functions don't collide."""
child1 = FastMCP("child1")
child2 = FastMCP("child2")
@child1.tool(task=True)
async def process(value: int) -> int:
return value * 2 # Double
@child2.tool(task=True)
async def process(value: int) -> int: # noqa: F811
return value * 3 # Triple
parent = FastMCP("parent")
parent.mount(child1, namespace="c1")
parent.mount(child2, namespace="c2")
async with Client(parent) as client:
# Both should execute their own implementation
task1 = await client.call_tool("c1_process", {"value": 10}, task=True)
task2 = await client.call_tool("c2_process", {"value": 10}, task=True)
result1 = await task1.result()
result2 = await task2.result()
assert result1.data == 20 # child1's process (doubles)
assert result2.data == 30 # child2's process (triples)
async def test_no_prefix_mount_collision(self):
"""No-prefix mounts with same tool name - last mount wins."""
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.mount(child1) # No prefix
parent.mount(child2) # No prefix - overwrites child1's "process"
async with Client(parent) as client:
# Last mount wins - child2's process should execute
task = await client.call_tool("process", {"value": 10}, task=True)
result = await task.result()
assert result.data == 30 # child2's process (triples)
async def test_nested_mount_prefix_accumulation(self):
"""Nested mounts accumulate prefixes correctly for tasks."""
grandchild = FastMCP("gc")
child = FastMCP("child")
parent = FastMCP("parent")
@grandchild.tool(task=True)
async def deep_tool() -> str:
return "deep"
child.mount(grandchild, namespace="gc")
parent.mount(child, namespace="child")
async with Client(parent) as client:
# Tool should be accessible and execute correctly
task = await client.call_tool("child_gc_deep_tool", {}, task=True)
result = await task.result()
assert result.data == "deep"
class TestMountedTaskList:
"""Test task listing with mounted servers."""
async def test_list_tasks_includes_mounted_tasks(self, parent_server):
"""Task list includes tasks from mounted server tools."""
async with Client(parent_server) as client:
# Create tasks on both parent and mounted tools
parent_task = await client.call_tool("parent_tool", {"value": 1}, task=True)
child_task = await client.call_tool(
"child_multiply", {"a": 2, "b": 2}, task=True
)
# Wait for completion
await parent_task.wait(timeout=2.0)
await child_task.wait(timeout=2.0)
# List all tasks - returns dict with "tasks" key
tasks_response = await client.list_tasks()
task_ids = [t["taskId"] for t in tasks_response["tasks"]]
assert parent_task.task_id in task_ids
assert child_task.task_id in task_ids
class TestMountedTaskMetadata:
"""Test task metadata exposure for mounted tools."""
async def test_mounted_tool_list_preserves_task_support_metadata(self):
"""Mounted tools should preserve execution.taskSupport in tools/list."""
child = FastMCP("child")
@child.tool(task=True)
async def foo() -> dict[str, bool]:
return {"ok": True}
parent = FastMCP("parent")
parent.mount(child)
child_tools = await child.list_tools()
parent_tools = await parent.list_tools()
child_tool = next(t for t in child_tools if t.name == "foo")
parent_tool = next(t for t in parent_tools if t.name == "foo")
child_mcp_tool = child_tool.to_mcp_tool(name=child_tool.name)
parent_mcp_tool = parent_tool.to_mcp_tool(name=parent_tool.name)
assert child_mcp_tool.execution is not None
assert parent_mcp_tool.execution is not None
assert child_mcp_tool.execution.taskSupport == "optional"
assert parent_mcp_tool.execution.taskSupport == "optional"
async def test_proxy_tool_preserves_execution_metadata(self):
"""ProxyTool.from_mcp_tool should propagate execution.taskSupport (#3569)."""
mcp_tool = MCPTool(
name="remote_task_tool",
description="A remote tool that supports tasks",
inputSchema={"type": "object", "properties": {}},
execution=ToolExecution(taskSupport="optional"),
)
proxy = ProxyTool.from_mcp_tool(lambda: None, mcp_tool)
result = proxy.to_mcp_tool(name=proxy.name)
assert result.execution is not None
assert result.execution.taskSupport == "optional"
class TestMountedTaskConfigModes:
"""Test TaskConfig mode enforcement for mounted tools."""
@pytest.fixture
def child_with_modes(self):
"""Create a child server with tools in all three TaskConfig modes."""
mcp = FastMCP("child-modes", tasks=False)
@mcp.tool(task=TaskConfig(mode="optional"))
async def optional_tool() -> str:
"""Tool that supports both sync and task execution."""
return "optional result"
@mcp.tool(task=TaskConfig(mode="required"))
async def required_tool() -> str:
"""Tool that requires task execution."""
return "required result"
@mcp.tool(task=TaskConfig(mode="forbidden"))
async def forbidden_tool() -> str:
"""Tool that forbids task execution."""
return "forbidden result"
return mcp
@pytest.fixture
def parent_with_modes(self, child_with_modes):
"""Create a parent server with the child mounted."""
parent = FastMCP("parent-modes")
parent.mount(child_with_modes, namespace="child")
return parent
async def test_optional_mode_sync_through_mount(self, parent_with_modes):
"""Optional mode tool works without task through mount."""
async with Client(parent_with_modes) as client:
result = await client.call_tool("child_optional_tool", {})
assert "optional result" in str(result)
async def test_optional_mode_task_through_mount(self, parent_with_modes):
"""Optional mode tool works with task through mount."""
async with Client(parent_with_modes) as client:
task = await client.call_tool("child_optional_tool", {}, task=True)
assert task is not None
result = await task.result()
assert result.data == "optional result"
async def test_required_mode_with_task_through_mount(self, parent_with_modes):
"""Required mode tool succeeds with task through mount."""
async with Client(parent_with_modes) as client:
task = await client.call_tool("child_required_tool", {}, task=True)
assert task is not None
result = await task.result()
assert result.data == "required result"
async def test_required_mode_without_task_through_mount(self, parent_with_modes):
"""Required mode tool errors without task through mount."""
from fastmcp.exceptions import ToolError
async with Client(parent_with_modes) as client:
with pytest.raises(ToolError) as exc_info:
await client.call_tool("child_required_tool", {})
assert "requires task-augmented execution" in str(exc_info.value)
async def test_forbidden_mode_sync_through_mount(self, parent_with_modes):
"""Forbidden mode tool works without task through mount."""
async with Client(parent_with_modes) as client:
result = await client.call_tool("child_forbidden_tool", {})
assert "forbidden result" in str(result)
async def test_forbidden_mode_with_task_through_mount(self, parent_with_modes):
"""Forbidden mode tool degrades gracefully with task through mount."""
async with Client(parent_with_modes) as client:
task = await client.call_tool("child_forbidden_tool", {}, task=True)
# Should return immediately (graceful degradation)
assert task.returned_immediately
result = await task.result()
# Result is available but may indicate error or sync execution
assert result is not None
# -----------------------------------------------------------------------------
# Middleware classes for tracing tests
# -----------------------------------------------------------------------------
class ToolTracingMiddleware(Middleware):
"""Middleware that traces tool calls."""
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 ResourceTracingMiddleware(Middleware):
"""Middleware that traces resource reads."""
def __init__(self, name: str, calls: list[str]):
super().__init__()
self._name = name
self._calls = calls
async def on_read_resource(
self,
context: MiddlewareContext[mt.ReadResourceRequestParams],
call_next: CallNext[mt.ReadResourceRequestParams, ResourceResult],
) -> ResourceResult:
self._calls.append(f"{self._name}:before")
result = await call_next(context)
self._calls.append(f"{self._name}:after")
return result
class PromptTracingMiddleware(Middleware):
"""Middleware that traces prompt gets."""
def __init__(self, name: str, calls: list[str]):
super().__init__()
self._name = name
self._calls = calls
async def on_get_prompt(
self,
context: MiddlewareContext[mt.GetPromptRequestParams],
call_next: CallNext[mt.GetPromptRequestParams, PromptResult],
) -> PromptResult:
self._calls.append(f"{self._name}:before")
result = await call_next(context)
self._calls.append(f"{self._name}:after")
return result
class TestMiddlewareWithMountedTasks:
"""Test that middleware runs at all levels when executing background tasks.
For background tasks, middleware runs during task submission (wrapping the MCP
request handling that queues to Docket). The actual function execution happens
later in the Docket worker, after the middleware chain completes.
"""
async def test_tool_middleware_runs_with_background_task(self):
"""Middleware runs at parent, child, and grandchild levels for tool tasks."""
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.mount(child, namespace="c")
parent.add_middleware(ToolTracingMiddleware("parent", calls))
async with Client(parent) as client:
task = await client.call_tool("c_gc_compute", {"x": 5}, task=True)
result = await task.result()
assert result.data == 10
# Middleware runs during task submission (before/after queuing to Docket)
# Function executes later in Docket worker
assert calls == [
"parent:before",
"child:before",
"grandchild:before",
"grandchild:after",
"child:after",
"parent:after",
"grandchild:tool", # Executes in Docket after middleware completes
]
async def test_resource_middleware_runs_with_background_task(self):
"""Middleware runs at parent, child, and grandchild levels for resource tasks."""
calls: list[str] = []
grandchild = FastMCP("Grandchild")
@grandchild.resource("data://value", task=True)
async def get_data() -> str:
calls.append("grandchild:resource")
return "result"
grandchild.add_middleware(ResourceTracingMiddleware("grandchild", calls))
child = FastMCP("Child")
child.mount(grandchild, namespace="gc")
child.add_middleware(ResourceTracingMiddleware("child", calls))
parent = FastMCP("Parent")
parent.mount(child, namespace="c")
parent.add_middleware(ResourceTracingMiddleware("parent", calls))
async with Client(parent) as client:
task = await client.read_resource("data://c/gc/value", task=True)
result = await task.result()
assert result[0].text == "result"
# Middleware runs during task submission, function in Docket
assert calls == [
"parent:before",
"child:before",
"grandchild:before",
"grandchild:after",
"child:after",
"parent:after",
"grandchild:resource",
]
async def test_prompt_middleware_runs_with_background_task(self):
"""Middleware runs at parent, child, and grandchild levels for prompt tasks."""
calls: list[str] = []
grandchild = FastMCP("Grandchild")
@grandchild.prompt(task=True)
async def greet(name: str) -> str:
calls.append("grandchild:prompt")
return f"Hello, {name}!"
grandchild.add_middleware(PromptTracingMiddleware("grandchild", calls))
child = FastMCP("Child")
child.mount(grandchild, namespace="gc")
child.add_middleware(PromptTracingMiddleware("child", calls))
parent = FastMCP("Parent")
parent.mount(child, namespace="c")
parent.add_middleware(PromptTracingMiddleware("parent", calls))
async with Client(parent) as client:
task = await client.get_prompt("c_gc_greet", {"name": "World"}, task=True)
result = await task.result()
assert result.messages[0].content.text == "Hello, World!"
# Middleware runs during task submission, function in Docket
assert calls == [
"parent:before",
"child:before",
"grandchild:before",
"grandchild:after",
"child:after",
"parent:after",
"grandchild:prompt",
]
async def test_resource_template_middleware_runs_with_background_task(self):
"""Middleware runs at all levels for resource template tasks."""
calls: list[str] = []
grandchild = FastMCP("Grandchild")
@grandchild.resource("item://{id}", task=True)
async def get_item(id: str) -> str:
calls.append("grandchild:template")
return f"item-{id}"
grandchild.add_middleware(ResourceTracingMiddleware("grandchild", calls))
child = FastMCP("Child")
child.mount(grandchild, namespace="gc")
child.add_middleware(ResourceTracingMiddleware("child", calls))
parent = FastMCP("Parent")
parent.mount(child, namespace="c")
parent.add_middleware(ResourceTracingMiddleware("parent", calls))
async with Client(parent) as client:
task = await client.read_resource("item://c/gc/42", task=True)
result = await task.result()
assert result[0].text == "item-42"
# Middleware runs during task submission, function in Docket
assert calls == [
"parent:before",
"child:before",
"grandchild:before",
"grandchild:after",
"child:after",
"parent:after",
"grandchild:template",
]
class TestMountedTasksWithTaskMetaParameter:
"""Test mounted components called directly with task_meta parameter.
These tests verify the programmatic API where server.call_tool() or
server.read_resource() is called with an explicit task_meta parameter,
as opposed to using the Client with task=True.
Direct server calls require a running server context, so we use an outer
tool that makes the direct call internally.
"""
async def test_mounted_tool_with_task_meta_creates_task(self):
"""Mounted tool called with task_meta returns CreateTaskResult."""
from fastmcp.server.tasks.config import TaskMeta
child = FastMCP("Child")
@child.tool(task=True)
async def add(a: int, b: int) -> int:
return a + b
parent = FastMCP("Parent")
parent.mount(child, namespace="child")
@parent.tool
async def outer() -> str:
# Direct call with task_meta from within server context
result = await parent.call_tool(
"child_add", {"a": 2, "b": 3}, task_meta=TaskMeta(ttl=300)
)
return f"task:{result.task.taskId}"
async with Client(parent) as client:
result = await client.call_tool("outer", {})
assert "task:" in str(result)
async def test_mounted_resource_with_task_meta_creates_task(self):
"""Mounted resource called with task_meta returns CreateTaskResult."""
from fastmcp.server.tasks.config import TaskMeta
child = FastMCP("Child")
@child.resource("data://info", task=True)
async def get_info() -> str:
return "child info"
parent = FastMCP("Parent")
parent.mount(child, namespace="child")
@parent.tool
async def outer() -> str:
result = await parent.read_resource(
"data://child/info", task_meta=TaskMeta(ttl=300)
)
return f"task:{result.task.taskId}"
async with Client(parent) as client:
result = await client.call_tool("outer", {})
assert "task:" in str(result)
async def test_mounted_template_with_task_meta_creates_task(self):
"""Mounted resource template with task_meta returns CreateTaskResult."""
from fastmcp.server.tasks.config import TaskMeta
child = FastMCP("Child")
@child.resource("item://{id}", task=True)
async def get_item(id: str) -> str:
return f"item-{id}"
parent = FastMCP("Parent")
parent.mount(child, namespace="child")
@parent.tool
async def outer() -> str:
result = await parent.read_resource(
"item://child/42", task_meta=TaskMeta(ttl=300)
)
return f"task:{result.task.taskId}"
async with Client(parent) as client:
result = await client.call_tool("outer", {})
assert "task:" in str(result)
async def test_deeply_nested_tool_with_task_meta(self):
"""Three-level nested tool works with task_meta."""
from fastmcp.server.tasks.config import TaskMeta
grandchild = FastMCP("Grandchild")
@grandchild.tool(task=True)
async def compute(n: int) -> int:
return n * 3
child = FastMCP("Child")
child.mount(grandchild, namespace="gc")
parent = FastMCP("Parent")
parent.mount(child, namespace="c")
@parent.tool
async def outer() -> str:
result = await parent.call_tool(
"c_gc_compute", {"n": 7}, task_meta=TaskMeta(ttl=300)
)
return f"task:{result.task.taskId}"
async with Client(parent) as client:
result = await client.call_tool("outer", {})
assert "task:" in str(result)
async def test_deeply_nested_template_with_task_meta(self):
"""Three-level nested template works with task_meta."""
from fastmcp.server.tasks.config import TaskMeta
grandchild = FastMCP("Grandchild")
@grandchild.resource("doc://{name}", task=True)
async def get_doc(name: str) -> str:
return f"doc: {name}"
child = FastMCP("Child")
child.mount(grandchild, namespace="gc")
parent = FastMCP("Parent")
parent.mount(child, namespace="c")
@parent.tool
async def outer() -> str:
result = await parent.read_resource(
"doc://c/gc/readme", task_meta=TaskMeta(ttl=300)
)
return f"task:{result.task.taskId}"
async with Client(parent) as client:
result = await client.call_tool("outer", {})
assert "task:" in str(result)
async def test_mounted_prompt_with_task_meta_creates_task(self):
"""Mounted prompt called with task_meta returns CreateTaskResult."""
from fastmcp.server.tasks.config import TaskMeta
child = FastMCP("Child")
@child.prompt(task=True)
async def greet(name: str) -> str:
return f"Hello, {name}!"
parent = FastMCP("Parent")
parent.mount(child, namespace="child")
@parent.tool
async def outer() -> str:
result = await parent.render_prompt(
"child_greet", {"name": "World"}, task_meta=TaskMeta(ttl=300)
)
return f"task:{result.task.taskId}"
async with Client(parent) as client:
result = await client.call_tool("outer", {})
assert "task:" in str(result)
async def test_deeply_nested_prompt_with_task_meta(self):
"""Three-level nested prompt works with task_meta."""
from fastmcp.server.tasks.config import TaskMeta
grandchild = FastMCP("Grandchild")
@grandchild.prompt(task=True)
async def describe(topic: str) -> str:
return f"Information about {topic}"
child = FastMCP("Child")
child.mount(grandchild, namespace="gc")
parent = FastMCP("Parent")
parent.mount(child, namespace="c")
@parent.tool
async def outer() -> str:
result = await parent.render_prompt(
"c_gc_describe", {"topic": "FastMCP"}, task_meta=TaskMeta(ttl=300)
)
return f"task:{result.task.taskId}"
async with Client(parent) as client:
result = await client.call_tool("outer", {})
assert "task:" in str(result)
Reference in a new issue View git blame Copy permalink