mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-19 12:04:18 +02:00
110 lines
5 KiB
Python
110 lines
5 KiB
Python
"""Regression test for driving the FastMCP lifespan through the SDK session manager.
|
|
|
|
Before FastMCP handed its lifespan to the SDK lowlevel Server (PR #4446), the
|
|
SDK v1 lifespan was effectively session-scoped and FastMCP worked around it by
|
|
driving its own ``_lifespan_manager`` beside the session manager. The SDK v2
|
|
``StreamableHTTPSessionManager`` now enters ``app.lifespan(app)`` -- FastMCP's
|
|
``_lifespan_proxy`` -- exactly once for the manager's lifetime and reuses the
|
|
yielded state for every session. The user lifespan must therefore fire once per
|
|
process and persist across HTTP client sessions, not once per session.
|
|
|
|
The invariant that actually matters is "the session manager drives the lifespan
|
|
exactly once, regardless of how many client sessions connect." A plain
|
|
user-lifespan enter/exit counter cannot guard it: ``_lifespan_manager`` is
|
|
ref-counted, and ``run_http_async`` opens an *outer* ``_lifespan_manager``
|
|
around uvicorn. That outer entry holds the ref count at >= 1 for the whole
|
|
server lifetime, so even if the session manager regressed to re-entering
|
|
``_lifespan_proxy`` once per session, the user lifespan would still be entered
|
|
exactly once (the proxy's nested ``_lifespan_manager`` entries would all reuse
|
|
the outer result). The user counter would stay ``1`` while the behavior it
|
|
claims to guard was broken.
|
|
|
|
So this test spies on the session-manager entry point directly -- it counts how
|
|
many times the SDK enters ``server._mcp_server.lifespan`` (the ``_lifespan_proxy``
|
|
wrapper) -- and asserts that count is exactly one across sequential and
|
|
overlapping sessions. A regression that moves ``app.lifespan(app)`` into the
|
|
per-session code path makes this count grow with the session count and fails
|
|
loudly. The user enter/exit counter is kept as a secondary check on the
|
|
"entered once, exited once at shutdown" shape.
|
|
"""
|
|
|
|
from collections.abc import AsyncIterator
|
|
from contextlib import asynccontextmanager
|
|
from typing import Any
|
|
|
|
from mcp.server.lowlevel.server import Server as LowLevelServer
|
|
|
|
from fastmcp import Client, FastMCP
|
|
from fastmcp.client.transports import StreamableHttpTransport
|
|
from fastmcp.utilities.tests import run_server_async
|
|
|
|
|
|
async def test_http_user_lifespan_fires_once_across_sessions():
|
|
"""The session manager must drive the user lifespan exactly once for the
|
|
server process, even when several independent HTTP client sessions connect
|
|
and disconnect.
|
|
"""
|
|
enter_count = 0
|
|
exit_count = 0
|
|
|
|
@asynccontextmanager
|
|
async def counting_lifespan(mcp: FastMCP) -> AsyncIterator[dict[str, Any]]:
|
|
nonlocal enter_count, exit_count
|
|
enter_count += 1
|
|
try:
|
|
yield {"initialized": True}
|
|
finally:
|
|
exit_count += 1
|
|
|
|
server = FastMCP("LifespanOnceServer", lifespan=counting_lifespan)
|
|
|
|
@server.tool
|
|
def ping() -> str:
|
|
return "pong"
|
|
|
|
# Spy on the session manager's single entry point into FastMCP's lifespan.
|
|
# `StreamableHTTPSessionManager.run()` calls `self.app.lifespan(self.app)`
|
|
# exactly once and reuses the yielded state per session; `self.app` is
|
|
# `server._mcp_server`, so `server._mcp_server.lifespan` is the
|
|
# `_lifespan_proxy` wrapper. Counting entries here asserts the invariant
|
|
# directly, independent of `_lifespan_manager`'s ref-count masking.
|
|
proxy_enter_count = 0
|
|
original_lifespan = server._mcp_server.lifespan
|
|
|
|
@asynccontextmanager
|
|
async def counting_proxy(app: LowLevelServer[Any]) -> AsyncIterator[Any]:
|
|
nonlocal proxy_enter_count
|
|
proxy_enter_count += 1
|
|
async with original_lifespan(app) as state:
|
|
yield state
|
|
|
|
server._mcp_server.lifespan = counting_proxy
|
|
|
|
async with run_server_async(server, transport="http") as mcp_url:
|
|
# `run_server_async` yields a URL that already includes the `/mcp` path.
|
|
# Three separate, sequential client sessions against the same process.
|
|
for _ in range(3):
|
|
async with Client(StreamableHttpTransport(mcp_url)) as client:
|
|
result = await client.call_tool("ping", {})
|
|
assert result.data == "pong"
|
|
# The session manager must not re-drive the lifespan when a session
|
|
# closes -- it owns a single entry for the whole process lifetime.
|
|
assert proxy_enter_count == 1
|
|
assert exit_count == 0
|
|
|
|
# Overlapping sessions must also observe a single, still-open lifespan.
|
|
async with (
|
|
Client(StreamableHttpTransport(mcp_url)) as c1,
|
|
Client(StreamableHttpTransport(mcp_url)) as c2,
|
|
):
|
|
assert (await c1.call_tool("ping", {})).data == "pong"
|
|
assert (await c2.call_tool("ping", {})).data == "pong"
|
|
assert proxy_enter_count == 1
|
|
assert exit_count == 0
|
|
|
|
# The session manager entered the lifespan exactly once across every
|
|
# session, and the user lifespan was entered once and exited once at
|
|
# process shutdown.
|
|
assert proxy_enter_count == 1
|
|
assert enter_count == 1
|
|
assert exit_count == 1
|