Joseph Doherty
89 lines
3.5 KiB
Python
89 lines
3.5 KiB
Python
"""Server-Sent Events endpoint for per-chat live updates.
|
|
|
|
Each browser tab on ``/chats/<id>`` opens an SSE connection here. On connect:
|
|
|
|
1. We verify the chat exists (404 otherwise).
|
|
2. We subscribe to the chat's pub/sub channel.
|
|
3. We emit a ``snapshot`` event with the current state. T16 only provides a
|
|
stub payload (``{"chat_id": <id>, "ready": true}``) so the client can
|
|
confirm the channel is live; T19+ will populate it with real state.
|
|
4. We loop, awaiting events from the queue and yielding them as SSE frames.
|
|
A 15-second keepalive comment is emitted on idle to defeat intermediary
|
|
timeouts.
|
|
5. When the client disconnects we unsubscribe so the registry doesn't leak.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
|
|
from fastapi import APIRouter, Depends, HTTPException, Request
|
|
from fastapi.responses import StreamingResponse
|
|
|
|
from chat.state.world import get_chat
|
|
from chat.web.bots import get_conn
|
|
from chat.web.pubsub import subscribe, unsubscribe
|
|
|
|
router = APIRouter()
|
|
|
|
# Heartbeat cadence. Long enough to avoid chattiness; short enough that most
|
|
# HTTP intermediaries won't close an idle connection.
|
|
_KEEPALIVE_SECONDS = 15.0
|
|
|
|
|
|
def _format_sse(event: str, data: dict | str) -> bytes:
|
|
"""Format a single SSE frame: ``event: <name>\\ndata: <body>\\n\\n``.
|
|
|
|
``data`` may be a dict (JSON-serialized) or a raw string. The string
|
|
form is used for HTMX SSE swaps where the payload is an HTML
|
|
fragment that the client splices into the DOM verbatim.
|
|
"""
|
|
if isinstance(data, str):
|
|
payload = data
|
|
else:
|
|
payload = json.dumps(data)
|
|
return f"event: {event}\ndata: {payload}\n\n".encode("utf-8")
|
|
|
|
|
|
@router.get("/chats/{chat_id}/events")
|
|
async def chat_events(chat_id: str, request: Request, conn=Depends(get_conn)):
|
|
chat = get_chat(conn, chat_id)
|
|
if chat is None:
|
|
raise HTTPException(status_code=404, detail="chat not found")
|
|
|
|
async def stream():
|
|
queue = await subscribe(chat_id)
|
|
try:
|
|
# Initial snapshot — T19 will fill in real state.
|
|
yield _format_sse("snapshot", {"chat_id": chat_id, "ready": True})
|
|
while True:
|
|
if await request.is_disconnected():
|
|
break
|
|
try:
|
|
event = await asyncio.wait_for(
|
|
queue.get(), timeout=_KEEPALIVE_SECONDS
|
|
)
|
|
except asyncio.TimeoutError:
|
|
# SSE comment line (per spec, lines starting with ":" are
|
|
# ignored by the client) — keeps the connection warm.
|
|
yield b": keepalive\n\n"
|
|
continue
|
|
# Allow publishers to set the SSE event name via "event" key;
|
|
# default to "message" if omitted. When the remaining payload
|
|
# is a single ``data`` string, send it verbatim — that lets
|
|
# turn-flow publishers ship pre-rendered HTML fragments that
|
|
# HTMX's SSE extension can swap into the DOM directly.
|
|
event = dict(event) # don't mutate the published dict
|
|
kind = event.pop("event", "message")
|
|
if set(event.keys()) == {"data"} and isinstance(
|
|
event["data"], str
|
|
):
|
|
yield _format_sse(kind, event["data"])
|
|
else:
|
|
yield _format_sse(kind, event)
|
|
finally:
|
|
await unsubscribe(chat_id, queue)
|
|
|
|
return StreamingResponse(stream(), media_type="text/event-stream")
|