Add close_sse_stream callback to request context for ergonomic stream closing

felixweinberger · felixweinberger · commit 426773f5178d · 2025-11-26T22:36:24.000Z
This change refactors the SSE stream closing API to follow the TypeScript SDK pattern (PR #1166), making it easier for tool handlers to trigger client reconnection during long-running operations. Changes: - Add CloseSSEStreamCallback type alias to message.py - Add close_sse_stream field to ServerMessageMetadata and RequestContext - Create callback closure in StreamableHTTPServerTransport._handle_post_request - Thread callback through lowlevel/server.py to RequestContext - Expose ctx.close_sse_stream() method in FastMCP Context class - Update tests to use new callback API instead of session_manager_ref hack - Fix dead code: remove unused return values from client _handle_sse_response Usage in FastMCP tools: @mcp.tool() async def long_running_tool(ctx: Context) -> str: await ctx.close_sse_stream() # Trigger client reconnection # Continue processing... return "Done"
diff --git a/src/mcp/client/streamable_http.py b/src/mcp/client/streamable_http.py
@@ -380,14 +380,8 @@ async def _handle_sse_response(
         ctx: RequestContext,
         is_initialization: bool = False,
         attempt: int = 0,
-    ) -> tuple[bool, str | None]:
-        """Handle SSE response from the server with automatic reconnection.
-
-        Returns:
-            Tuple of (has_priming_event, last_event_id) where:
-            - has_priming_event: True if any event had an ID (priming event received)
-            - last_event_id: The last event ID received, for resumption
-        """
+    ) -> None:
+        """Handle SSE response from the server with automatic reconnection."""
         has_priming_event = False
         last_event_id: str | None = None
         is_complete = False
