Add SSE polling support (SEP-1699) (#1654)

This commit is contained in:
Felix Weinberger 2025-12-02 11:44:49 +00:00 • committed by GitHub
commit 281fd4765e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 1461 additions and 45 deletions

View file

@ -153,6 +153,7 @@ class FastMCP(Generic[LifespanResultT]):
auth_server_provider: (OAuthAuthorizationServerProvider[Any, Any, Any] | None) = None,
token_verifier: TokenVerifier | None = None,
event_store: EventStore | None = None,
retry_interval: int | None = None,
*,
tools: list[Tool] | None = None,
debug: bool = False,
@ -221,6 +222,7 @@ class FastMCP(Generic[LifespanResultT]):
if auth_server_provider and not token_verifier: # pragma: no cover
self._token_verifier = ProviderTokenVerifier(auth_server_provider)
self._event_store = event_store
self._retry_interval = retry_interval
self._custom_starlette_routes: list[Route] = []
self.dependencies = self.settings.dependencies
self._session_manager: StreamableHTTPSessionManager | None = None
@ -940,6 +942,7 @@ class FastMCP(Generic[LifespanResultT]):
self._session_manager = StreamableHTTPSessionManager(
app=self._mcp_server,
event_store=self._event_store,
retry_interval=self._retry_interval,
json_response=self.settings.json_response,
stateless=self.settings.stateless_http, # Use the stateless setting
security_settings=self.settings.transport_security,
@ -1282,6 +1285,38 @@ class Context(BaseModel, Generic[ServerSessionT, LifespanContextT, RequestT]):
"""Access to the underlying session for advanced usage."""
return self.request_context.session
async def close_sse_stream(self) -> None:
"""Close the SSE stream to trigger client reconnection.
This method closes the HTTP connection for the current request, triggering
client reconnection. Events continue to be stored in the event store and will
be replayed when the client reconnects with Last-Event-ID.
Use this to implement polling behavior during long-running operations -
client will reconnect after the retry interval specified in the priming event.
Note:
This is a no-op if not using StreamableHTTP transport with event_store.
The callback is only available when event_store is configured.
"""
if self._request_context and self._request_context.close_sse_stream: # pragma: no cover
await self._request_context.close_sse_stream()
async def close_standalone_sse_stream(self) -> None:
"""Close the standalone GET SSE stream to trigger client reconnection.
This method closes the HTTP connection for the standalone GET stream used
for unsolicited server-to-client notifications. The client SHOULD reconnect
with Last-Event-ID to resume receiving notifications.
Note:
This is a no-op if not using StreamableHTTP transport with event_store.
Currently, client reconnection for standalone GET streams is NOT
implemented - this is a known gap.
"""
if self._request_context and self._request_context.close_standalone_sse_stream: # pragma: no cover
await self._request_context.close_standalone_sse_stream()
# Convenience methods for common log levels
async def debug(self, message: str, **extra: Any) -> None:
"""Send a debug log message."""

View file

@ -713,12 +713,16 @@ class Server(Generic[LifespanResultT, RequestT]):
token = None
try:
# Extract request context from message metadata
# Extract request context and close_sse_stream from message metadata
request_data = None
close_sse_stream_cb = None
close_standalone_sse_stream_cb = None
if message.message_metadata is not None and isinstance(
message.message_metadata, ServerMessageMetadata
): # pragma: no cover
request_data = message.message_metadata.request_context
close_sse_stream_cb = message.message_metadata.close_sse_stream
close_standalone_sse_stream_cb = message.message_metadata.close_standalone_sse_stream
# Set our global state that can be retrieved via
# app.get_request_context()
@ -741,6 +745,8 @@ class Server(Generic[LifespanResultT, RequestT]):
_task_support=task_support,
),
request=request_data,
close_sse_stream=close_sse_stream_cb,
close_standalone_sse_stream=close_standalone_sse_stream_cb,
)
)
response = await handler(req)

View file

