From 0fc231965a634c931e4fd2d5c2d3aad9fe01c5d1 Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Sat, 14 Mar 2026 15:26:32 -0400 Subject: [PATCH 1/2] fix: task test teardown hanging 5s per test Fix context manager ordering in FastMCPTransport so the task group (server run + subscriptions) is cancelled before the lifespan (Docket Worker) tears down. Also break subscription loop on terminal task states. Closes #3498 --- src/fastmcp/client/transports/memory.py | 51 +++++++++-------- src/fastmcp/server/tasks/subscriptions.py | 11 +++- .../transports/test_memory_transport.py | 56 +++++++++++++++++++ 3 files changed, 95 insertions(+), 23 deletions(-) create mode 100644 tests/client/transports/test_memory_transport.py diff --git a/src/fastmcp/client/transports/memory.py b/src/fastmcp/client/transports/memory.py index 26b4e934a..7b2cbc8e4 100644 --- a/src/fastmcp/client/transports/memory.py +++ b/src/fastmcp/client/transports/memory.py @@ -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: diff --git a/src/fastmcp/server/tasks/subscriptions.py b/src/fastmcp/server/tasks/subscriptions.py index 9c6c8d59e..488b56a3e 100644 --- a/src/fastmcp/server/tasks/subscriptions.py +++ b/src/fastmcp/server/tasks/subscriptions.py @@ -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( diff --git a/tests/client/transports/test_memory_transport.py b/tests/client/transports/test_memory_transport.py new file mode 100644 index 000000000..5abdbdfb3 --- /dev/null +++ b/tests/client/transports/test_memory_transport.py @@ -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." + ) From 86fd5585d1f718421b8a926a0368c821a4a04bc5 Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Sat, 14 Mar 2026 15:41:50 -0400 Subject: [PATCH 2/2] perf: expose minimum_check_interval, reduce task pickup latency The Docket Worker polls for new tasks every minimum_check_interval (previously hardcoded to 250ms in pydocket). Expose this setting so users can tune it, default to 50ms, and override to 10ms in tests. This cuts average task pickup latency from ~125ms to ~5ms per task. --- src/fastmcp/server/mixins/lifespan.py | 1 + src/fastmcp/settings.py | 15 +++++++++++++++ tests/conftest.py | 10 +++++++++- 3 files changed, 25 insertions(+), 1 deletion(-) diff --git a/src/fastmcp/server/mixins/lifespan.py b/src/fastmcp/server/mixins/lifespan.py index 583944d61..a6c0cb9ed 100644 --- a/src/fastmcp/server/mixins/lifespan.py +++ b/src/fastmcp/server/mixins/lifespan.py @@ -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 diff --git a/src/fastmcp/settings.py b/src/fastmcp/settings.py index d98d7ed00..bf15aa488 100644 --- a/src/fastmcp/settings.py +++ b/src/fastmcp/settings.py @@ -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.""" diff --git a/tests/conftest.py b/tests/conftest.py index 1401d4141..9dd55bcb9 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -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