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/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/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/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/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." + ) 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