62 lines
1.9 KiB
Python
62 lines
1.9 KiB
Python
"""Server-sent events helpers shared by the markdown and API frontends."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import contextlib
|
|
import json
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import AsyncIterator
|
|
|
|
__all__ = ["HEARTBEAT_INTERVAL", "SSE_HEADERS", "events_with_heartbeat", "sse_pack"]
|
|
|
|
HEARTBEAT_INTERVAL = 15.0
|
|
"""Seconds of backend silence before a comment frame keeps the socket warm."""
|
|
|
|
SSE_HEADERS = {
|
|
"Cache-Control": "no-cache, no-transform",
|
|
"Connection": "keep-alive",
|
|
"X-Accel-Buffering": "no",
|
|
}
|
|
|
|
KEEPALIVE = b": keepalive\n\n"
|
|
|
|
|
|
async def events_with_heartbeat(
|
|
events: AsyncIterator[Any], interval: float = HEARTBEAT_INTERVAL
|
|
) -> AsyncIterator[Any]:
|
|
"""Pass ``events`` through, yielding ``None`` after ``interval`` seconds of silence.
|
|
|
|
Only one consumer may iterate the result at a time; cancelling the
|
|
outer scope cancels the in-flight upstream fetch instead of leaking it.
|
|
"""
|
|
src = events.__aiter__()
|
|
next_task: asyncio.Task[Any] | None = None
|
|
try:
|
|
while True:
|
|
if next_task is None:
|
|
next_task = asyncio.ensure_future(src.__anext__())
|
|
done, _pending = await asyncio.wait({next_task}, timeout=interval)
|
|
if not done:
|
|
yield None
|
|
continue
|
|
task = next_task
|
|
next_task = None
|
|
try:
|
|
result = task.result()
|
|
except StopAsyncIteration:
|
|
return
|
|
yield result
|
|
finally:
|
|
if next_task is not None and not next_task.done():
|
|
next_task.cancel()
|
|
with contextlib.suppress(BaseException):
|
|
await next_task
|
|
|
|
|
|
def sse_pack(event: str, data: dict[str, Any]) -> bytes:
|
|
body = json.dumps(data, ensure_ascii=False)
|
|
return f"event: {event}\ndata: {body}\n\n".encode()
|