Add SSE polling support (SEP-1699) (#1654)
This commit is contained in:
parent
2cd178a962
commit
281fd4765e
4 changed files with 1461 additions and 45 deletions
|
|
@ -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."""
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Reference in a new issue