arborist-viz/arborist_viz/views/sse.py

"""Server-Sent Events endpoints — §3-corrected replacement for WebSockets.

Phase 4 tails arborist's audit log by polling ``arborist.read.Shards``
for events newer than the last opaque cursor. Each new event is emitted
as a ``patch`` event shaped like a §8.3 ``GraphPatch`` so the dashboard
can apply it incrementally.

Pyramid streams via ``Response.app_iter`` (a generator). Each event is
flushed immediately — no line buffering — matching the CLAUDE.md
"never buffer long-running output" rule.
"""

from __future__ import annotations

import json
import logging
import time
from typing import Iterator, Optional

from pyramid.response import Response
from pyramid.view import view_config


log = logging.getLogger(__name__)


def _format_event(event: str | None, data: dict, event_id: str | None = None) -> bytes:
    lines = []
    if event_id is not None:
        lines.append("id: {0}".format(event_id))
    if event is not None:
        lines.append("event: {0}".format(event))
    lines.append("data: {0}".format(json.dumps(data, separators=(",", ":"), default=str)))
    lines.append("")  # blank line terminates the event
    lines.append("")
    return ("\n".join(lines)).encode("utf-8")


def _audit_tail(
    shards_handle,
    *,
    heartbeat_seconds: int = 15,
    poll_seconds: float = 1.0,
    subject_root_filter: Optional[str] = None,
    run_id_filter: Optional[str] = None,
) -> Iterator[bytes]:
    """Yield SSE events tailing arborist's audit log.

    ``shards_handle`` is an ``arborist.read.Shards`` instance or ``None``.
    Polls ``audit_since(cursor)`` once per ``poll_seconds``; emits a
    ``patch`` event per new row, then a ``heartbeat`` if the poll period
    exceeded ``heartbeat_seconds`` with no activity. Handles
    shard-not-configured by emitting one ``ready`` event and continuing
    to heartbeat.
    """
    paths = list(shards_handle.paths) if shards_handle is not None else []
    cursor: dict = shards_handle.audit_cursor() if shards_handle is not None else {}

    seq = 0
    yield _format_event(
        "ready",
        {
            "seq": seq,
            "stream": "audit.tail",
            "shards": paths,
            "shard_open": shards_handle is not None,
            "shard_count": len(paths),
            "subject_root_filter": subject_root_filter,
            "run_id_filter": run_id_filter,
            "start_seqs": dict(cursor),
        },
        event_id=str(seq),
    )

    if shards_handle is None:
        while True:
            seq += 1
            time.sleep(heartbeat_seconds)
            try:
                yield _format_event(
                    "heartbeat",
                    {"seq": seq, "ts": time.time(), "shard_open": False},
                    event_id=str(seq),
                )
            except GeneratorExit:
                return

    last_emit = time.time()
    log.info("SSE audit tail across %d shard(s) (filter root=%s run=%s)",
             len(paths), subject_root_filter, run_id_filter)

    while True:
        try:
            new_events, cursor = shards_handle.audit_since(cursor, limit=200)
            new_any = False
            for ev in new_events:
                if subject_root_filter and ev.subject_root != subject_root_filter:
                    continue
                # body.run_id is best-effort — not all events carry one.
                if run_id_filter and ev.body.get("run_id") != run_id_filter:
                    continue
                new_any = True
                seq += 1
                last_emit = time.time()
                yield _format_event(
                    "patch",
                    {
                        "seq": seq,
                        "shard": ev.shard_path,
                        "event_hash": ev.event_hash,
                        "prev_event_hash": ev.prev_event_hash,
                        "event_type": ev.event_type,
                        "subject_root": ev.subject_root,
                        "ts": ev.ts,
                        "body": ev.body,
                        "patches": [
                            ["node.upsert", ev.subject_root or ev.event_hash,
                             {"kind": "audit_event" if not ev.subject_root else "document_root",
                              "hash": ev.subject_root or ev.event_hash,
                              "status": "live"}],
                        ],
                    },
                    event_id=str(seq),
                )
            if not new_any and time.time() - last_emit >= heartbeat_seconds:
                seq += 1
                last_emit = time.time()
                yield _format_event(
                    "heartbeat",
                    {"seq": seq, "ts": time.time(),
                     "last_audit_seqs": dict(cursor)},
                    event_id=str(seq),
                )
            time.sleep(poll_seconds)
        except GeneratorExit:
            log.info("SSE audit tail closed at seq=%s", seq)
            return
        except Exception as exc:  # pragma: no cover — defensive
            log.warning("SSE audit tail error: %s", exc)
            time.sleep(poll_seconds)


def _sse_response(request, generator: Iterator[bytes]) -> Response:
    response = Response(
        content_type="text/event-stream",
        charset="utf-8",
    )
    # SSE spec + reverse-proxy hardening. "Connection" is hop-by-hop
    # (PEP 3333) — leave it to the WSGI server.
    response.headers["Cache-Control"] = "no-cache, no-transform"
    response.headers["X-Accel-Buffering"] = "no"
    response.app_iter = generator
    return response


def _settings_hb(request) -> int:
    try:
        return int(request.registry.settings.get("sse.heartbeat_seconds", "15"))
    except (TypeError, ValueError):
        return 15


@view_config(route_name="sse_audit_live")
def sse_audit_live(request):
    return _sse_response(
        request,
        _audit_tail(request.shards, heartbeat_seconds=_settings_hb(request)),
    )


@view_config(route_name="sse_run")
def sse_run(request):
    return _sse_response(
        request,
        _audit_tail(
            request.shards,
            heartbeat_seconds=_settings_hb(request),
            run_id_filter=request.matchdict["run_id"],
        ),
    )


@view_config(route_name="sse_root")
def sse_root(request):
    return _sse_response(
        request,
        _audit_tail(
            request.shards,
            heartbeat_seconds=_settings_hb(request),
            subject_root_filter=request.matchdict["root_hash"],
        ),
    )


@view_config(route_name="sse_claim")
def sse_claim(request):
    # Claims live on top of arborist's providence/audit projection;
    # for now we tail the same audit stream.
    return _sse_response(
        request,
        _audit_tail(request.shards, heartbeat_seconds=_settings_hb(request)),
    )


@view_config(route_name="sse_dashboard")
def sse_dashboard(request):
    return _sse_response(
        request,
        _audit_tail(request.shards, heartbeat_seconds=_settings_hb(request)),
    )


@view_config(route_name="sse_mesh_live")
def sse_mesh_live(request):
    # Mesh roster + epochs land in a future phase; heartbeat-only for now.
    return _sse_response(
        request,
        _audit_tail(request.shards, heartbeat_seconds=_settings_hb(request)),
    )