Jeremiah Lowin commited on
Commit
bd66e60
·
1 Parent(s): 9d54cca

Enter mounted app lifespans

Browse files
src/fastmcp/server/server.py CHANGED
@@ -6,6 +6,7 @@ import re
6
  from collections.abc import AsyncIterator, Callable
7
  from contextlib import (
8
  AbstractAsyncContextManager,
 
9
  asynccontextmanager,
10
  )
11
  from typing import TYPE_CHECKING, Any, Generic, Literal
@@ -18,7 +19,6 @@ from fastapi import FastAPI
18
  from mcp.server.lowlevel.helper_types import ReadResourceContents
19
  from mcp.server.lowlevel.server import LifespanResultT
20
  from mcp.server.lowlevel.server import Server as MCPServer
21
- from mcp.server.lowlevel.server import lifespan as default_lifespan
22
  from mcp.server.session import ServerSession
23
  from mcp.server.sse import SseServerTransport
24
  from mcp.server.stdio import stdio_server
@@ -56,6 +56,19 @@ if TYPE_CHECKING:
56
  logger = get_logger(__name__)
57
 
58
 
 
 
 
 
 
 
 
 
 
 
 
 
 
59
  def lifespan_wrapper(
60
  app: "FastMCP",
61
  lifespan: Callable[["FastMCP"], AbstractAsyncContextManager[LifespanResultT]],
@@ -64,7 +77,18 @@ def lifespan_wrapper(
64
  ]:
65
  @asynccontextmanager
66
  async def wrap(s: MCPServer[LifespanResultT]) -> AsyncIterator[LifespanResultT]:
67
- async with lifespan(app) as context:
 
 
 
 
 
 
 
 
 
 
 
68
  yield context
69
 
70
  return wrap
@@ -84,10 +108,16 @@ class FastMCP(Generic[LifespanResultT]):
84
  self.tags: set[str] = tags or set()
85
  self.settings = fastmcp.settings.ServerSettings(**settings)
86
 
 
 
 
 
 
 
87
  self._mcp_server = MCPServer[LifespanResultT](
88
  name=name or "FastMCP",
89
  instructions=instructions,
90
- lifespan=lifespan_wrapper(self, lifespan) if lifespan else default_lifespan, # type: ignore
91
  )
92
  self._tool_manager = ToolManager(
93
  duplicate_behavior=self.settings.on_duplicate_tools
@@ -100,9 +130,6 @@ class FastMCP(Generic[LifespanResultT]):
100
  )
101
  self.dependencies = self.settings.dependencies
102
 
103
- # Setup for mounted apps
104
- self._mounted_apps: dict[str, FastMCP] = {}
105
-
106
  # Set up MCP protocol handlers
107
  self._setup_handlers()
108
 
@@ -554,10 +581,15 @@ class FastMCP(Generic[LifespanResultT]):
554
  Example: If app has a template with URI "weather://location/{id}", it will be available as "weather+weather://location/{id}"
555
  - The prompts are imported with prefixed names using the prompt_separator
556
  Example: If app has a prompt named "weather_prompt", it will be available as "weather_weather_prompt"
 
 
557
 
558
  Args:
559
  prefix: The prefix to use for the mounted application
560
  app: The FastMCP application to mount
 
 
 
561
  """
562
  if tool_separator is None:
563
  tool_separator = "_"
 
6
  from collections.abc import AsyncIterator, Callable
7
  from contextlib import (
8
  AbstractAsyncContextManager,
9
+ AsyncExitStack,
10
  asynccontextmanager,
11
  )
12
  from typing import TYPE_CHECKING, Any, Generic, Literal
 
19
  from mcp.server.lowlevel.helper_types import ReadResourceContents
20
  from mcp.server.lowlevel.server import LifespanResultT
21
  from mcp.server.lowlevel.server import Server as MCPServer
 
22
  from mcp.server.session import ServerSession
23
  from mcp.server.sse import SseServerTransport
24
  from mcp.server.stdio import stdio_server
 
56
  logger = get_logger(__name__)
57
 
58
 
59
+ @asynccontextmanager
60
+ async def default_lifespan(server: "FastMCP") -> AsyncIterator[Any]:
61
+ """Default lifespan context manager that does nothing.
62
+
63
+ Args:
64
+ server: The server instance this lifespan is managing
65
+
66
+ Returns:
67
+ An empty context object
68
+ """
69
+ yield {}
70
+
71
+
72
  def lifespan_wrapper(
73
  app: "FastMCP",
74
  lifespan: Callable[["FastMCP"], AbstractAsyncContextManager[LifespanResultT]],
 
77
  ]:
78
  @asynccontextmanager
79
  async def wrap(s: MCPServer[LifespanResultT]) -> AsyncIterator[LifespanResultT]:
80
+ async with AsyncExitStack() as stack:
81
+ # enter main app's lifespan
82
+ context = await stack.enter_async_context(lifespan(app))
83
+
84
+ # Enter all mounted app lifespans
85
+ for prefix, mounted_app in app._mounted_apps.items():
86
+ mounted_context = mounted_app._mcp_server.lifespan(
87
+ mounted_app._mcp_server
88
+ )
89
+ await stack.enter_async_context(mounted_context)
90
+ logger.debug(f"Prepared lifespan for mounted app '{prefix}'")
91
+
92
  yield context
93
 
94
  return wrap
 
108
  self.tags: set[str] = tags or set()
109
  self.settings = fastmcp.settings.ServerSettings(**settings)
110
 
111
+ # Setup for mounted apps - must be initialized before _mcp_server
112
+ self._mounted_apps: dict[str, FastMCP] = {}
113
+
114
+ if lifespan is None:
115
+ lifespan = default_lifespan
116
+
117
  self._mcp_server = MCPServer[LifespanResultT](
118
  name=name or "FastMCP",
119
  instructions=instructions,
120
+ lifespan=lifespan_wrapper(self, lifespan),
121
  )
122
  self._tool_manager = ToolManager(
123
  duplicate_behavior=self.settings.on_duplicate_tools
 
130
  )
131
  self.dependencies = self.settings.dependencies
132
 
 
 
 
133
  # Set up MCP protocol handlers
134
  self._setup_handlers()
135
 
 
581
  Example: If app has a template with URI "weather://location/{id}", it will be available as "weather+weather://location/{id}"
582
  - The prompts are imported with prefixed names using the prompt_separator
583
  Example: If app has a prompt named "weather_prompt", it will be available as "weather_weather_prompt"
584
+ - The mounted app's lifespan will be executed when the parent app's lifespan runs,
585
+ ensuring that any setup needed by the mounted app is performed
586
 
587
  Args:
588
  prefix: The prefix to use for the mounted application
589
  app: The FastMCP application to mount
590
+ tool_separator: Separator for tool names (defaults to "_")
591
+ resource_separator: Separator for resource URIs (defaults to "+")
592
+ prompt_separator: Separator for prompt names (defaults to "_")
593
  """
594
  if tool_separator is None:
595
  tool_separator = "_"
tests/server/test_mount.py CHANGED
@@ -1,3 +1,7 @@
 
 
 
 
1
  from fastmcp.server.server import FastMCP
2
 
3
 
@@ -182,3 +186,47 @@ async def test_mount_multiple_prompts():
182
  # Verify prompts were imported with correct prefixes
183
  assert "python_review_python" in main_app._prompt_manager._prompts
184
  assert "sql_explain_sql" in main_app._prompt_manager._prompts
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import contextlib
2
+
3
+ import pytest
4
+
5
  from fastmcp.server.server import FastMCP
6
 
7
 
 
186
  # Verify prompts were imported with correct prefixes
187
  assert "python_review_python" in main_app._prompt_manager._prompts
188
  assert "sql_explain_sql" in main_app._prompt_manager._prompts
189
+
190
+
191
+ @pytest.mark.anyio
192
+ async def test_mount_lifespan():
193
+ """Test that the lifespan of a mounted app is properly handled."""
194
+ # Create apps
195
+
196
+ lifespan_checkpoints = []
197
+
198
+ @contextlib.asynccontextmanager
199
+ async def lifespan(app: FastMCP):
200
+ lifespan_checkpoints.append(f"enter {app.name}")
201
+ try:
202
+ yield
203
+ finally:
204
+ lifespan_checkpoints.append(f"exit {app.name}")
205
+
206
+ main_app = FastMCP("MainApp", lifespan=lifespan)
207
+ sub_app = FastMCP("SubApp", lifespan=lifespan)
208
+ sub_app_2 = FastMCP("SubApp2", lifespan=lifespan)
209
+
210
+ main_app.mount("sub", sub_app)
211
+ main_app.mount("sub2", sub_app_2)
212
+
213
+ low_level_server = main_app._mcp_server
214
+ async with contextlib.AsyncExitStack() as stack:
215
+ # Note: this imitates the way that lifespans are entered for mounted
216
+ # apps It is presently difficult to stop a running server
217
+ # programmatically without error in order to test the exit conditions,
218
+ # so this is the next best thing
219
+ await stack.enter_async_context(low_level_server.lifespan(low_level_server))
220
+ assert lifespan_checkpoints == [
221
+ "enter MainApp",
222
+ "enter SubApp",
223
+ "enter SubApp2",
224
+ ]
225
+ assert lifespan_checkpoints == [
226
+ "enter MainApp",
227
+ "enter SubApp",
228
+ "enter SubApp2",
229
+ "exit SubApp2",
230
+ "exit SubApp",
231
+ "exit MainApp",
232
+ ]