From 218c49ce8def7522cb80240b21e2ce243c5da93c Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Sat, 12 Apr 2025 21:58:06 -0400 Subject: [PATCH] Enter mounted app lifespans --- src/fastmcp/server/server.py | 44 ++++++++++++++++++++++++++++----- tests/server/test_mount.py | 48 ++++++++++++++++++++++++++++++++++++ 2 files changed, 86 insertions(+), 6 deletions(-) diff --git a/src/fastmcp/server/server.py b/src/fastmcp/server/server.py index 3c3014518..875fc3d94 100644 --- a/src/fastmcp/server/server.py +++ b/src/fastmcp/server/server.py @@ -6,6 +6,7 @@ import re from collections.abc import AsyncIterator, Callable from contextlib import ( AbstractAsyncContextManager, + AsyncExitStack, asynccontextmanager, ) from typing import TYPE_CHECKING, Any, Generic, Literal @@ -18,7 +19,6 @@ from fastapi import FastAPI from mcp.server.lowlevel.helper_types import ReadResourceContents from mcp.server.lowlevel.server import LifespanResultT from mcp.server.lowlevel.server import Server as MCPServer -from mcp.server.lowlevel.server import lifespan as default_lifespan from mcp.server.session import ServerSession from mcp.server.sse import SseServerTransport from mcp.server.stdio import stdio_server @@ -56,6 +56,19 @@ if TYPE_CHECKING: logger = get_logger(__name__) +@asynccontextmanager +async def default_lifespan(server: "FastMCP") -> AsyncIterator[Any]: + """Default lifespan context manager that does nothing. + + Args: + server: The server instance this lifespan is managing + + Returns: + An empty context object + """ + yield {} + + def lifespan_wrapper( app: "FastMCP", lifespan: Callable[["FastMCP"], AbstractAsyncContextManager[LifespanResultT]], @@ -64,7 +77,18 @@ def lifespan_wrapper( ]: @asynccontextmanager async def wrap(s: MCPServer[LifespanResultT]) -> AsyncIterator[LifespanResultT]: - async with lifespan(app) as context: + async with AsyncExitStack() as stack: + # enter main app's lifespan + context = await stack.enter_async_context(lifespan(app)) + + # Enter all mounted app lifespans + for prefix, mounted_app in app._mounted_apps.items(): + mounted_context = mounted_app._mcp_server.lifespan( + mounted_app._mcp_server + ) + await stack.enter_async_context(mounted_context) + logger.debug(f"Prepared lifespan for mounted app '{prefix}'") + yield context return wrap @@ -84,10 +108,16 @@ class FastMCP(Generic[LifespanResultT]): self.tags: set[str] = tags or set() self.settings = fastmcp.settings.ServerSettings(**settings) + # Setup for mounted apps - must be initialized before _mcp_server + self._mounted_apps: dict[str, FastMCP] = {} + + if lifespan is None: + lifespan = default_lifespan + self._mcp_server = MCPServer[LifespanResultT]( name=name or "FastMCP", instructions=instructions, - lifespan=lifespan_wrapper(self, lifespan) if lifespan else default_lifespan, # type: ignore + lifespan=lifespan_wrapper(self, lifespan), ) self._tool_manager = ToolManager( duplicate_behavior=self.settings.on_duplicate_tools @@ -100,9 +130,6 @@ class FastMCP(Generic[LifespanResultT]): ) self.dependencies = self.settings.dependencies - # Setup for mounted apps - self._mounted_apps: dict[str, FastMCP] = {} - # Set up MCP protocol handlers self._setup_handlers() @@ -554,10 +581,15 @@ class FastMCP(Generic[LifespanResultT]): Example: If app has a template with URI "weather://location/{id}", it will be available as "weather+weather://location/{id}" - The prompts are imported with prefixed names using the prompt_separator Example: If app has a prompt named "weather_prompt", it will be available as "weather_weather_prompt" + - The mounted app's lifespan will be executed when the parent app's lifespan runs, + ensuring that any setup needed by the mounted app is performed Args: prefix: The prefix to use for the mounted application app: The FastMCP application to mount + tool_separator: Separator for tool names (defaults to "_") + resource_separator: Separator for resource URIs (defaults to "+") + prompt_separator: Separator for prompt names (defaults to "_") """ if tool_separator is None: tool_separator = "_" diff --git a/tests/server/test_mount.py b/tests/server/test_mount.py index 74a94f3dd..b48f15739 100644 --- a/tests/server/test_mount.py +++ b/tests/server/test_mount.py @@ -1,3 +1,7 @@ +import contextlib + +import pytest + from fastmcp.server.server import FastMCP @@ -182,3 +186,47 @@ async def test_mount_multiple_prompts(): # Verify prompts were imported with correct prefixes assert "python_review_python" in main_app._prompt_manager._prompts assert "sql_explain_sql" in main_app._prompt_manager._prompts + + +@pytest.mark.anyio +async def test_mount_lifespan(): + """Test that the lifespan of a mounted app is properly handled.""" + # Create apps + + lifespan_checkpoints = [] + + @contextlib.asynccontextmanager + async def lifespan(app: FastMCP): + lifespan_checkpoints.append(f"enter {app.name}") + try: + yield + finally: + lifespan_checkpoints.append(f"exit {app.name}") + + main_app = FastMCP("MainApp", lifespan=lifespan) + sub_app = FastMCP("SubApp", lifespan=lifespan) + sub_app_2 = FastMCP("SubApp2", lifespan=lifespan) + + main_app.mount("sub", sub_app) + main_app.mount("sub2", sub_app_2) + + low_level_server = main_app._mcp_server + async with contextlib.AsyncExitStack() as stack: + # Note: this imitates the way that lifespans are entered for mounted + # apps It is presently difficult to stop a running server + # programmatically without error in order to test the exit conditions, + # so this is the next best thing + await stack.enter_async_context(low_level_server.lifespan(low_level_server)) + assert lifespan_checkpoints == [ + "enter MainApp", + "enter SubApp", + "enter SubApp2", + ] + assert lifespan_checkpoints == [ + "enter MainApp", + "enter SubApp", + "enter SubApp2", + "exit SubApp2", + "exit SubApp", + "exit MainApp", + ]