Implement enhanced SSE support with keep-alive, retries, and proper encoding

Add keep-alive pings, configurable retry intervals, and a full
EventSourceResponse.encode() API to fastapi/sse.py with comprehensive
test coverage.

Co-Authored-By: Claude Opus 4.6 <[email protected]>

fastapi/sse.py

@@ -1,8 +1,13 @@
+from __future__ import annotations
+
+from collections.abc import AsyncIterator
 from typing import Annotated, Any
 
 from annotated_doc import Doc
 from pydantic import AfterValidator, BaseModel, Field, model_validator
+from starlette.background import BackgroundTask
 from starlette.responses import StreamingResponse
+from starlette.types import Receive, Send, Scope
 
 # Canonical SSE event schema matching the OpenAPI 3.2 spec
 # (Section 4.14.4 "Special Considerations for Server-Sent Events")
@@ -16,22 +21,270 @@ _SSE_EVENT_SCHEMA: dict[str, Any] = {
     },
 }
 
+# Sentinel object to represent an unset retry value.
+_UNSET = object()
 
-class EventSourceResponse(StreamingResponse):
-    """Streaming response with `text/event-stream` media type.
+# Seconds between keep-alive pings when a generator is idle.
+# Private but importable so tests can monkeypatch it.
+_PING_INTERVAL: float = 15.0
 
-    Use as `response_class=EventSourceResponse` on a *path operation* that uses `yield`
-    to enable Server Sent Events (SSE) responses.
+# Keep-alive comment, per the SSE spec recommendation
+KEEPALIVE_COMMENT = b": ping\n\n"
 
-    Works with **any HTTP method** (`GET`, `POST`, etc.), which makes it compatible
-    with protocols like MCP that stream SSE over `POST`.
 
-    The actual encoding logic lives in the FastAPI routing layer. This class
-    serves mainly as a marker and sets the correct `Content-Type`.
+class EventSourceResponse(StreamingResponse):
+    """Streaming response implementing the Server-Sent Events (SSE) protocol.
+
+    Use as ``response_class=EventSourceResponse`` on a *path operation* that uses ``yield``
+    to enable SSE responses.
+
+    Works with **any HTTP method** (``GET``, ``POST``, etc.), which makes it compatible
+    with protocols like MCP that stream SSE over ``POST``.
+
+    When used directly (not via the routing-layer generator shortcut), pass an async
+    iterator of :class:`ServerSentEvent` objects or raw bytes::
+
+        @app.get("/feed")
+        async def feed():
+            return EventSourceResponse(_event_stream())
+
+        async def _event_stream() -> AsyncIterator[ServerSentEvent]:
+            yield ServerSentEvent(data="hello", event="greeting")
+
+    Parameters
+    ----------
+    content:
+        An async or sync iterator yielding :class:`ServerSentEvent` objects or raw
+        bytes.  When yielding ``ServerSentEvent`` objects they are automatically
+        encoded into SSE wire format.
+    retry:
+        Default reconnection time in **milliseconds** sent to the client after every
+        event that does not override it.  Tells the browser how long to wait before
+        reconnecting if the connection is lost.
+    ping_interval:
+        Seconds between keep-alive comment pings when the generator is idle.
+        Set to ``0`` to disable.  Default is 15 seconds.
     """
 
     media_type = "text/event-stream"
 
+    def __init__(
+        self,
+        content: Any = None,
+        status_code: int = 200,
+        headers: dict[str, str] | None = None,
+        background: BackgroundTask | None = None,
+        retry: int | None | object = _UNSET,
+        ping_interval: float = _PING_INTERVAL,
+    ) -> None:
+        # Wrap the content iterator to auto-encode ServerSentEvent objects.
+        self._retry: int | None = (
+            retry if retry is not _UNSET and retry is not None else None
+        )
+        self._ping_interval: float = ping_interval
+
+        encoded_content = self._encode_content(content) if content is not None else ()
+
+        merged_headers: dict[str, str] = {
+            "Cache-Control": "no-cache",
+            "X-Accel-Buffering": "no",
+        }
+        if headers:
+            merged_headers.update(headers)
+
+        super().__init__(
+            content=encoded_content,
+            status_code=status_code,
+            headers=merged_headers,
+            media_type=self.media_type,
+            background=background,
+        )
+
+    # ------------------------------------------------------------------
+    # Encoding helpers
+    # ------------------------------------------------------------------
+
+    async def _encode_content(
+        self,
+        content: Any,
+    ) -> AsyncIterator[bytes]:
+        """Iterate *content* and yield SSE-encoded bytes.
+
+        Each item is expected to be a :class:`ServerSentEvent` (encoded
+        automatically) or already-raw bytes (passed through).
+        """
+        from fastapi.sse import ServerSentEvent, format_sse_event
+
+        if hasattr(content, "__aiter__"):
+            async for item in content:
+                yield self._encode_one(item)
+        else:
+            for item in content:
+                yield self._encode_one(item)
+
+    def _encode_one(self, item: Any) -> bytes:
+        """Encode a single item into SSE wire format bytes."""
+        from fastapi.sse import ServerSentEvent, format_sse_event
+
+        if isinstance(item, ServerSentEvent):
+            return self._encode_sse(item)
+        elif isinstance(item, bytes):
+            return item
+        elif isinstance(item, str):
+            return item.encode("utf-8")
+        else:
+            # Fallback: treat as a dict-like object — JSON-encode as data.
+            data_str = self._serialize_data(item)
+            return format_sse_event(
+                data_str=data_str,
+                retry=self._retry,
+            )
+
+    def _encode_sse(self, event: ServerSentEvent) -> bytes:
+        """Encode a :class:`ServerSentEvent` into SSE wire-format bytes."""
+        from fastapi.sse import format_sse_event
+
+        if event.raw_data is not None:
+            data_str: str | None = event.raw_data
+        elif event.data is not None:
+            if hasattr(event.data, "model_dump_json"):
+                data_str = event.data.model_dump_json()
+            else:
+                data_str = self._serialize_data(event.data)
+        else:
+            data_str = None
+
+        # Use the event's own retry if set, otherwise fall back to the
+        # response-level default.
+        retry: int | None = event.retry if event.retry is not None else self._retry
+
+        return format_sse_event(
+            data_str=data_str,
+            event=event.event,
+            id=event.id,
+            retry=retry,
+            comment=event.comment,
+        )
+
+    @staticmethod
+    def _serialize_data(data: Any) -> str:
+        """Serialize an arbitrary data payload to a JSON string."""
+        import json
+
+        from fastapi.encoders import jsonable_encoder
+
+        data = jsonable_encoder(data)
+        return json.dumps(data)
+
+    # ------------------------------------------------------------------
+    # Disconnect detection
+    # ------------------------------------------------------------------
+
+    async def disconnect(self) -> bool:
+        """Wait for the client to disconnect and return ``True``.
+
+        Use this inside an endpoint to wait until the client closes the
+        connection, allowing clean shutdown of background tasks::
+
+            @app.get("/stream")
+            async def stream():
+                response = EventSourceResponse(_events())
+                await response.disconnect()  # blocks until client leaves
+                await cleanup()
+                return response
+        """
+        await self.listen_for_disconnect(self._receive)
+        return True
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        """ASGI entry-point.  Stores *receive* so :meth:`disconnect` works."""
+        self._receive = receive
+        await super().__call__(scope, receive, send)
+
+    # ------------------------------------------------------------------
+    # Static convenience methods (automatic event formatting)
+    # ------------------------------------------------------------------
+
+    @staticmethod
+    def encode(
+        data: Any | None = None,
+        *,
+        event: str | None = None,
+        id: str | None = None,
+        retry: int | None = None,
+        comment: str | None = None,
+    ) -> bytes:
+        """Encode an event into SSE wire-format bytes.
+
+        This is a convenience wrapper around :func:`format_sse_event` that
+        automatically JSON-serializes the ``data`` argument.
+
+        Parameters
+        ----------
+        data:
+            Any JSON-serializable value.  Strings are quoted on the wire.
+        event:
+            Optional event type name.
+        id:
+            Optional event ID.
+        retry:
+            Optional reconnection time in milliseconds.
+        comment:
+            Optional comment text (ignored by browsers).
+        """
+        from fastapi.sse import format_sse_event
+
+        if data is not None:
+            if hasattr(data, "model_dump_json"):
+                data_str = data.model_dump_json()
+            else:
+                data_str = EventSourceResponse._serialize_data(data)
+        else:
+            data_str = None
+
+        return format_sse_event(
+            data_str=data_str,
+            event=event,
+            id=id,
+            retry=retry,
+            comment=comment,
+        )
+
+    @staticmethod
+    def encode_raw(
+        raw_data: str,
+        *,
+        event: str | None = None,
+        id: str | None = None,
+        retry: int | None = None,
+        comment: str | None = None,
+    ) -> bytes:
+        """Encode raw (non-JSON) data into SSE wire-format bytes.
+
+        The ``raw_data`` string is placed into the ``data:`` field as-is
+        without JSON encoding.
+        """
+        from fastapi.sse import format_sse_event
+
+        return format_sse_event(
+            data_str=raw_data,
+            event=event,
+            id=id,
+            retry=retry,
+            comment=comment,
+        )
+
+    @staticmethod
+    def encode_comment(text: str) -> bytes:
+        """Encode an SSE comment.
+
+        Comments start with ``:`` on the wire and are ignored by
+        ``EventSource`` clients.  Useful for keep-alive pings.
+        """
+        from fastapi.sse import format_sse_event
+
+        return format_sse_event(comment=text)
+
 
 def _check_id_no_null(v: str | None) -> str | None:
     if v is not None and "\0" in v:
@@ -212,11 +465,3 @@ def format_sse_event(
     lines.append("")
     lines.append("")
     return "\n".join(lines).encode("utf-8")
-
-
-# Keep-alive comment, per the SSE spec recommendation
-KEEPALIVE_COMMENT = b": ping\n\n"
-
-# Seconds between keep-alive pings when a generator is idle.
-# Private but importable so tests can monkeypatch it.
-_PING_INTERVAL: float = 15.0

tests/test_sse.py

@@ -316,3 +316,296 @@ def test_no_keepalive_when_fast(client: TestClient):
     assert response.status_code == 200
     # KEEPALIVE_COMMENT is ": ping\n\n".
     assert ": ping\n" not in response.text
+
+
+# ── Tests for enhanced EventSourceResponse class ──────────────────────
+
+
+def test_encode_static_method():
+    """EventSourceResponse.encode() JSON-serializes data automatically."""
+    result = EventSourceResponse.encode(data="hello", event="greeting", id="1")
+    assert b'data: "hello"\n' in result
+    assert b"event: greeting\n" in result
+    assert b"id: 1\n" in result
+
+
+def test_encode_with_dict():
+    """encode() handles dict data."""
+    result = EventSourceResponse.encode(data={"key": "value"}, event="json")
+    assert b'data: {"key": "value"}\n' in result or b'data: {"key":"value"}\n' in result
+
+
+def test_encode_with_model():
+    """encode() handles Pydantic model data."""
+    result = EventSourceResponse.encode(data=Item(name="Test"), event="model")
+    assert b'"name":"Test"' in result or b'"name": "Test"' in result
+
+
+def test_encode_with_retry():
+    """encode() includes retry field."""
+    result = EventSourceResponse.encode(data="test", retry=5000)
+    assert b"retry: 5000\n" in result
+
+
+def test_encode_comment():
+    """encode_comment() creates a comment event."""
+    result = EventSourceResponse.encode_comment("keepalive")
+    assert b": keepalive\n" in result
+
+
+def test_encode_raw():
+    """encode_raw() places raw data without JSON encoding."""
+    result = EventSourceResponse.encode_raw(
+        raw_data="<div>hello</div>", event="html"
+    )
+    assert b"data: <div>hello</div>\n" in result
+    assert b'data: "<div>hello</div>"' not in result
+
+
+def test_encode_no_data():
+    """encode() with no data still produces valid SSE (event-only)."""
+    result = EventSourceResponse.encode(event="heartbeat")
+    assert b"event: heartbeat\n" in result
+    assert b"data:" not in result
+
+
+def test_direct_instantiation_with_sse_events():
+    """EventSourceResponse can be used directly with an async iterator of SSE events."""
+    app = FastAPI()
+
+    async def event_stream():
+        yield ServerSentEvent(data="first", event="start")
+        yield ServerSentEvent(data="second", event="middle")
+        yield ServerSentEvent(data="last", event="end")
+
+    @app.get("/direct")
+    async def direct_endpoint():
+        return EventSourceResponse(event_stream())
+
+    with TestClient(app) as c:
+        response = c.get("/direct")
+        assert response.status_code == 200
+        assert response.headers["content-type"] == "text/event-stream; charset=utf-8"
+        assert response.headers["cache-control"] == "no-cache"
+        assert response.headers["x-accel-buffering"] == "no"
+        text = response.text
+        assert "event: start\n" in text
+        assert 'data: "first"\n' in text
+        assert "event: middle\n" in text
+        assert "event: end\n" in text
+
+
+def test_direct_instantiation_with_dict_items():
+    """EventSourceResponse auto-JSON-encodes plain dict items."""
+    app = FastAPI()
+
+    async def event_stream():
+        yield {"msg": "hello"}
+        yield {"msg": "world"}
+
+    @app.get("/direct-dict")
+    async def direct_endpoint():
+        return EventSourceResponse(event_stream())
+
+    with TestClient(app) as c:
+        response = c.get("/direct-dict")
+        assert response.status_code == 200
+        text = response.text
+        assert '"msg": "hello"' in text or '"msg":"hello"' in text
+        assert '"msg": "world"' in text or '"msg":"world"' in text
+
+
+def test_direct_instantiation_with_raw_bytes():
+    """EventSourceResponse passes through raw bytes unchanged."""
+    app = FastAPI()
+
+    async def event_stream():
+        yield b"data: raw bytes\n\n"
+        yield b"data: more bytes\n\n"
+
+    @app.get("/direct-bytes")
+    async def direct_endpoint():
+        return EventSourceResponse(event_stream())
+
+    with TestClient(app) as c:
+        response = c.get("/direct-bytes")
+        assert response.status_code == 200
+        assert response.text == "data: raw bytes\n\ndata: more bytes\n\n"
+
+
+def test_direct_instantiation_with_retry():
+    """EventSourceResponse retry parameter is applied to events."""
+    app = FastAPI()
+
+    async def event_stream():
+        yield ServerSentEvent(data="msg1")
+        yield ServerSentEvent(data="msg2")
+
+    @app.get("/direct-retry")
+    async def direct_endpoint():
+        return EventSourceResponse(event_stream(), retry=3000)
+
+    with TestClient(app) as c:
+        response = c.get("/direct-retry")
+        assert response.status_code == 200
+        text = response.text
+        # Each event should have the default retry
+        assert text.count("retry: 3000\n") == 2
+
+
+def test_direct_instantiation_retry_override():
+    """Event-level retry overrides response-level retry."""
+    app = FastAPI()
+
+    async def event_stream():
+        yield ServerSentEvent(data="default-retry")
+        yield ServerSentEvent(data="override-retry", retry=7000)
+
+    @app.get("/direct-retry-override")
+    async def direct_endpoint():
+        return EventSourceResponse(event_stream(), retry=3000)
+
+    with TestClient(app) as c:
+        response = c.get("/direct-retry-override")
+        assert response.status_code == 200
+        text = response.text
+        assert "retry: 3000\n" in text
+        assert "retry: 7000\n" in text
+
+
+def test_direct_instantiation_with_models():
+    """EventSourceResponse auto-encodes Pydantic models."""
+    app = FastAPI()
+
+    async def event_stream():
+        yield Item(name="Widget", description="A widget")
+        yield Item(name="Gadget")
+
+    @app.get("/direct-models")
+    async def direct_endpoint():
+        return EventSourceResponse(event_stream())
+
+    with TestClient(app) as c:
+        response = c.get("/direct-models")
+        assert response.status_code == 200
+        text = response.text
+        assert '"name": "Widget"' in text or '"name":"Widget"' in text
+        assert '"name": "Gadget"' in text or '"name":"Gadget"' in text
+
+
+def test_direct_instantiation_mixed_content():
+    """EventSourceResponse handles mixed ServerSentEvent and raw bytes."""
+    app = FastAPI()
+
+    async def event_stream():
+        yield ServerSentEvent(data="hello", event="greet")
+        yield b"data: raw event\n\n"
+        yield ServerSentEvent(data="world", event="farewell")
+
+    @app.get("/direct-mixed")
+    async def direct_endpoint():
+        return EventSourceResponse(event_stream())
+
+    with TestClient(app) as c:
+        response = c.get("/direct-mixed")
+        assert response.status_code == 200
+        text = response.text
+        assert "event: greet\n" in text
+        assert "data: raw event\n" in text
+        assert "event: farewell\n" in text
+
+
+def test_direct_instantiation_custom_headers():
+    """EventSourceResponse merges custom headers with defaults."""
+    app = FastAPI()
+
+    async def event_stream():
+        yield ServerSentEvent(data="test")
+
+    @app.get("/direct-headers")
+    async def direct_endpoint():
+        return EventSourceResponse(
+            event_stream(),
+            headers={"X-Custom-Header": "custom-value"},
+        )
+
+    with TestClient(app) as c:
+        response = c.get("/direct-headers")
+        assert response.status_code == 200
+        assert response.headers["cache-control"] == "no-cache"
+        assert response.headers["x-accel-buffering"] == "no"
+        assert response.headers["x-custom-header"] == "custom-value"
+
+
+def test_direct_instantiation_custom_headers_override():
+    """Custom headers can override default headers."""
+    app = FastAPI()
+
+    async def event_stream():
+        yield ServerSentEvent(data="test")
+
+    @app.get("/direct-headers-override")
+    async def direct_endpoint():
+        return EventSourceResponse(
+            event_stream(),
+            headers={"Cache-Control": "custom-cache"},
+        )
+
+    with TestClient(app) as c:
+        response = c.get("/direct-headers-override")
+        assert response.status_code == 200
+        assert response.headers["cache-control"] == "custom-cache"
+
+
+def test_direct_instantiation_sync_iterator():
+    """EventSourceResponse works with sync iterators too."""
+    app = FastAPI()
+
+    def event_stream():
+        yield ServerSentEvent(data="one")
+        yield ServerSentEvent(data="two")
+
+    @app.get("/direct-sync")
+    def direct_endpoint():
+        return EventSourceResponse(event_stream())
+
+    with TestClient(app) as c:
+        response = c.get("/direct-sync")
+        assert response.status_code == 200
+        text = response.text
+        assert 'data: "one"\n' in text
+        assert 'data: "two"\n' in text
+
+
+def test_direct_instantiation_with_comment():
+    """EventSourceResponse encodes comment events."""
+    app = FastAPI()
+
+    async def event_stream():
+        yield ServerSentEvent(comment="keep-alive")
+        yield ServerSentEvent(data="data-event")
+
+    @app.get("/direct-comment")
+    async def direct_endpoint():
+        return EventSourceResponse(event_stream())
+
+    with TestClient(app) as c:
+        response = c.get("/direct-comment")
+        assert response.status_code == 200
+        text = response.text
+        assert ": keep-alive\n" in text
+        assert 'data: "data-event"\n' in text
+
+
+def test_encode_json_encodes_multiline_data():
+    """encode() JSON-encodes data, so newlines are escaped inside the JSON string."""
+    result = EventSourceResponse.encode(data="line1\nline2\nline3")
+    # Since data is JSON-encoded, newlines appear as \n inside the JSON string.
+    assert b'"line1\\nline2\\nline3"' in result
+
+
+def test_encode_raw_preserves_multiline():
+    """encode_raw() splits multiline raw data across multiple data: lines."""
+    result = EventSourceResponse.encode_raw(raw_data="row1\ncsv,row2")
+    assert b"data: row1\n" in result
+    assert b"data: csv,row2\n" in result