Implement close_sse_stream and client auto-reconnect (SEP-1699)

Server now supports closing SSE streams mid-operation via close_sse_stream(),
which triggers client reconnection. Client automatically reconnects when the
stream closes after receiving a priming event.

Changes:
- Server transport: Implement close_sse_stream() to close SSE writer
- Server transport: Create callback and pass via ServerMessageMetadata
- Lowlevel server: Thread close_sse_stream callback to RequestContext
- FastMCP Context: Wire close_sse_stream() to call the callback
- Client: Track priming events and auto-reconnect with Last-Event-ID

Author: felixweinberger

src/mcp/client/streamable_http.py

@@ -329,9 +329,15 @@ async def _handle_sse_response(
         is_initialization: bool = False,
     ) -> None:
         """Handle SSE response from the server."""
+        last_event_id: str | None = None
+
         try:
             event_source = EventSource(response)
             async for sse in event_source.aiter_sse():  # pragma: no branch
+                # Track last event ID for potential reconnection
+                if sse.id:
+                    last_event_id = sse.id
+
                 is_complete = await self._handle_sse_event(
                     sse,
                     ctx.read_stream_writer,
@@ -342,10 +348,63 @@ async def _handle_sse_response(
                 # break the loop
                 if is_complete:
                     await response.aclose()
-                    break
+                    return  # Normal completion, no reconnect needed
+        except Exception as e:
+            logger.debug(f"SSE stream ended: {e}")
+
+        # Stream ended without response - reconnect if we have priming event
+        if last_event_id is not None:
+            await self._handle_reconnection(ctx, last_event_id)
+
+    async def _handle_reconnection(
+        self,
+        ctx: RequestContext,
+        last_event_id: str,
+    ) -> None:
+        """Reconnect with Last-Event-ID to resume stream after server disconnect."""
+        headers = self._prepare_request_headers(ctx.headers)
+        headers[LAST_EVENT_ID] = last_event_id
+
+        # Extract original request ID to map responses
+        original_request_id = None
+        if isinstance(ctx.session_message.message.root, JSONRPCRequest):
+            original_request_id = ctx.session_message.message.root.id
+
+        try:
+            async with aconnect_sse(
+                ctx.client,
+                "GET",
+                self.url,
+                headers=headers,
+                timeout=httpx.Timeout(self.timeout, read=self.sse_read_timeout),
+            ) as event_source:
+                event_source.response.raise_for_status()
+                logger.debug("Reconnection GET SSE connection established")
+
+                # Track for potential further reconnection
+                reconnect_last_event_id: str | None = last_event_id
+
+                async for sse in event_source.aiter_sse():
+                    if sse.id:
+                        reconnect_last_event_id = sse.id
+
+                    is_complete = await self._handle_sse_event(
+                        sse,
+                        ctx.read_stream_writer,
+                        original_request_id,
+                        ctx.metadata.on_resumption_token_update if ctx.metadata else None,
+                    )
+                    if is_complete:
+                        await event_source.response.aclose()
+                        return
+
+                # Stream ended again without response - reconnect again
+                if reconnect_last_event_id is not None:
+                    await self._handle_reconnection(ctx, reconnect_last_event_id)
         except Exception as e:
-            logger.exception("Error reading SSE stream:")  # pragma: no cover
-            await ctx.read_stream_writer.send(e)  # pragma: no cover
+            logger.debug(f"Reconnection failed: {e}")
+            # Try to reconnect again if we still have an event ID
+            await self._handle_reconnection(ctx, last_event_id)
 
     async def _handle_unexpected_content_type(
         self,

src/mcp/server/fastmcp/server.py

@@ -1295,11 +1295,9 @@ async def close_sse_stream(self) -> None:
         Note:
             This is a no-op if not using StreamableHTTP transport with event_store.
             The callback is only available when event_store is configured.
-
-        Raises:
-            NotImplementedError: Feature not yet implemented.
         """
-        raise NotImplementedError("close_sse_stream not yet implemented")
+        if self._request_context and self._request_context.close_sse_stream:
+            await self._request_context.close_sse_stream()
 
     # Convenience methods for common log levels
     async def debug(self, message: str, **extra: Any) -> None:

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 from message metadata
                 request_data = None
+                close_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
 
                 # 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_cb,
                     )
                 )
                 response = await handler(req)

src/mcp/server/streamable_http.py

@@ -177,6 +177,7 @@ def __init__(
                 MemoryObjectReceiveStream[EventMessage],
             ],
         ] = {}
+        self._sse_stream_writers: dict[RequestId, MemoryObjectSendStream[dict[str, str]]] = {}
         self._terminated = False
 
     @property
@@ -202,7 +203,26 @@ def close_sse_stream(self, request_id: RequestId) -> None:
             Requires event_store to be configured for events to be stored during
             the disconnect.
         """
-        raise NotImplementedError("close_sse_stream not yet implemented")
+        writer = self._sse_stream_writers.pop(request_id, None)
+        if writer:
+            writer.close()
+
+    def _create_session_message(
+        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)
+
+        metadata = ServerMessageMetadata(
+            request_context=request,
+            close_sse_stream=close_stream_callback,
+        )
+        return SessionMessage(message, metadata=metadata)
 
     def _create_error_response(
         self,
@@ -485,6 +505,9 @@ async def _handle_post_request(self, scope: Scope, request: Request, receive: Re
                 # 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:
@@ -516,6 +539,7 @@ async def sse_writer():
                         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
@@ -539,8 +563,7 @@ 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)
-                        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")