Compare commits

...

2 commits

Author SHA1 Message Date
Jeremiah Lowin
86fd5585d1
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.
2026-03-14 15:41:50 -04:00
Jeremiah Lowin
0fc231965a
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
2026-03-14 15:26:46 -04:00
6 changed files with 120 additions and 24 deletions

View file

@ -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:

View file

@ -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

View file

@ -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(

View file

@ -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."""

View 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."
)

View file

@ -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