Add dual-mode direct-use to EventSourceResponse

EventSourceResponse can now be used directly as a return value (like
StreamingResponse) in addition to the existing response_class mode.

When used directly, it auto-formats ServerSentEvent objects into SSE
wire format, emits an initial retry field, and inserts keep-alive
comment pings on idle periods. Also supports disconnect callbacks
and sync iterables.

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

fastapi/sse.py

@@ -1,8 +1,27 @@
+from __future__ import annotations
+
+import json
+from collections.abc import AsyncIterable, Callable, Iterable
+from functools import partial
 from typing import Annotated, Any
 
+import anyio
 from annotated_doc import Doc
 from pydantic import AfterValidator, BaseModel, Field, model_validator
+from starlette.concurrency import iterate_in_threadpool
+from starlette.requests import ClientDisconnect
 from starlette.responses import StreamingResponse
+from starlette.types import Message, Receive, Scope, Send
+
+try:
+    from starlette._utils import collapse_excgroups
+except ImportError:
+    # Fallback for older starlette versions
+    import contextlib
+
+    @contextlib.contextmanager
+    def collapse_excgroups():
+        yield
 
 # Canonical SSE event schema matching the OpenAPI 3.2 spec
 # (Section 4.14.4 "Special Considerations for Server-Sent Events")