@ -15,6 +15,7 @@ from collections.abc import AsyncGenerator, Awaitable, Callable
from contextlib import asynccontextmanager
from dataclasses import dataclass
from http import HTTPStatus
from typing import Any
import anyio
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
@ -87,13 +88,13 @@ class EventStore(ABC):
"""
@abstractmethod
async def store_event(self, stream_id: StreamId, message: JSONRPCMessage) -> EventId:
async def store_event(self, stream_id: StreamId, message: JSONRPCMessage | None) -> EventId:
"""
Stores an event for later retrieval.
Args:
stream_id: ID of the stream the event belongs to
message: The JSON-RPC message to store
message: The JSON-RPC message to store, or None for priming events
Returns:
The generated event ID for the stored event
@ -140,6 +141,7 @@ class StreamableHTTPServerTransport:
is_json_response_enabled: bool = False,
event_store: EventStore | None = None,
security_settings: TransportSecuritySettings | None = None,
retry_interval: int | None = None,
) -> None:
"""
Initialize a new StreamableHTTP server transport.
@ -153,6 +155,10 @@ class StreamableHTTPServerTransport:
resumability will be enabled, allowing clients to
reconnect and resume messages.
security_settings: Optional security settings for DNS rebinding protection.
retry_interval: Retry interval in milliseconds to suggest to clients in SSE
retry field. When set, the server will send a retry field in
SSE priming events to control client reconnection timing for
polling behavior. Only used when event_store is provided.
Raises:
ValueError: If the session ID contains invalid characters.
@ -164,6 +170,7 @@ class StreamableHTTPServerTransport:
self.is_json_response_enabled = is_json_response_enabled
self._event_store = event_store
self._security = TransportSecurityMiddleware(security_settings)
self._retry_interval = retry_interval
self._request_streams: dict[
RequestId,
tuple[
@ -171,6 +178,7 @@ class StreamableHTTPServerTransport:
MemoryObjectReceiveStream[EventMessage],
],
] = {}
self._sse_stream_writers: dict[RequestId, MemoryObjectSendStream[dict[str, str]]] = {}
self._terminated = False
@property
@ -178,6 +186,91 @@ class StreamableHTTPServerTransport:
"""Check if this transport has been explicitly terminated."""
return self._terminated
def close_sse_stream(self, request_id: RequestId) -> None: # pragma: no cover
"""Close SSE connection for a specific request without terminating the stream.
This method closes the HTTP connection for the specified request, triggering
client reconnection. Events continue to be stored in the event store and will
be replayed when the client reconnects with Last-Event-ID.
Use this to implement polling behavior during long-running operations -
client will reconnect after the retry interval specified in the priming event.
Args:
request_id: The request ID whose SSE stream should be closed.
Note:
This is a no-op if there is no active stream for the request ID.
Requires event_store to be configured for events to be stored during
the disconnect.
"""
writer = self._sse_stream_writers.pop(request_id, None)
if writer:
writer.close()
# Also close and remove request streams
if request_id in self._request_streams:
send_stream, receive_stream = self._request_streams.pop(request_id)
send_stream.close()
receive_stream.close()
def close_standalone_sse_stream(self) -> None: # pragma: no cover
"""Close the standalone GET SSE stream, triggering client reconnection.
This method closes the HTTP connection for the standalone GET stream used
for unsolicited server-to-client notifications. The client SHOULD reconnect
with Last-Event-ID to resume receiving notifications.
Use this to implement polling behavior for the notification stream -
client will reconnect after the retry interval specified in the priming event.
Note:
This is a no-op if there is no active standalone SSE stream.
Requires event_store to be configured for events to be stored during
the disconnect.
Currently, client reconnection for standalone GET streams is NOT
implemented - this is a known gap (see test_standalone_get_stream_reconnection).
"""
self.close_sse_stream(GET_STREAM_KEY)
def _create_session_message( # pragma: no cover
self,
message: JSONRPCMessage,
request: Request,
request_id: RequestId,
) -> SessionMessage:
"""Create a session message with metadata including close_sse_stream callback."""
async def close_stream_callback() -> None:
self.close_sse_stream(request_id)
async def close_standalone_stream_callback() -> None:
self.close_standalone_sse_stream()
metadata = ServerMessageMetadata(
request_context=request,
close_sse_stream=close_stream_callback,
close_standalone_sse_stream=close_standalone_stream_callback,
)
return SessionMessage(message, metadata=metadata)
async def _send_priming_event( # pragma: no cover
self,
request_id: RequestId,
sse_stream_writer: MemoryObjectSendStream[dict[str, Any]],
) -> None:
"""Send priming event for SSE resumability if event_store is configured."""
if not self._event_store:
return
priming_event_id = await self._event_store.store_event(
str(request_id), # Convert RequestId to StreamId (str)
None, # Priming event has no payload
)
priming_event: dict[str, str | int] = {"id": priming_event_id, "data": ""}
if self._retry_interval is not None:
priming_event["retry"] = self._retry_interval
await sse_stream_writer.send(priming_event)
def _create_error_response(
self,
error_message: str,
@ -459,10 +552,16 @@ class StreamableHTTPServerTransport:
# Create SSE stream
sse_stream_writer, sse_stream_reader = anyio.create_memory_object_stream[dict[str, str]](0)
# Store writer reference so close_sse_stream() can close it
self._sse_stream_writers[request_id] = sse_stream_writer
async def sse_writer():
# Get the request ID from the incoming request message
try:
async with sse_stream_writer, request_stream_reader:
# Send priming event for SSE resumability
await self._send_priming_event(request_id, sse_stream_writer)
# Process messages from the request-specific stream
async for event_message in request_stream_reader:
# Build the event data
@ -475,10 +574,14 @@ class StreamableHTTPServerTransport:
JSONRPCResponse | JSONRPCError,
):
break
except anyio.ClosedResourceError:
# Expected when close_sse_stream() is called
logger.debug("SSE stream closed by close_sse_stream()")
except Exception:
logger.exception("Error in SSE writer")
finally:
logger.debug("Closing SSE writer")
self._sse_stream_writers.pop(request_id, None)
await self._clean_up_memory_streams(request_id)
# Create and start EventSourceResponse
@ -502,8 +605,7 @@ class StreamableHTTPServerTransport:
async with anyio.create_task_group() as tg:
tg.start_soon(response, scope, receive, send)
# Then send the message to be processed by the server
metadata = ServerMessageMetadata(request_context=request)
session_message = SessionMessage(message, metadata=metadata)
session_message = self._create_session_message(message, request, request_id)
await writer.send(session_message)
except Exception:
logger.exception("SSE response error")
@ -778,6 +880,13 @@ class StreamableHTTPServerTransport:
# If stream ID not in mapping, create it
if stream_id and stream_id not in self._request_streams:
# Register SSE writer so close_sse_stream() can close it
self._sse_stream_writers[stream_id] = sse_stream_writer
# Send priming event for this new connection
await self._send_priming_event(stream_id, sse_stream_writer)
# Create new request streams for this connection
self._request_streams[stream_id] = anyio.create_memory_object_stream[EventMessage](0)
msg_reader = self._request_streams[stream_id][1]
@ -787,6 +896,9 @@ class StreamableHTTPServerTransport:
event_data = self._create_event_data(event_message)
await sse_stream_writer.send(event_data)
except anyio.ClosedResourceError:
# Expected when close_sse_stream() is called
logger.debug("Replay SSE stream closed by close_sse_stream()")
except Exception:
logger.exception("Error in replay sender")

View file

@ -51,6 +51,9 @@ class StreamableHTTPSessionManager:
json_response: Whether to use JSON responses instead of SSE streams
stateless: If True, creates a completely fresh transport for each request
with no session tracking or state persistence between requests.
security_settings: Optional transport security settings.
retry_interval: Retry interval in milliseconds to suggest to clients in SSE
retry field. Used for SSE polling behavior.
"""
def __init__(
@ -60,12 +63,14 @@ class StreamableHTTPSessionManager:
json_response: bool = False,
stateless: bool = False,
security_settings: TransportSecuritySettings | None = None,
retry_interval: int | None = None,
):
self.app = app
self.event_store = event_store
self.json_response = json_response
self.stateless = stateless
self.security_settings = security_settings
self.retry_interval = retry_interval
# Session tracking (only used if not stateless)
self._session_creation_lock = anyio.Lock()
@ -226,6 +231,7 @@ class StreamableHTTPSessionManager:
is_json_response_enabled=self.json_response,
event_store=self.event_store, # May be None (no resumability)
security_settings=self.security_settings,
retry_interval=self.retry_interval,
)
assert http_transport.mcp_session_id is not None