@@ -422,8 +416,6 @@ async def _handle_sse_response(
         if not is_complete and has_priming_event and last_event_id:  # pragma: no cover
             await self._attempt_sse_reconnection(ctx, last_event_id, attempt)
 
-        return has_priming_event, last_event_id
-
     async def _attempt_sse_reconnection(  # pragma: no cover
         self,
         ctx: RequestContext,
diff --git a/src/mcp/server/fastmcp/server.py b/src/mcp/server/fastmcp/server.py
@@ -1303,3 +1303,24 @@ async def warning(self, message: str, **extra: Any) -> None:
     async def error(self, message: str, **extra: Any) -> None:
         """Send an error log message."""
         await self.log("error", message, **extra)
+
+    async def close_sse_stream(self, retry_interval_ms: int | None = None) -> None:
+        """Close the SSE stream for this request, triggering client reconnection.
+
+        Use this to implement polling behavior during long-running operations.
+        The client will reconnect after the retry interval and receive any events
+        that were stored while disconnected.
+
+        This is only available when using StreamableHTTP transport with an
+        event store configured for resumability. It is a no-op otherwise,
+        allowing portable tool code.
+
+        Args:
+            retry_interval_ms: Optional retry interval in milliseconds to suggest
+                              to the client before closing (currently unused,
+                              reserved for future use - the retry interval is
+                              configured at the transport level)
+        """
+        callback = self._request_context.close_sse_stream if self._request_context else None
+        if callback:
+            await callback(retry_interval_ms)
diff --git a/src/mcp/server/lowlevel/server.py b/src/mcp/server/lowlevel/server.py
@@ -680,12 +680,14 @@ async def _handle_request(
 
             token = None
             try:
-                # Extract request context from message metadata
+                # Extract request context and close_sse_stream callback from metadata
                 request_data = None
+                close_sse_stream_callback = 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_callback = message.message_metadata.close_sse_stream
 
                 # Set our global state that can be retrieved via
                 # app.get_request_context()
@@ -696,6 +698,7 @@ async def _handle_request(
                         session,
                         lifespan_context,
                         request=request_data,
+                        close_sse_stream=close_sse_stream_callback,
                     )
                 )
                 response = await handler(req)
diff --git a/src/mcp/server/streamable_http.py b/src/mcp/server/streamable_http.py
@@ -28,7 +28,7 @@
     TransportSecurityMiddleware,
     TransportSecuritySettings,
 )
-from mcp.shared.message import ServerMessageMetadata, SessionMessage
+from mcp.shared.message import CloseSSEStreamCallback, ServerMessageMetadata, SessionMessage
 from mcp.shared.version import SUPPORTED_PROTOCOL_VERSIONS
 from mcp.types import (
     DEFAULT_NEGOTIATED_VERSION,
@@ -449,9 +449,15 @@ async def _handle_post_request(self, scope: Scope, request: Request, receive: Re
             self._request_streams[request_id] = anyio.create_memory_object_stream[EventMessage](0)  # pragma: no cover
             request_stream_reader = self._request_streams[request_id][1]  # pragma: no cover
 
+            # Create close_sse_stream callback (only if event store is configured)
+            close_sse_callback = self._create_close_sse_callback(request_id)  # pragma: no cover
+
             if self.is_json_response_enabled:  # pragma: no cover
                 # Process the message
-                metadata = ServerMessageMetadata(request_context=request)
+                metadata = ServerMessageMetadata(
+                    request_context=request,
+                    close_sse_stream=close_sse_callback,
+                )
                 session_message = SessionMessage(message, metadata=metadata)
                 await writer.send(session_message)
                 try:
@@ -544,7 +550,10 @@ async def sse_writer():
                     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)
+                        metadata = ServerMessageMetadata(
+                            request_context=request,
+                            close_sse_stream=close_sse_callback,
+                        )
                         session_message = SessionMessage(message, metadata=metadata)
                         await writer.send(session_message)
                 except Exception:
@@ -716,6 +725,28 @@ async def terminate(self) -> None:
             # During cleanup, we catch all exceptions since streams might be in various states
             logger.debug(f"Error closing streams: {e}")
 
+    def _create_close_sse_callback(self, request_id: str) -> CloseSSEStreamCallback | None:
+        """Create a callback to close the SSE stream for a request.
+
+        Only creates a callback if event store is configured (resumability enabled).
+        The callback allows handlers to trigger client reconnection during long-running ops.
+
+        Args:
+            request_id: The request ID to create the callback for
+
+        Returns:
+            The callback function, or None if event store is not configured
+        """
+        if not self._event_store:
+            return None
+
+        async def close_callback(retry_interval_ms: int | None = None) -> None:
+            # Note: retry_interval_ms is accepted for API compatibility but not used
+            # The retry interval is set via the priming event
+            await self.close_sse_stream(request_id)
+
+        return close_callback
+
     async def close_sse_stream(self, request_id: RequestId) -> None:
         """Close an SSE stream for a specific request, triggering client reconnection.
 
diff --git a/src/mcp/shared/context.py b/src/mcp/shared/context.py
@@ -3,6 +3,7 @@
 
 from typing_extensions import TypeVar
 
+from mcp.shared.message import CloseSSEStreamCallback
 from mcp.shared.session import BaseSession
 from mcp.types import RequestId, RequestParams
 
@@ -18,3 +19,4 @@ class RequestContext(Generic[SessionT, LifespanContextT, RequestT]):
     session: SessionT
     lifespan_context: LifespanContextT
     request: RequestT | None = None
+    close_sse_stream: CloseSSEStreamCallback | None = None
diff --git a/src/mcp/shared/message.py b/src/mcp/shared/message.py
@@ -14,6 +14,9 @@
 
 ResumptionTokenUpdateCallback = Callable[[ResumptionToken], Awaitable[None]]
 
+# Callback type for closing SSE streams (takes optional retry_interval_ms)
+CloseSSEStreamCallback = Callable[[int | None], Awaitable[None]]
+
 
 @dataclass
 class ClientMessageMetadata:
@@ -30,6 +33,8 @@ class ServerMessageMetadata:
     related_request_id: RequestId | None = None
     # Request-specific context (e.g., headers, auth info)
     request_context: object | None = None
+    # Callback to close SSE stream for this request (triggers client reconnection)
+    close_sse_stream: CloseSSEStreamCallback | None = None
 
 
 MessageMetadata = ClientMessageMetadata | ServerMessageMetadata | None
diff --git a/tests/shared/test_streamable_http.py b/tests/shared/test_streamable_http.py
@@ -269,14 +269,9 @@ async def handle_call_tool(name: str, args: dict[str, Any]) -> list[TextContent]
                     related_request_id=ctx.request_id,
                 )
 
-                # Trigger server-initiated SSE disconnect
-                if self._session_manager_ref:
-                    session_manager = self._session_manager_ref[0]
-                    request = ctx.request
-                    if isinstance(request, Request):
-                        session_id = request.headers.get("mcp-session-id")
-                        if session_id:
-                            await session_manager.close_sse_stream(session_id, ctx.request_id)
+                # Trigger server-initiated SSE disconnect using the callback
+                if ctx.close_sse_stream:
+                    await ctx.close_sse_stream(None)
 
                 # Wait a bit for client to reconnect
                 await anyio.sleep(0.2)