@@ -18,20 +37,256 @@ _SSE_EVENT_SCHEMA: dict[str, Any] = {
 
 
 class EventSourceResponse(StreamingResponse):
-    """Streaming response with `text/event-stream` media type.
+    """Streaming response with ``text/event-stream`` media type.
-
+
-    Use as `response_class=EventSourceResponse` on a *path operation* that uses `yield`
+    Use as ``response_class=EventSourceResponse`` on a *path operation* that
-    to enable Server Sent Events (SSE) responses.
+    uses ``yield`` to enable Server-Sent Events (SSE) responses via the
-
+    routing layer's automatic encoding.
-    Works with **any HTTP method** (`GET`, `POST`, etc.), which makes it compatible
+
-    with protocols like MCP that stream SSE over `POST`.
+    Can also be used **directly as a return value**, similar to
-
+    ``StreamingResponse``, by passing an async iterable of pre-formatted
-    The actual encoding logic lives in the FastAPI routing layer. This class
+    SSE bytes or ``ServerSentEvent`` objects::
-    serves mainly as a marker and sets the correct `Content-Type`.
+
+        @app.get("/stream")
+        def stream_events():
+            return EventSourceResponse(my_event_generator(), retry=3000)
+
+    When used directly the class formats ``ServerSentEvent`` objects into
+    SSE wire format, emits an optional initial ``retry:`` field, and
+    inserts keep-alive comment pings on idle periods.
+
+    **Parameters**
+
+    * ``content`` – async iterable of ``ServerSentEvent`` objects or raw
+      ``bytes``.  Required for direct-use mode; omit when used as
+      ``response_class`` (the routing layer provides content).
+    * ``retry`` – default reconnection time in **milliseconds**. Emitted
+      as an initial ``retry:`` field.  Individual ``ServerSentEvent``
+      objects can override per-event.
+    * ``ping_interval`` – seconds between keep-alive ``: ping`` comments
+      in direct-use mode.  Set to ``0`` to disable.  Default ``15.0``.
+    * ``disconnect_callback`` – async callable invoked when the client
+      disconnects.  Useful for resource cleanup.
     """
 
     media_type = "text/event-stream"
 
+    def __init__(
+        self,
+        content: AsyncIterable[Any] | None = None,
+        status_code: int = 200,
+        headers: dict[str, str] | None = None,
+        media_type: str | None = None,
+        background: Any = None,
+        retry: int | None = None,
+        ping_interval: float = 15.0,
+        disconnect_callback: (
+            Callable[[], Any] | None  # sync or async callable
+        ) = None,
+    ) -> None:
+        if retry is not None and retry < 0:
+            raise ValueError("retry must be non-negative")
+
+        # When content is provided (direct-use mode), wrap it to auto-format
+        # ServerSentEvent objects, emit an initial retry field, and add
+        # keepalive pings.  Also set SSE-appropriate headers.
+        if content is not None:
+            content = _wrap_content(
+                content, retry=retry, ping_interval=ping_interval
+            )
+            if headers is None:
+                headers = {}
+            headers.setdefault("Cache-Control", "no-cache")
+            # Prevent Nginx and other proxies from buffering SSE
+            headers.setdefault("X-Accel-Buffering", "no")
+
+            # Store callback and content flag for __call__
+            self._disconnect_callback = disconnect_callback
+            self._direct_use = True
+        else:
+            self._disconnect_callback = None
+            self._direct_use = False
+
+        super().__init__(
+            content=content,
+            status_code=status_code,
+            headers=headers,
+            media_type=media_type,
+            background=background,
+        )
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        if scope["type"] == "websocket":
+            send = self._wrap_websocket_denial_send(send)
+            await self.stream_response(send)
+            if self.background is not None:
+                await self.background()
+            return
+
+        if not self._direct_use or self._disconnect_callback is None:
+            # Routing-layer mode or no callback: delegate to parent
+            await super().__call__(scope, receive, send)
+            return
+
+        # Direct-use mode with disconnect callback
+        spec_version = tuple(
+            map(int, scope.get("asgi", {}).get("spec_version", "2.0").split("."))
+        )
+
+        if spec_version >= (2, 4):
+            try:
+                await self.stream_response(send)
+            except OSError:
+                cb = self._disconnect_callback
+                if cb is not None:
+                    result = cb()
+                    if result is not None and hasattr(result, "__await__"):
+                        await result
+                raise ClientDisconnect()
+        else:
+            with collapse_excgroups():
+                async with anyio.create_task_group() as task_group:
+
+                    async def wrap(
+                        func: Callable[[], Any],
+                    ) -> None:
+                        await func()
+                        task_group.cancel_scope.cancel()
+
+                    task_group.start_soon(wrap, partial(self.stream_response, send))
+                    await wrap(
+                        partial(
+                            _listen_for_disconnect,
+                            receive,
+                            task_group,
+                            self._disconnect_callback,
+                        )
+                    )
+
+        if self.background is not None:
+            await self.background()
+
+
+async def _listen_for_disconnect(
+    receive: Receive,
+    task_group: anyio.abc.TaskGroup,
+    callback: Callable[[], Any] | None,
+) -> None:
+    """Listen for disconnect, invoke callback, then cancel the task group."""
+    while True:
+        message = await receive()
+        if message["type"] == "http.disconnect":
+            break
+    if callback is not None:
+        result = callback()
+        if result is not None and hasattr(result, "__await__"):
+            await result
+    task_group.cancel_scope.cancel()
+
+
+def _wrap_content(
+    content: AsyncIterable[Any] | Iterable[Any],
+    *,
+    retry: int | None,
+    ping_interval: float,
+) -> AsyncIterable[bytes]:
+    """Wrap content to auto-format ServerSentEvent objects and emit an
+    initial ``retry:`` field if configured.
+
+    When ``ping_interval > 0``, a keep-alive comment is inserted whenever
+    the source generator is idle for longer than the interval.
+    """
+    if isinstance(content, AsyncIterable):
+        ait: AsyncIterable[Any] = content
+    else:
+        ait = iterate_in_threadpool(content)
+
+    if ping_interval <= 0:
+        # No keepalive needed: simple generator
+        async def _simple_generator() -> AsyncIterable[bytes]:
+            if retry is not None:
+                yield format_sse_event(retry=retry)
+            async for event in ait:
+                yield _format_one(event)
+                await anyio.sleep(0)
+
+        return _simple_generator()
+
+    # Keepalive-enabled: use anyio memory stream to decouple
+    # the producer from the timeout logic.
+    send_stream, receive_stream = anyio.create_memory_object_stream[bytes](
+        max_buffer_size=1
+    )
+
+    async def _format_and_send() -> None:
+        async with send_stream:
+            if retry is not None:
+                await send_stream.send(format_sse_event(retry=retry))
+            async for event in ait:
+                await send_stream.send(_format_one(event))
+                await anyio.sleep(0)
+
+    async def _insert_pings() -> AsyncIterable[bytes]:
+        async with receive_stream:
+            try:
+                while True:
+                    try:
+                        with anyio.fail_after(ping_interval):
+                            data = await receive_stream.receive()
+                        yield data
+                    except TimeoutError:
+                        yield KEEPALIVE_COMMENT
+            except anyio.EndOfStream:
+                pass
+
+    # We need to return an async iterable that starts the producer task
+    # when iterated. Use a wrapper async generator.
+    async def _keepalive_generator() -> AsyncIterable[bytes]:
+        async with anyio.create_task_group() as tg:
+            tg.start_soon(_format_and_send)
+            async for chunk in _insert_pings():
+                yield chunk
+                await anyio.sleep(0)
+            tg.cancel_scope.cancel()
+
+    return _keepalive_generator()
+
+
+def _format_one(event: Any) -> bytes:
+    """Format a single event into SSE wire-format bytes."""
+    if isinstance(event, ServerSentEvent):
+        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 = json.dumps(jsonable_encoder(event.data))
+        else:
+            data_str = None
+        return format_sse_event(
+            data_str=data_str,
+            event=event.event,
+            id=event.id,
+            retry=event.retry,
+            comment=event.comment,
+        )
+    else:
+        # Plain bytes or string (already formatted by user)
+        if isinstance(event, bytes):
+            return event
+        if isinstance(event, str):
+            return event.encode("utf-8")
+        # Unexpected type — try to convert
+        return str(event).encode("utf-8")
+
+
+def jsonable_encoder(obj: Any) -> Any:
+    """Minimal jsonable_encoder fallback for standalone use."""
+    from fastapi.encoders import jsonable_encoder as _enc
+
+    return _enc(obj)
+
 
 def _check_id_no_null(v: str | None) -> str | None:
     if v is not None and "\0" in v:
@@ -42,18 +297,19 @@ def _check_id_no_null(v: str | None) -> str | None:
 class ServerSentEvent(BaseModel):
     """Represents a single Server-Sent Event.
 
-    When `yield`ed from a *path operation function* that uses
+    When ``yield``ed from a *path operation function* that uses
-    `response_class=EventSourceResponse`, each `ServerSentEvent` is encoded
+    ``response_class=EventSourceResponse``, each ``ServerSentEvent`` is
-    into the [SSE wire format](https://html.spec.whatwg.org/multipage/server-sent-events.html#parsing-an-event-stream)
+    encoded into the `SSE wire format
-    (`text/event-stream`).
+    <https://html.spec.whatwg.org/multipage/server-sent-events.html#parsing-an-event-stream>`_
+    (``text/event-stream``).
 
     If you yield a plain object (dict, Pydantic model, etc.) instead, it is
-    automatically JSON-encoded and sent as the `data:` field.
+    automatically JSON-encoded and sent as the ``data:`` field.
 
-    All `data` values **including plain strings** are JSON-serialized.
+    All ``data`` values **including plain strings** are JSON-serialized.
 
-    For example, `data="hello"` produces `data: "hello"` on the wire (with
+    For example, ``data="hello"`` produces ``data: "hello"`` on the wire
-    quotes).
+    (with quotes).
     """
 
     data: Annotated[
@@ -64,9 +320,9 @@ class ServerSentEvent(BaseModel):
 
             Can be any JSON-serializable value: a Pydantic model, dict, list,
             string, number, etc. It is **always** serialized to JSON: strings
-            are quoted (`"hello"` becomes `data: "hello"` on the wire).
+            are quoted (``"hello"`` becomes ``data: "hello"`` on the wire).
 
-            Mutually exclusive with `raw_data`.
+            Mutually exclusive with ``raw_data``.
             """
         ),
     ] = None
@@ -74,13 +330,14 @@ class ServerSentEvent(BaseModel):
         str | None,
         Doc(
             """
-            Raw string to send as the `data:` field **without** JSON encoding.
+            Raw string to send as the ``data:`` field **without** JSON
+            encoding.
 
-            Use this when you need to send pre-formatted text, HTML fragments,
+            Use this when you need to send pre-formatted text, HTML
-            CSV lines, or any non-JSON payload. The string is placed directly
+            fragments, CSV lines, or any non-JSON payload. The string is
-            into the `data:` field as-is.
+            placed directly into the ``data:`` field as-is.
 
-            Mutually exclusive with `data`.
+            Mutually exclusive with ``data``.
             """
         ),
     ] = None
@@ -90,8 +347,9 @@ class ServerSentEvent(BaseModel):
             """
             Optional event type name.
 
-            Maps to `addEventListener(event, ...)` on the browser. When omitted,
+            Maps to ``addEventListener(event, ...)`` on the browser. When
-            the browser dispatches on the generic `message` event.
+            omitted, the browser dispatches on the generic ``message``
+            event.
             """
         ),
     ] = None
@@ -102,8 +360,9 @@ class ServerSentEvent(BaseModel):
             """
             Optional event ID.
 
-            The browser sends this value back as the `Last-Event-ID` header on
+            The browser sends this value back as the ``Last-Event-ID``
-            automatic reconnection. **Must not contain null (`\\0`) characters.**
+            header on automatic reconnection. **Must not contain null
+            (``\\0``) characters.**
             """
         ),
     ] = None
@@ -114,8 +373,8 @@ class ServerSentEvent(BaseModel):
             """
             Optional reconnection time in **milliseconds**.
 
-            Tells the browser how long to wait before reconnecting after the
+            Tells the browser how long to wait before reconnecting after
-            connection is lost. Must be a non-negative integer.
+            the connection is lost. Must be a non-negative integer.
             """
         ),
     ] = None
@@ -125,9 +384,9 @@ class ServerSentEvent(BaseModel):
             """
             Optional comment line(s).
 
-            Comment lines start with `:` in the SSE wire format and are ignored by
+            Comment lines start with ``:`` in the SSE wire format and are
-            `EventSource` clients. Useful for keep-alive pings to prevent
+            ignored by ``EventSource`` clients. Useful for keep-alive
-            proxy/load-balancer timeouts.
+            pings to prevent proxy/load-balancer timeouts.
             """
         ),
     ] = None
@@ -149,7 +408,7 @@ def format_sse_event(
         str | None,
         Doc(
             """
-            Pre-serialized data string to use as the `data:` field.
+            Pre-serialized data string to use as the ``data:`` field.
             """
         ),
     ] = None,
@@ -157,7 +416,7 @@ def format_sse_event(
         str | None,
         Doc(
             """
-            Optional event type name (`event:` field).
+            Optional event type name (``event:`` field).
             """
         ),
     ] = None,
@@ -165,7 +424,7 @@ def format_sse_event(
         str | None,
         Doc(
             """
-            Optional event ID (`id:` field).
+            Optional event ID (``id:`` field).
             """
         ),
     ] = None,
@@ -173,7 +432,7 @@ def format_sse_event(
         int | None,
         Doc(
             """
-            Optional reconnection time in milliseconds (`retry:` field).
+            Optional reconnection time in milliseconds (``retry:`` field).
             """
         ),
     ] = None,
@@ -181,14 +440,14 @@ def format_sse_event(
         str | None,
         Doc(
             """
-            Optional comment line(s) (`:` prefix).
+            Optional comment line(s) (``:`` prefix).
             """
         ),
     ] = None,
 ) -> bytes:
     """Build SSE wire-format bytes from **pre-serialized** data.
 
-    The result always ends with `\n\n` (the event terminator).
+    The result always ends with ``\\n\\n`` (the event terminator).
     """
     lines: list[str] = []
 

tests/test_sse.py

@@ -316,3 +316,178 @@ 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
+
+
+# ─── Direct-use tests (EventSourceResponse as return value) ──────────────
+
+
+direct_app = FastAPI()
+
+
+@direct_app.get("/direct-events")
+def direct_sse():
+    async def gen():
+        yield ServerSentEvent(data="hello", event="greeting")
+        yield ServerSentEvent(data={"key": "value"}, event="json")
+        yield ServerSentEvent(comment="keepalive")
+
+    return EventSourceResponse(gen())
+
+
+@direct_app.get("/direct-raw-bytes")
+def direct_raw():
+    async def gen():
+        yield b"data: raw line 1\n\n"
+        yield b"data: raw line 2\n\n"
+
+    return EventSourceResponse(gen())
+
+
+@direct_app.get("/direct-with-retry")
+def direct_with_retry():
+    async def gen():
+        yield ServerSentEvent(data="event1")
+        yield ServerSentEvent(data="event2")
+
+    return EventSourceResponse(gen(), retry=5000)
+
+
+@direct_app.get("/direct-with-mixed-types")
+def direct_mixed():
+    async def gen():
+        yield ServerSentEvent(data="sse-event", event="typed")
+        yield b"data: raw bytes\n\n"
+        yield "data: raw string\n\n"
+
+    return EventSourceResponse(gen())
+
+
+@direct_app.get("/direct-negative-retry")
+def direct_negative_retry():
+    return EventSourceResponse(iter([]), retry=-1)  # type: ignore[arg-type]
+
+
+@direct_app.get("/direct-with-disconnect-callback")
+async def direct_with_disconnect_callback():
+    disconnected = []
+
+    async def on_disconnect():
+        disconnected.append(True)
+
+    async def gen():
+        yield ServerSentEvent(data="msg")
+        # Keep streaming so test client can detect disconnect
+        yield ServerSentEvent(data="msg2")
+        yield ServerSentEvent(data="msg3")
+
+    return EventSourceResponse(gen(), disconnect_callback=on_disconnect)
+
+
+@direct_app.get("/direct-with-ping")
+async def direct_with_ping():
+    async def gen():
+        yield ServerSentEvent(data="before-ping")
+        await asyncio.sleep(0.2)  # longer than ping interval
+        yield ServerSentEvent(data="after-ping")
+
+    return EventSourceResponse(gen(), ping_interval=0.05)
+
+
+@direct_app.get("/direct-no-ping")
+async def direct_no_ping():
+    async def gen():
+        yield ServerSentEvent(data="msg1")
+        await asyncio.sleep(0.2)
+        yield ServerSentEvent(data="msg2")
+
+    return EventSourceResponse(gen(), ping_interval=0)
+
+
+direct_client = TestClient(direct_app)
+
+
+def test_direct_use_server_sent_events():
+    """Direct-use: ServerSentEvent objects are auto-formatted."""
+    response = direct_client.get("/direct-events")
+    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: greeting\n" in text
+    assert 'data: "hello"\n' in text
+    assert "event: json\n" in text
+    assert ": keepalive\n" in text
+
+
+def test_direct_use_raw_bytes():
+    """Direct-use: raw bytes are passed through unchanged."""
+    response = direct_client.get("/direct-raw-bytes")
+    assert response.status_code == 200
+    assert "data: raw line 1\n" in response.text
+    assert "data: raw line 2\n" in response.text
+
+
+def test_direct_use_retry_field():
+    """Direct-use: retry parameter emits an initial retry: field."""
+    response = direct_client.get("/direct-with-retry")
+    assert response.status_code == 200
+    text = response.text
+    # The initial retry field should appear before the data events
+    retry_pos = text.find("retry: 5000")
+    data_pos = text.find("data:")
+    assert retry_pos >= 0
+    assert data_pos > retry_pos
+
+
+def test_direct_use_mixed_types():
+    """Direct-use: ServerSentEvent, bytes, and strings handled correctly."""
+    response = direct_client.get("/direct-with-mixed-types")
+    assert response.status_code == 200
+    text = response.text
+    assert "event: typed\n" in text
+    assert "data: raw bytes\n" in text
+    assert "data: raw string\n" in text
+
+
+def test_direct_use_negative_retry_rejected():
+    """Direct-use: negative retry raises ValueError."""
+    with pytest.raises(ValueError, match="non-negative"):
+        direct_client.get("/direct-negative-retry")
+
+
+def test_direct_use_keepalive_ping():
+    """Direct-use with ping_interval emits : ping comments on idle."""
+    response = direct_client.get("/direct-with-ping")
+    assert response.status_code == 200
+    text = response.text
+    assert ": ping\n" in text
+    assert "data: \"before-ping\"\n" in text
+    assert "data: \"after-ping\"\n" in text
+
+
+def test_direct_use_no_ping_when_disabled():
+    """Direct-use with ping_interval=0 emits no keepalive pings."""
+    response = direct_client.get("/direct-no-ping")
+    assert response.status_code == 200
+    assert ": ping\n" not in response.text
+    assert "data: \"msg1\"\n" in response.text
+    assert "data: \"msg2\"\n" in response.text
+
+
+def test_direct_use_sync_iterable():
+    """Direct-use works with sync iterables (wrapped to async)."""
+
+    @direct_app.get("/direct-sync")
+    def direct_sync():
+        def gen():
+            yield ServerSentEvent(data="sync1")
+            yield ServerSentEvent(data="sync2")
+
+        return EventSourceResponse(gen())
+
+    response = direct_client.get("/direct-sync")
+    assert response.status_code == 200
+    assert 'data: "sync1"\n' in response.text
+    assert 'data: "sync2"\n' in response.text