ceos/fastmcp
1
0
Fork 0
mirror of https://github.com/PrefectHQ/fastmcp.git synced 2026-08-22 13:34:17 +02:00
Code Issues Projects Releases Packages Wiki Activity

Make Docket reentrant; mounted servers enter their own lifespan (#4095)

Browse source
This commit is contained in:
Jeremiah Lowin 2026-05-04 17:36:50 -04:00 • committed by GitHub
commit 4719f3055a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 226 additions and 24 deletions
Download patch file Download diff file Expand all files Collapse all files

26
src/fastmcp/server/context.py
View file

@ -354,13 +354,14 @@ class Context:
def lifespan_context(self) -> dict[str, Any]:
"""Access the server's lifespan context.
Returns the context dict yielded by the server's lifespan function.
Returns an empty dict if no lifespan was configured or if the MCP
session is not yet established.
Returns the context dict yielded by *this* server's lifespan function.
For a mounted child this is the child's own lifespan, not the parent's
— the MCP session always belongs to the parent, so reading from the
request context would return the parent's. We read directly from the
server's cached lifespan result instead, which is set by the
per-server ``_lifespan_manager`` regardless of mount position.
In background tasks (Docket workers), where request_context is not
available, falls back to reading from the FastMCP server's lifespan
result directly.
Returns an empty dict if no lifespan was configured.
Example:
```python
@ -372,13 +373,16 @@ class Context:
return "No database connection"
```
"""
result = self.fastmcp._lifespan_result
if result is not None:
return result
# Server's lifespan was never entered for this Context's server (or
# yielded None). Fall back to the request context's lifespan, which
# for a mounted child will be the parent's — preserved for parity
# with prior behavior, but in normal operation a child's own
# lifespan populates `_lifespan_result` and short-circuits above.
rc = self.request_context
if rc is None:
# In background tasks, request_context is not available.
# Fall back to the server's lifespan result directly (#3095).
result = self.fastmcp._lifespan_result
if result is not None:
return result
return {}
return rc.lifespan_context

56
src/fastmcp/server/mixins/lifespan.py
View file

@ -6,6 +6,7 @@ import asyncio
import weakref
from collections.abc import AsyncIterator
from contextlib import AsyncExitStack, asynccontextmanager, suppress
from contextvars import ContextVar
from typing import TYPE_CHECKING, Any
import anyio
@ -22,14 +23,35 @@ if TYPE_CHECKING:
logger = get_logger(__name__)
# Set True by `FastMCPProvider.lifespan` immediately before it enters the
# wrapped (mounted) server's `_lifespan_manager`, and reset on exit. The
# mounted server's `_docket_lifespan` reads this and becomes a no-op so that
# Docket / Worker / SharedContext are not re-initialized — there's one set
# per runtime tree, owned by the root.
#
# Independent servers entered as siblings (e.g. via `AsyncExitStack` in the
# same async context) are NOT in a parent/child relationship; the flag is not
# set in that case, so each independently establishes its own Docket and
# server context.
_lifespan_root_active: ContextVar[bool] = ContextVar(
"fastmcp_lifespan_root_active", default=False
)
class LifespanMixin:
"""Mixin providing lifespan and Docket task infrastructure for FastMCP."""
@property
def docket(self: FastMCP) -> Docket | None:
"""Get the Docket instance if Docket support is enabled.
"""The Docket instance owned by this server.
Returns None if Docket is not enabled or server hasn't been started yet.
Returns the Docket that this server initialized as the root of a
runtime tree. Mounted children do not own their own Docket — they
share the root's via ``_current_docket`` ContextVar inheritance —
so accessing ``.docket`` on a mounted child returns None even while
its tasks run on the root's Docket. For "the Docket in scope right
now," prefer reading ``_current_docket`` directly or use the
``CurrentDocket`` dependency injection.
"""
return self._docket
@ -37,13 +59,37 @@ class LifespanMixin:
async def _docket_lifespan(self: FastMCP) -> AsyncIterator[None]:
"""Manage Docket instance and Worker for background task execution.
Docket infrastructure is only initialized if:
Docket is process-level, not server-level: only the first server in a
runtime tree starts Docket and the Worker. Mounted children entered
via ``FastMCPProvider.lifespan`` see ``_lifespan_root_active=True``
(set by the provider before delegating to ``_lifespan_manager``) and
become no-ops, sharing the root's Docket via ``_current_docket``.
Independent servers entered as siblings — for example two unrelated
``FastMCP`` instances each entered through ``AsyncExitStack`` in the
same async context — are not in a parent/child relationship; no
provider has set the flag for them, so each runs the full root setup.
Docket infrastructure is only initialized at the root if:
1. pydocket is installed (fastmcp[tasks] extra)
2. There are task-enabled components (task_config.mode != 'forbidden')
This means users with pydocket installed but no task-enabled components
won't spin up Docket/Worker infrastructure.
Users with pydocket installed but no task-enabled components won't spin
up Docket / Worker infrastructure even at the root.
"""
# Nested entry: a parent in this runtime tree already owns Docket and
# SharedContext (the FastMCPProvider that mounted us set the flag).
# Stay out of their way and inherit via ContextVars.
if _lifespan_root_active.get():
yield
return
async with self._docket_lifespan_root():
yield
@asynccontextmanager
async def _docket_lifespan_root(self: FastMCP) -> AsyncIterator[None]:
"""Root-only Docket lifecycle. See _docket_lifespan for the dispatch."""
from fastmcp.server.dependencies import _current_server, is_docket_available
# Set FastMCP server in ContextVar so CurrentFastMCP can access it

35
src/fastmcp/server/providers/fastmcp_provider.py
View file

@ -698,11 +698,34 @@ class FastMCPProvider(Provider):
@asynccontextmanager
async def lifespan(self) -> AsyncIterator[None]:
"""Start the mounted server's user lifespan.
"""Start the mounted server's lifespan.
This starts only the wrapped server's user-defined lifespan, NOT its
full _lifespan_manager() (which includes Docket). The parent server's
Docket handles all background tasks.
Sets ``_lifespan_root_active=True`` to signal to the wrapped server's
``_docket_lifespan`` that it is running below an existing root in the
same runtime tree, then delegates to its full ``_lifespan_manager``.
The root's Docket / Worker / SharedContext are reused through
ContextVars (``_current_docket`` etc.); the mounted server's user
lifespan, ``_lifespan_result`` cache, and its own sub-providers
(nested mounts) all run normally.
The flag is reset as soon as ``_lifespan_manager`` finishes entering,
so it doesn't leak into the caller's async scope. Unrelated servers
entered later in the same task (e.g. siblings via ``AsyncExitStack``)
correctly see no active root and start their own infrastructure.
"""
async with self.server._lifespan(self.server):
yield
from fastmcp.server.mixins.lifespan import _lifespan_root_active
token = _lifespan_root_active.set(True)
flag_active = True
try:
async with self.server._lifespan_manager():
# Inner entry is complete; the flag's job (telling _docket_lifespan
# to no-op during _lifespan_manager's setup) is done. Reset now so
# unrelated lifespans entered later in this task aren't misclassified
# as nested.
_lifespan_root_active.reset(token)
flag_active = False
yield
finally:
if flag_active:
_lifespan_root_active.reset(token)

129
tests/server/mount/test_advanced.py
View file

@ -460,6 +460,135 @@ class TestToolNameOverrides:
)
class TestMountedServerLifespanContext:
"""Regression tests for mounted server lifespan_context resolution.
A mounted server's tools should see *their own* lifespan context via
``ctx.lifespan_context``, not the parent's. Previously the result yielded
by a mounted server's lifespan was discarded and ``ctx.lifespan_context``
fell through to the parent's MCP-session lifespan context.
"""
async def test_mounted_child_sees_own_lifespan_context(self):
"""A tool on a mounted child reads its own lifespan result."""
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from fastmcp.server.context import Context
@asynccontextmanager
async def parent_lifespan(_mcp: FastMCP) -> AsyncIterator[dict]:
yield {"who": "parent", "parent_only": "P"}
@asynccontextmanager
async def child_lifespan(_mcp: FastMCP) -> AsyncIterator[dict]:
yield {"who": "child", "child_only": "C"}
parent = FastMCP("Parent", lifespan=parent_lifespan)
child = FastMCP("Child", lifespan=child_lifespan)
@child.tool
def whoami(ctx: Context) -> dict:
return ctx.lifespan_context
parent.mount(child, "child")
async with Client(parent) as client:
result = await client.call_tool("child_whoami", {})
assert result.data == {"who": "child", "child_only": "C"}
async def test_independent_sibling_servers_each_get_own_lifecycle(self):
"""Two unrelated servers entered in the same async context are not nested.
Each must run its own ``_docket_lifespan_root`` and establish its own
Docket / server-context — they are not in a parent/child relationship,
only the same async stack.
"""
from contextlib import AsyncExitStack
from fastmcp.server.dependencies import _current_server
server_a = FastMCP("A")
server_b = FastMCP("B")
async with AsyncExitStack() as stack:
await stack.enter_async_context(server_a._lifespan_manager())
# While only A is active, _current_server resolves to A.
ref = _current_server.get()
assert ref is not None and ref() is server_a
await stack.enter_async_context(server_b._lifespan_manager())
# B established its own lifecycle; _current_server now resolves to B.
ref = _current_server.get()
assert ref is not None and ref() is server_b
async def test_unrelated_server_after_parent_with_mount(self):
"""A parent with a mounted child must not contaminate later entries.
``FastMCPProvider`` sets ``_lifespan_root_active`` while entering the
wrapped child. The flag must be cleared by the time the parent's
``_lifespan_manager`` yields, otherwise an unrelated server entered
later in the same task would see the flag and skip its own root setup.
"""
from contextlib import AsyncExitStack
from fastmcp.server.dependencies import _current_server
parent = FastMCP("Parent")
child = FastMCP("Child")
parent.mount(child, "child")
unrelated = FastMCP("Unrelated")
async with AsyncExitStack() as stack:
await stack.enter_async_context(parent._lifespan_manager())
# Parent's mount has run FastMCPProvider.lifespan; the flag set
# during child entry must have been reset by now.
await stack.enter_async_context(unrelated._lifespan_manager())
# Unrelated established its own lifecycle.
ref = _current_server.get()
assert ref is not None and ref() is unrelated
async def test_nested_grandchild_lifespan_runs(self):
"""A grandchild's lifespan is entered exactly once and visible to its tools."""
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from fastmcp.server.context import Context
events: list[str] = []
@asynccontextmanager
async def grandchild_lifespan(_mcp: FastMCP) -> AsyncIterator[dict]:
events.append("enter")
try:
yield {"who": "grandchild"}
finally:
events.append("exit")
parent = FastMCP("Parent")
child = FastMCP("Child")
grandchild = FastMCP("Grandchild", lifespan=grandchild_lifespan)
@grandchild.tool
def whoami(ctx: Context) -> dict:
return ctx.lifespan_context
child.mount(grandchild, "g")
parent.mount(child, "c")
async with Client(parent) as client:
result1 = await client.call_tool("c_g_whoami", {})
result2 = await client.call_tool("c_g_whoami", {})
assert result1.data == {"who": "grandchild"}
assert result2.data == {"who": "grandchild"}
# Lifespan is entered once at startup and exited once at teardown,
# not per-call and not per-mount level.
assert events == ["enter", "exit"]
class TestMountedServerDocketBehavior:
"""Regression tests for mounted server lifecycle behavior.