mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-09 15:19:10 +02:00
Compare commits
2 commits
main
...
fix/task-t
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
86fd5585d1 |
||
|
|
0fc231965a |
6 changed files with 120 additions and 24 deletions
|
|
@ -45,30 +45,37 @@ class FastMCPTransport(ClientTransport):
|
|||
# is called during cleanup, so we capture and re-raise manually.
|
||||
exception_to_raise: BaseException | None = None
|
||||
|
||||
async with (
|
||||
anyio.create_task_group() as tg,
|
||||
_enter_server_lifespan(server=self.server),
|
||||
):
|
||||
tg.start_soon(
|
||||
lambda: self.server._mcp_server.run(
|
||||
server_read,
|
||||
server_write,
|
||||
self.server._mcp_server.create_initialization_options(),
|
||||
raise_exceptions=self.raise_exceptions,
|
||||
# IMPORTANT: The lifespan MUST be the outer context and the task
|
||||
# group MUST be the inner context. This ensures the task group
|
||||
# (containing the server's run() and all its pub/sub subscriptions)
|
||||
# is cancelled and fully drained BEFORE the lifespan tears down
|
||||
# the Docket Worker and closes Redis connections. Reversing this
|
||||
# order (e.g. via `async with (tg, lifespan):`) causes the Worker
|
||||
# shutdown to hang for 5 seconds per test because fakeredis
|
||||
# blocking operations hold references that prevent clean
|
||||
# cancellation.
|
||||
async with _enter_server_lifespan(server=self.server): # noqa: SIM117
|
||||
async with anyio.create_task_group() as tg:
|
||||
tg.start_soon(
|
||||
lambda: self.server._mcp_server.run(
|
||||
server_read,
|
||||
server_write,
|
||||
self.server._mcp_server.create_initialization_options(),
|
||||
raise_exceptions=self.raise_exceptions,
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
try:
|
||||
async with ClientSession(
|
||||
read_stream=client_read,
|
||||
write_stream=client_write,
|
||||
**session_kwargs,
|
||||
) as client_session:
|
||||
yield client_session
|
||||
except BaseException as e:
|
||||
exception_to_raise = e
|
||||
finally:
|
||||
tg.cancel_scope.cancel()
|
||||
try:
|
||||
async with ClientSession(
|
||||
read_stream=client_read,
|
||||
write_stream=client_write,
|
||||
**session_kwargs,
|
||||
) as client_session:
|
||||
yield client_session
|
||||
except BaseException as e:
|
||||
exception_to_raise = e
|
||||
finally:
|
||||
tg.cancel_scope.cancel()
|
||||
|
||||
# Re-raise after task group has exited cleanly
|
||||
if exception_to_raise is not None:
|
||||
|
|
|
|||
|
|
@ -105,6 +105,7 @@ class LifespanMixin:
|
|||
"concurrency": settings.docket.concurrency,
|
||||
"redelivery_timeout": settings.docket.redelivery_timeout,
|
||||
"reconnection_delay": settings.docket.reconnection_delay,
|
||||
"minimum_check_interval": settings.docket.minimum_check_interval,
|
||||
}
|
||||
if settings.docket.worker_name:
|
||||
worker_kwargs["name"] = settings.docket.worker_name
|
||||
|
|
|
|||
|
|
@ -55,17 +55,26 @@ async def subscribe_to_task_updates(
|
|||
return
|
||||
|
||||
# Subscribe to state and progress events from Docket
|
||||
terminal_states = {
|
||||
ExecutionState.COMPLETED,
|
||||
ExecutionState.FAILED,
|
||||
ExecutionState.CANCELLED,
|
||||
}
|
||||
async for event in execution.subscribe():
|
||||
if event["type"] == "state":
|
||||
state = ExecutionState(event["state"])
|
||||
# Send notifications/tasks/status when state changes
|
||||
await _send_status_notification(
|
||||
session=session,
|
||||
task_id=task_id,
|
||||
task_key=task_key,
|
||||
docket=docket,
|
||||
state=ExecutionState(event["state"]),
|
||||
state=state,
|
||||
poll_interval_ms=poll_interval_ms,
|
||||
)
|
||||
# Stop subscribing once the task reaches a terminal state
|
||||
if state in terminal_states:
|
||||
break
|
||||
elif event["type"] == "progress":
|
||||
# Send notification when progress message changes
|
||||
await _send_progress_notification(
|
||||
|
|
|
|||
|
|
@ -117,6 +117,21 @@ class DocketSettings(BaseSettings):
|
|||
),
|
||||
] = timedelta(seconds=5)
|
||||
|
||||
minimum_check_interval: Annotated[
|
||||
timedelta,
|
||||
Field(
|
||||
description=inspect.cleandoc(
|
||||
"""
|
||||
How frequently the worker polls for new tasks. Lower
|
||||
values reduce latency for task pickup at the cost of
|
||||
more CPU usage. The default of 50ms is a good balance;
|
||||
increase for high-volume production deployments where
|
||||
tasks are long-running.
|
||||
"""
|
||||
),
|
||||
),
|
||||
] = timedelta(milliseconds=50)
|
||||
|
||||
|
||||
class Settings(BaseSettings):
|
||||
"""FastMCP settings."""
|
||||
|
|
|
|||
56
tests/client/transports/test_memory_transport.py
Normal file
56
tests/client/transports/test_memory_transport.py
Normal file
|
|
@ -0,0 +1,56 @@
|
|||
"""Tests for the in-memory FastMCPTransport.
|
||||
|
||||
These tests verify transport-level behavior that affects all tests using
|
||||
Client(server) with an in-process FastMCP server.
|
||||
"""
|
||||
|
||||
import time
|
||||
|
||||
import pytest
|
||||
|
||||
from fastmcp import Client, FastMCP
|
||||
|
||||
|
||||
@pytest.mark.timeout(10)
|
||||
async def test_task_teardown_does_not_hang():
|
||||
"""In-memory transport must tear down in under 2 seconds after a task call.
|
||||
|
||||
This is a regression test for a teardown ordering bug where the Docket
|
||||
Worker shutdown would hang for 5 seconds on every test that used
|
||||
task=True. The root cause was the server lifespan (which owns the Docket
|
||||
Worker) being torn down BEFORE the task group (which owns the server's
|
||||
run() and all its pub/sub subscriptions). Fakeredis blocking operations
|
||||
held by those subscriptions prevented the Worker's internal TaskGroup
|
||||
from cancelling its children, causing a 5-second stall until the
|
||||
Client's move_on_after(5) timeout fired.
|
||||
|
||||
The fix is to nest the task group INSIDE the lifespan context so that
|
||||
all server tasks (and their fakeredis resources) are cancelled and
|
||||
drained before Docket teardown begins.
|
||||
|
||||
If this test takes ~5 seconds, the context manager nesting in
|
||||
FastMCPTransport.connect_session() has been reversed — the lifespan
|
||||
must be the OUTER context and the task group must be the INNER context.
|
||||
"""
|
||||
mcp = FastMCP("teardown-test")
|
||||
|
||||
@mcp.tool(task=True)
|
||||
async def fast_tool(x: int) -> int:
|
||||
return x * 2
|
||||
|
||||
t0 = time.monotonic()
|
||||
|
||||
async with Client(mcp) as client:
|
||||
task = await client.call_tool("fast_tool", {"x": 21}, task=True)
|
||||
result = await task.result()
|
||||
assert result.data == 42
|
||||
|
||||
elapsed = time.monotonic() - t0
|
||||
|
||||
assert elapsed < 2.0, (
|
||||
f"Client teardown took {elapsed:.1f}s — expected <2s. "
|
||||
f"This usually means the context manager nesting in "
|
||||
f"FastMCPTransport.connect_session() is wrong: the lifespan "
|
||||
f"must be the OUTER context and the task group the INNER context. "
|
||||
f"See the comment in memory.py for details."
|
||||
)
|
||||
|
|
@ -3,6 +3,7 @@ import logging
|
|||
import socket
|
||||
import sys
|
||||
from collections.abc import Callable, Generator
|
||||
from datetime import timedelta
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
|
|
@ -60,11 +61,18 @@ def isolate_settings_home(tmp_path: Path):
|
|||
|
||||
This prevents file locking issues when multiple tests share the same
|
||||
storage directory in settings.home / "oauth-proxy".
|
||||
|
||||
Also sets a fast Docket polling interval for tests — the default 50ms
|
||||
is fine for production but still adds ~25ms average pickup latency per
|
||||
task. 10ms makes task tests near-instant.
|
||||
"""
|
||||
test_home = tmp_path / "fastmcp-test-home"
|
||||
test_home.mkdir(exist_ok=True)
|
||||
|
||||
with temporary_settings(home=test_home):
|
||||
with temporary_settings(
|
||||
home=test_home,
|
||||
docket__minimum_check_interval=timedelta(milliseconds=10),
|
||||
):
|
||||
yield
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue