From 8b3588f3bfbba4fab49cf42201601c7c1f6b9df6 Mon Sep 17 00:00:00 2001 From: ubuntu Date: Sat, 25 Apr 2026 11:49:55 +0800 Subject: [PATCH] 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 --- fastapi/sse.py | 339 ++++++++++++++++++++++++++++++++++++++++------ tests/test_sse.py | 175 ++++++++++++++++++++++++ 2 files changed, 474 insertions(+), 40 deletions(-) diff --git a/fastapi/sse.py b/fastapi/sse.py index 901d824964..d34f618620 100644 --- a/fastapi/sse.py +++ b/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. - - Use as `response_class=EventSourceResponse` on a *path operation* that uses `yield` - to enable Server Sent Events (SSE) responses. - - 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`. + """Streaming response with ``text/event-stream`` media type. + + Use as ``response_class=EventSourceResponse`` on a *path operation* that + uses ``yield`` to enable Server-Sent Events (SSE) responses via the + routing layer's automatic encoding. + + Can also be used **directly as a return value**, similar to + ``StreamingResponse``, by passing an async iterable of pre-formatted + SSE bytes or ``ServerSentEvent`` objects:: + + @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 - `response_class=EventSourceResponse`, each `ServerSentEvent` is encoded - into the [SSE wire format](https://html.spec.whatwg.org/multipage/server-sent-events.html#parsing-an-event-stream) - (`text/event-stream`). + When ``yield``ed from a *path operation function* that uses + ``response_class=EventSourceResponse``, each ``ServerSentEvent`` is + encoded into the `SSE wire format + `_ + (``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 - quotes). + For example, ``data="hello"`` produces ``data: "hello"`` on the wire + (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, - CSV lines, or any non-JSON payload. The string is placed directly - into the `data:` field as-is. + Use this when you need to send pre-formatted text, HTML + fragments, CSV lines, or any non-JSON payload. The string 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, - the browser dispatches on the generic `message` event. + Maps to ``addEventListener(event, ...)`` on the browser. When + 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 - automatic reconnection. **Must not contain null (`\\0`) characters.** + The browser sends this value back as the ``Last-Event-ID`` + 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 - connection is lost. Must be a non-negative integer. + Tells the browser how long to wait before reconnecting after + 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 - `EventSource` clients. Useful for keep-alive pings to prevent - proxy/load-balancer timeouts. + Comment lines start with ``:`` in the SSE wire format and are + ignored by ``EventSource`` clients. Useful for keep-alive + 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] = [] diff --git a/tests/test_sse.py b/tests/test_sse.py index 6dfec61838..67bf637c48 100644 --- a/tests/test_sse.py +++ b/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