mirror
/
tiangolo.fastapi
mirror of https://github.com/tiangolo/fastapi
1
0
Fork 0
Code Issues Projects Releases Wiki Activity
Browse Source

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]>
pull/15428/head
ubuntu 3 months ago
parent
4fa826ce0a
commit
8b3588f3bf
2 changed files with 474 additions and 40 deletions
Whitespace
Show all changes Ignore whitespace when comparing lines Ignore changes in amount of whitespace Ignore changes in whitespace at EOL
Split View
Diff Options
Show Stats Download Patch File Download Diff File
  1. 339
      fastapi/sse.py
  2. 175
      tests/test_sse.py

339
fastapi/sse.py
View File

@ -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
<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
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] = []

175
tests/test_sse.py
View File

@ -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

Write Preview
Loading…
Cancel
Save