mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-21 21:14:17 +02:00
Enter mounted app lifespans
This commit is contained in:
parent
6c4d86fa13
commit
218c49ce8d
2 changed files with 86 additions and 6 deletions
|
|
@ -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 = "_"
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
]
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue