Sebastián Ramírez
4 months ago
4 changed files with 389 additions and 35 deletions
@ -0,0 +1,222 @@ |
|||
from typing import Annotated, Any |
|||
|
|||
from annotated_doc import Doc |
|||
from pydantic import AfterValidator, BaseModel, Field, model_validator |
|||
from starlette.responses import StreamingResponse |
|||
|
|||
# Canonical SSE event schema matching the OpenAPI 3.2 spec |
|||
# (Section 4.14.4 "Special Considerations for Server-Sent Events") |
|||
_SSE_EVENT_SCHEMA: dict[str, Any] = { |
|||
"type": "object", |
|||
"properties": { |
|||
"data": {"type": "string"}, |
|||
"event": {"type": "string"}, |
|||
"id": {"type": "string"}, |
|||
"retry": {"type": "integer", "minimum": 0}, |
|||
}, |
|||
} |
|||
|
|||
|
|||
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`. |
|||
""" |
|||
|
|||
media_type = "text/event-stream" |
|||
|
|||
|
|||
def _check_id_no_null(v: int | str | None) -> int | str | None: |
|||
if v is not None and "\0" in str(v): |
|||
raise ValueError("SSE 'id' must not contain null characters") |
|||
return v |
|||
|
|||
|
|||
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`). |
|||
|
|||
If you yield a plain object (dict, Pydantic model, etc.) instead, it is |
|||
automatically JSON-encoded and sent as the `data:` field. |
|||
|
|||
All `data` values **including plain strings** are JSON-serialized. |
|||
|
|||
For example, `data="hello"` produces `data: "hello"` on the wire (with |
|||
quotes). |
|||
""" |
|||
|
|||
data: Annotated[ |
|||
Any, |
|||
Doc( |
|||
""" |
|||
The event payload. |
|||
|
|||
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). |
|||
|
|||
Mutually exclusive with `raw_data`. |
|||
""" |
|||
), |
|||
] = None |
|||
raw_data: Annotated[ |
|||
str | None, |
|||
Doc( |
|||
""" |
|||
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. |
|||
|
|||
Mutually exclusive with `data`. |
|||
""" |
|||
), |
|||
] = None |
|||
event: Annotated[ |
|||
str | None, |
|||
Doc( |
|||
""" |
|||
Optional event type name. |
|||
|
|||
Maps to `addEventListener(event, ...)` on the browser. When omitted, |
|||
the browser dispatches on the generic `message` event. |
|||
""" |
|||
), |
|||
] = None |
|||
id: Annotated[ |
|||
int | str | None, |
|||
AfterValidator(_check_id_no_null), |
|||
Doc( |
|||
""" |
|||
Optional event ID. |
|||
|
|||
The browser sends this value back as the `Last-Event-ID` header on |
|||
automatic reconnection. **Must not contain null (`\\0`) characters.** |
|||
""" |
|||
), |
|||
] = None |
|||
retry: Annotated[ |
|||
int | None, |
|||
Field(ge=0), |
|||
Doc( |
|||
""" |
|||
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. |
|||
""" |
|||
), |
|||
] = None |
|||
comment: Annotated[ |
|||
str | None, |
|||
Doc( |
|||
""" |
|||
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. |
|||
""" |
|||
), |
|||
] = None |
|||
|
|||
@model_validator(mode="after") |
|||
def _check_data_exclusive(self) -> "ServerSentEvent": |
|||
if self.data is not None and self.raw_data is not None: |
|||
raise ValueError( |
|||
"Cannot set both 'data' and 'raw_data' on the same " |
|||
"ServerSentEvent. Use 'data' for JSON-serialized payloads " |
|||
"or 'raw_data' for pre-formatted strings." |
|||
) |
|||
return self |
|||
|
|||
|
|||
def format_sse_event( |
|||
*, |
|||
data_str: Annotated[ |
|||
str | None, |
|||
Doc( |
|||
""" |
|||
Pre-serialized data string to use as the `data:` field. |
|||
""" |
|||
), |
|||
] = None, |
|||
event: Annotated[ |
|||
str | None, |
|||
Doc( |
|||
""" |
|||
Optional event type name (`event:` field). |
|||
""" |
|||
), |
|||
] = None, |
|||
id: Annotated[ |
|||
int | str | None, |
|||
Doc( |
|||
""" |
|||
Optional event ID (`id:` field). |
|||
""" |
|||
), |
|||
] = None, |
|||
retry: Annotated[ |
|||
int | None, |
|||
Doc( |
|||
""" |
|||
Optional reconnection time in milliseconds (`retry:` field). |
|||
""" |
|||
), |
|||
] = None, |
|||
comment: Annotated[ |
|||
str | None, |
|||
Doc( |
|||
""" |
|||
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). |
|||
""" |
|||
lines: list[str] = [] |
|||
|
|||
if comment is not None: |
|||
for line in comment.splitlines(): |
|||
lines.append(f": {line}") |
|||
|
|||
if event is not None: |
|||
lines.append(f"event: {event}") |
|||
|
|||
if data_str is not None: |
|||
for line in data_str.splitlines(): |
|||
lines.append(f"data: {line}") |
|||
|
|||
if id is not None: |
|||
lines.append(f"id: {id}") |
|||
|
|||
if retry is not None: |
|||
lines.append(f"retry: {retry}") |
|||
|
|||
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 |
|||
Loading…
Reference in new issue