server/sse.py

"""Asyncio broadcast hub for /api/stream (ARCHITECTURE.md §SSE protocol).

One hub, many subscribers; every subscriber receives every published event. Design
constraints, in order of importance:

  * The producer (the queue worker granting a wish) must NEVER block or fail because
    a browser tab is slow: each subscriber gets a bounded queue and the OLDEST event
    is dropped on overflow. A reconnecting client re-syncs via GET /api/state.
  * Idle connections receive ``{"type": "heartbeat"}`` every 15s so proxies (and the
    HF Space router) keep the stream open.
  * Disconnects clean up after themselves: the wire generator unsubscribes in a
    ``finally`` block, which Starlette triggers on client disconnect.
"""

from __future__ import annotations

import asyncio
import json
import logging
from typing import Any, AsyncIterator, Iterable, Optional

log = logging.getLogger("godseed.sse")

HEARTBEAT_SECONDS = 15.0
DEFAULT_QUEUE_SIZE = 256


def format_sse(event: dict[str, Any]) -> str:
    """Wire-encode one event. Events are plain JSON on the default `message` channel;
    clients dispatch on ``data.type`` (matches the contract's event objects)."""
    return f"data: {json.dumps(event, ensure_ascii=False, separators=(',', ':'))}\n\n"


class SSEHub:
    """Fan-out broadcast hub. ``publish`` is synchronous and non-blocking; ``emit`` is
    the async alias matching the planner's ``emit`` callback signature."""

    def __init__(
        self,
        heartbeat: float = HEARTBEAT_SECONDS,
        max_queue: int = DEFAULT_QUEUE_SIZE,
    ) -> None:
        self.heartbeat = heartbeat
        self.max_queue = max_queue
        self._subscribers: set[asyncio.Queue[dict[str, Any]]] = set()

    @property
    def subscriber_count(self) -> int:
        return len(self._subscribers)

    # ------------------------------------------------------------------ producer side
    def publish(self, event: dict[str, Any]) -> None:
        """Deliver ``event`` to every subscriber. Never blocks: when a subscriber's
        queue is full, its oldest pending event is dropped to make room."""
        for queue in tuple(self._subscribers):
            try:
                queue.put_nowait(event)
            except asyncio.QueueFull:
                try:
                    queue.get_nowait()  # drop the oldest event for this slow client
                except asyncio.QueueEmpty:  # pragma: no cover — only under a race
                    pass
                try:
                    queue.put_nowait(event)
                except asyncio.QueueFull:  # pragma: no cover — only under a race
                    log.debug("subscriber queue still full; event dropped")

    async def emit(self, event: dict[str, Any]) -> None:
        """Async alias of :meth:`publish` (the planner/queue-worker emit callback)."""
        self.publish(event)

    # ---------------------------------------------------------------- subscriber side
    def subscribe(self) -> asyncio.Queue[dict[str, Any]]:
        queue: asyncio.Queue[dict[str, Any]] = asyncio.Queue(maxsize=self.max_queue)
        self._subscribers.add(queue)
        return queue

    def unsubscribe(self, queue: asyncio.Queue[dict[str, Any]]) -> None:
        self._subscribers.discard(queue)

    async def event_stream(
        self,
        initial_events: Iterable[dict[str, Any]] = (),
        limit: Optional[int] = None,
    ) -> AsyncIterator[str]:
        """Async generator yielding wire-format SSE for one client.

        ``initial_events`` (hello / queue snapshot) are sent first, then every
        broadcast event; a heartbeat is sent whenever ``self.heartbeat`` seconds pass
        without traffic. Subscription is registered before the first yield so no
        event published after the request reaches this generator is missed.

        ``limit`` (optional) closes the stream after that many events — a debug/test
        aid (`curl '/api/stream?limit=20'`); browsers stream unbounded.
        """
        queue = self.subscribe()
        sent = 0
        try:
            yield "retry: 3000\n\n"  # client auto-reconnect hint
            for event in initial_events:
                yield format_sse(event)
                sent += 1
                if limit is not None and sent >= limit:
                    return
            while True:
                try:
                    event = await asyncio.wait_for(queue.get(), timeout=self.heartbeat)
                except asyncio.TimeoutError:
                    event = {"type": "heartbeat"}
                yield format_sse(event)
                sent += 1
                if limit is not None and sent >= limit:
                    return
        finally:
            self.unsubscribe(queue)