From 4719f3055abc6e90aee209d37f259aa0f8f5d3a2 Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Mon, 4 May 2026 17:36:50 -0400 Subject: [PATCH] Make Docket reentrant; mounted servers enter their own lifespan (#4095) --- src/fastmcp/server/context.py | 26 ++-- src/fastmcp/server/mixins/lifespan.py | 56 +++++++- .../server/providers/fastmcp_provider.py | 35 ++++- tests/server/mount/test_advanced.py | 129 ++++++++++++++++++ 4 files changed, 224 insertions(+), 22 deletions(-) diff --git a/src/fastmcp/server/context.py b/src/fastmcp/server/context.py index cc66419e7..604ce301f 100644 --- a/src/fastmcp/server/context.py +++ b/src/fastmcp/server/context.py @@ -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 diff --git a/src/fastmcp/server/mixins/lifespan.py b/src/fastmcp/server/mixins/lifespan.py index b64d725ed..3cbb4094c 100644 --- a/src/fastmcp/server/mixins/lifespan.py +++ b/src/fastmcp/server/mixins/lifespan.py @@ -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 diff --git a/src/fastmcp/server/providers/fastmcp_provider.py b/src/fastmcp/server/providers/fastmcp_provider.py index 8ae2c18d3..9ce1ef6b0 100644 --- a/src/fastmcp/server/providers/fastmcp_provider.py +++ b/src/fastmcp/server/providers/fastmcp_provider.py @@ -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) diff --git a/tests/server/mount/test_advanced.py b/tests/server/mount/test_advanced.py index f9697b527..31c3d6c89 100644 --- a/tests/server/mount/test_advanced.py +++ b/tests/server/mount/test_advanced.py @@ -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.