ceos/fastmcp
1
0
Fork 0
mirror of https://github.com/PrefectHQ/fastmcp.git synced 2026-08-28 10:18:08 +02:00
Code Issues Projects Releases Packages Wiki Activity

Spawn ping keepalive per Connection via asyncio task (SDK v2 has no per-connection task group)

Browse source
This commit is contained in:
Jeremiah Lowin 2026-07-05 23:33:11 -04:00
commit 72c636dd4c
No known key found for this signature in database
2 changed files with 58 additions and 60 deletions
Download patch file Download diff file Expand all files Collapse all files

43
fastmcp_slim/fastmcp/server/middleware/ping.py
View file

@ -1,5 +1,7 @@
"""Ping middleware for keeping client connections alive."""
import asyncio
import contextlib
from typing import Any
import anyio
@ -40,7 +42,7 @@ class PingMiddleware(Middleware):
self._lock = anyio.Lock()
async def on_message(self, context: MiddlewareContext, call_next: CallNext) -> Any:
"""Start ping task on first message from a session."""
"""Start ping task on first message from a connection."""
if (
context.fastmcp_context is None
or context.fastmcp_context.request_context is None
@ -48,29 +50,40 @@ class PingMiddleware(Middleware):
return await call_next(context)
session = context.fastmcp_context.session
session_id = id(session)
# SDK v2 constructs a ServerSession per request; the stable per-connection
# identity lives on the underlying Connection. Key the keepalive loop off
# it so one ping task runs for the whole connection and is torn down when
# the connection closes.
connection = getattr(session, "_connection", None)
connection_id = id(connection) if connection is not None else id(session)
async with self._lock:
if session_id not in self._active_sessions:
# SDK v2 constructs sessions per-request and no longer owns a
# per-connection task group FastMCP can borrow, so the keepalive
# ping loop cannot be spawned here. See Phase B report; this is a
# known regression pending a FastMCP-owned connection nursery.
tg = getattr(session, "_subscription_task_group", None)
if tg is not None:
self._active_sessions.add(session_id)
tg.start_soon(self._ping_loop, session, session_id)
if connection_id not in self._active_sessions:
self._active_sessions.add(connection_id)
ping_task = asyncio.create_task(
self._ping_loop(session, connection_id),
name=f"ping-keepalive-{connection_id}",
)
if connection is not None:
async def _cancel_ping() -> None:
ping_task.cancel()
with contextlib.suppress(asyncio.CancelledError):
await ping_task
connection.exit_stack.push_async_callback(_cancel_ping)
return await call_next(context)
async def _ping_loop(self, session: Any, session_id: int) -> None:
"""Send periodic pings until session ends."""
async def _ping_loop(self, session: Any, connection_id: int) -> None:
"""Send periodic pings until the connection ends."""
try:
while True:
await anyio.sleep(self.interval_ms / 1000)
try:
await session.send_ping()
except anyio.ClosedResourceError:
except (anyio.ClosedResourceError, anyio.BrokenResourceError):
return
finally:
self._active_sessions.discard(session_id)
self._active_sessions.discard(connection_id)

75
tests/server/middleware/test_ping.py
View file

@ -38,14 +38,26 @@ class TestPingMiddlewareInit:
class TestPingMiddlewareOnMessage:
"""Test on_message hook behavior."""
def _mock_session(self):
"""Build a mock session with a stable per-connection Connection.
SDK v2 constructs a ServerSession per request; PingMiddleware keys the
keepalive loop off the underlying Connection (and registers cleanup on
its exit stack), so tests supply a connection with an exit_stack.
"""
connection = MagicMock()
connection.exit_stack = MagicMock()
connection.exit_stack.push_async_callback = MagicMock()
session = MagicMock()
session._connection = connection
session.send_ping = AsyncMock()
return session, connection
async def test_starts_ping_task_on_first_message(self):
"""Test that ping task is started on first message from a session."""
"""Test that a ping task is started on first message from a connection."""
middleware = PingMiddleware(interval_ms=1000)
mock_session = MagicMock()
mock_session._subscription_task_group = MagicMock()
mock_session._subscription_task_group.start_soon = MagicMock()
mock_session, connection = self._mock_session()
mock_context = MagicMock()
mock_context.fastmcp_context.session = mock_session
@ -54,43 +66,34 @@ class TestPingMiddlewareOnMessage:
result = await middleware.on_message(mock_context, mock_call_next)
assert result == "result"
assert id(mock_session) in middleware._active_sessions
mock_session._subscription_task_group.start_soon.assert_called_once()
assert id(connection) in middleware._active_sessions
connection.exit_stack.push_async_callback.assert_called_once()
async def test_does_not_start_duplicate_task(self):
"""Test that duplicate messages from same session don't spawn duplicate tasks."""
"""Test that duplicate messages from same connection don't spawn duplicates."""
middleware = PingMiddleware(interval_ms=1000)
mock_session = MagicMock()
mock_session._subscription_task_group = MagicMock()
mock_session._subscription_task_group.start_soon = MagicMock()
mock_session, connection = self._mock_session()
mock_context = MagicMock()
mock_context.fastmcp_context.session = mock_session
mock_call_next = AsyncMock(return_value="result")
# First message
# Three messages from the same connection
await middleware.on_message(mock_context, mock_call_next)
# Second message from same session
await middleware.on_message(mock_context, mock_call_next)
# Third message from same session
await middleware.on_message(mock_context, mock_call_next)
# Should only start task once
assert mock_session._subscription_task_group.start_soon.call_count == 1
# Cleanup registered only once
assert connection.exit_stack.push_async_callback.call_count == 1
assert len(middleware._active_sessions) == 1
async def test_starts_separate_task_per_session(self):
"""Test that different sessions get separate ping tasks."""
"""Test that different connections get separate ping tasks."""
middleware = PingMiddleware(interval_ms=1000)
mock_session1 = MagicMock()
mock_session1._subscription_task_group = MagicMock()
mock_session1._subscription_task_group.start_soon = MagicMock()
mock_session2 = MagicMock()
mock_session2._subscription_task_group = MagicMock()
mock_session2._subscription_task_group.start_soon = MagicMock()
mock_session1, connection1 = self._mock_session()
mock_session2, connection2 = self._mock_session()
mock_context1 = MagicMock()
mock_context1.fastmcp_context.session = mock_session1
@ -103,28 +106,10 @@ class TestPingMiddlewareOnMessage:
await middleware.on_message(mock_context1, mock_call_next)
await middleware.on_message(mock_context2, mock_call_next)
mock_session1._subscription_task_group.start_soon.assert_called_once()
mock_session2._subscription_task_group.start_soon.assert_called_once()
connection1.exit_stack.push_async_callback.assert_called_once()
connection2.exit_stack.push_async_callback.assert_called_once()
assert len(middleware._active_sessions) == 2
async def test_skips_task_when_no_task_group(self):
"""Test graceful handling when session has no task group."""
middleware = PingMiddleware(interval_ms=1000)
mock_session = MagicMock()
mock_session._subscription_task_group = None
mock_context = MagicMock()
mock_context.fastmcp_context.session = mock_session
mock_call_next = AsyncMock(return_value="result")
result = await middleware.on_message(mock_context, mock_call_next)
assert result == "result"
# Session should NOT be added if task group is None
assert id(mock_session) not in middleware._active_sessions
async def test_skips_when_fastmcp_context_is_none(self):
"""Test that middleware passes through when fastmcp_context is None."""
middleware = PingMiddleware(interval_ms=1000)