src/genro_asgi/sse.py¶
Source from this local checkout, regenerated when the reader rebuilds.
Line links use #L<number>; a GitHub line range opens its first line.
1 # Copyright 2025 Softwell S.r.l.2 #3 # Licensed under the Apache License, Version 2.0 (the "License");4 # you may not use this file except in compliance with the License.5 # You may obtain a copy of the License at6 #7 # https://www.apache.org/licenses/LICENSE-2.08 #9 # Unless required by applicable law or agreed to in writing, software10 # distributed under the License is distributed on an "AS IS" BASIS,11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12 # See the License for the specific language governing permissions and13 # limitations under the License.14 15 """Server-Sent Events framing over ``StreamingResponse``.16 17 An event is a small dict — ``{"data": ..., "event": ..., "id": ...}`` — framed18 into the ``text/event-stream`` wire format: ``id:``/``event:``/``data:`` lines19 (``data`` split across lines on newlines, per the spec), ``retry:`` once at the20 start when configured, and a ``: keepalive`` comment when the source falls21 silent longer than the heartbeat interval (the comment keeps proxies from22 closing an idle connection; the client ignores it). Each event ends with a23 blank line. ``data`` that is not a string is JSON-encoded.24 25 The framing is shaped like ``channel/frame.py`` (a slotted codec, its own wire26 format) but has no bytes in common — SSE is a text protocol over HTTP, not the27 length-prefixed wsx envelope. ``SseStream`` is SELF-CONTAINED: it wraps ANY28 async source of event dicts (a user generator, a task hub subscription) and29 yields wire ``bytes``; the source is the caller's concern. Resumability30 (``Last-Event-ID`` → a snapshot baseline then the live source) is built by the31 consumer that owns the event source, not here.32 """33 34 from __future__ import annotations35 36 import asyncio37 import contextlib38 import json39 from collections.abc import AsyncIterable, AsyncIterator40 from typing import Any41 42 from .streaming import StreamingResponse43 44 __all__ = ["SseStream", "KEEPALIVE_SECONDS"]45 46 KEEPALIVE_SECONDS = 15.0 # idle interval after which a ``: keepalive`` comment goes out47 _KEEPALIVE_FRAME = b": keepalive\n\n"48 49 50 class SseStream:51 """Frames an async source of event dicts into ``text/event-stream`` bytes.52 53 Note:54 The stream is bound to one source (dual relationship: ``self.source``).55 Iterating it yields wire bytes; ``response()`` wraps it in a56 ``StreamingResponse`` with the SSE headers already set.57 """58 59 __slots__ = ("source", "retry_ms", "keepalive_seconds")60 61 def __init__(62 self,63 source: AsyncIterable[dict[str, Any]],64 *,65 retry_ms: int | None = None,66 keepalive_seconds: float = KEEPALIVE_SECONDS,67 ) -> None:68 """Bind to an async source of event dicts.69 70 Args:71 source: Any async iterable of events; each event is a dict with an72 optional ``id``/``event`` and a ``data`` payload.73 retry_ms: When set, a ``retry:`` line is emitted once at the start74 (the client's reconnection delay).75 keepalive_seconds: Idle interval after which a ``: keepalive``76 comment is sent to hold the connection open.77 """78 self.source = source79 self.retry_ms = retry_ms80 self.keepalive_seconds = keepalive_seconds81 82 def frame(self, event: dict[str, Any]) -> bytes:83 """Encode one event dict into an SSE record (ends with a blank line)."""84 lines: list[str] = []85 if event.get("id") is not None:86 lines.append(f"id: {event['id']}")87 if event.get("event") is not None:88 lines.append(f"event: {event['event']}")89 data = event.get("data")90 text = data if isinstance(data, str) else json.dumps(data)91 for line in text.split("\n"):92 lines.append(f"data: {line}")93 return ("\n".join(lines) + "\n\n").encode("utf-8")94 95 async def __aiter__(self) -> AsyncIterator[bytes]:96 """Yield wire bytes: an optional ``retry:``, then framed events.97 98 The source is consumed one event at a time; while it stays silent past99 ``keepalive_seconds`` a ``: keepalive`` comment is emitted so the100 connection is not reaped (the pending read is shielded across the101 timeout). The loop ends when the source is exhausted; when the CONSUMER102 goes away instead (cancellation / close), the shielded read is cancelled103 and awaited so the source's ``finally`` runs before the stream unwinds104 (a subscription source must get to unsubscribe).105 """106 if self.retry_ms is not None:107 yield f"retry: {self.retry_ms}\n\n".encode()108 iterator = self.source.__aiter__()109 nxt: asyncio.Future[dict[str, Any]] | None = None110 try:111 while True:112 nxt = asyncio.ensure_future(iterator.__anext__())113 while True:114 try:115 event = await asyncio.wait_for(116 asyncio.shield(nxt), self.keepalive_seconds117 )118 except asyncio.TimeoutError:119 yield _KEEPALIVE_FRAME # source silent: hold the connection120 continue121 except StopAsyncIteration:122 return # source exhausted: end the stream123 break124 nxt = None125 yield self.frame(event)126 finally:127 if nxt is not None and not nxt.done():128 nxt.cancel() # client gone: close the source129 with contextlib.suppress(asyncio.CancelledError, StopAsyncIteration):130 await nxt131 132 def response(self, status_code: int = 200) -> StreamingResponse:133 """A ``StreamingResponse`` over this stream with the SSE headers set."""134 return StreamingResponse(135 self,136 status_code=status_code,137 headers=[138 ("cache-control", "no-cache"),139 ("connection", "keep-alive"),140 ],141 media_type="text/event-stream",142 )