Import AITURK IDE 1.0.0-beta.1 from Hermes 63279301; preserve MIT license

tui_gateway/event_replay.py

@@ -0,0 +1,125 @@
+"""Per-session event sequencing + bounded replay for WS reconnects.
+
+Every gateway event frame that flows through :func:`server.write_json` (and
+therefore ``_emit``) is stamped with a per-session monotonic ``seq`` and
+appended to a small ring buffer keyed by session id. A reconnecting client
+calls the ``session.events.since`` RPC with its last observed seq; the server
+replays everything newer from the buffer, then live events resume seamlessly.
+
+Design constraints honored:
+- stdio TUI path unaffected: frames gain a ``seq`` field only on event frames;
+  Ink ignores unknown params keys.
+- Thread safety: a single module lock guards counters + buffers; write_json
+  already serializes per-transport writes, so stamping under the lock cannot
+  reorder frames relative to each other.
+- Memory bound: _REPLAY_BUFFER_MAX events / _REPLAY_SESSIONS_MAX sessions,
+  oldest session evicted FIFO.
+"""
+
+from __future__ import annotations
+
+import threading
+import uuid
+from collections import OrderedDict, deque
+
+# Process identity for the replay contract. Seq counters live in-process, so
+# a gateway restart silently resets them to 1 while clients still hold high
+# watermarks — events_since(sid, 97) then returns [] with truncated=False and
+# the client believes it missed nothing (and its stale watermark makes every
+# future replay empty too). The epoch lets clients detect the restart and
+# reset their watermarks.
+_REPLAY_EPOCH = uuid.uuid4().hex
+
+# Replay ring per session. A long turn emits ~hundreds of token events; this
+# covers several minutes of streaming plus all control events.
+_REPLAY_BUFFER_MAX = 512
+# Distinct sessions remembered. Desktop users rarely exceed a dozen live chats.
+_REPLAY_SESSIONS_MAX = 64
+
+_replay_lock = threading.Lock()
+# sid -> deque of (seq, event_object) where event_object is the frame's
+# ``params`` dict (bare event: type/session_id/seq/payload) — the exact shape
+# the client's dispatch path consumes.
+_replay_buffers: "OrderedDict[str, deque]" = OrderedDict()
+_replay_next_seq: dict[str, int] = {}
+
+
+def replay_epoch() -> str:
+    """Opaque token identifying this server process's seq numbering."""
+    return _REPLAY_EPOCH
+
+
+def _stamp_event(obj: dict) -> None:
+    """Stamp one outgoing event frame (mutates obj in place) and record it."""
+    if obj.get("method") != "event":
+        return
+    params = obj.get("params")
+    if not isinstance(params, dict):
+        return
+    sid = params.get("session_id") or ""
+    if not sid:
+        # Session-less global events (skin.changed etc.) are re-fetchable via
+        # their own RPCs; no replay contract for them.
+        return
+    with _replay_lock:
+        seq = _replay_next_seq.get(sid, 0) + 1
+        _replay_next_seq[sid] = seq
+        params["seq"] = seq
+        buf = _replay_buffers.get(sid)
+        if buf is None:
+            buf = deque(maxlen=_REPLAY_BUFFER_MAX)
+            _replay_buffers[sid] = buf
+            while len(_replay_buffers) > _REPLAY_SESSIONS_MAX:
+                _oldest_sid, _oldest_buf = _replay_buffers.popitem(last=False)
+                _replay_next_seq.pop(_oldest_sid, None)
+        buf.append((seq, params))
+
+
+def events_since(sid: str, last_seen: int) -> list[dict]:
+    """Return recorded EVENT OBJECTS with seq > last_seen for *sid*, in order.
+
+    Shape contract: each element is the frame's ``params`` dict — a bare event
+    object with top-level ``type`` / ``session_id`` / ``seq`` — because that is
+    exactly what the client's dispatch path consumes. Returning the full
+    JSON-RPC envelope here would make every replayed event fail the client's
+    ``event.type`` gate and be silently dropped.
+    """
+    with _replay_lock:
+        buf = _replay_buffers.get(sid or "")
+        if not buf:
+            return []
+        return [event for seq, event in buf if seq > last_seen]
+
+
+def is_truncated(sid: str, last_seen: int) -> bool:
+    """True when events between *last_seen* and the ring's oldest retained
+    seq were evicted — the client must refetch history instead of trusting
+    the replay to be gap-free."""
+    with _replay_lock:
+        buf = _replay_buffers.get(sid or "")
+        if not buf:
+            return False
+        return last_seen + 1 < buf[0][0]
+
+
+def latest_seq(sid: str) -> int:
+    """Current highest stamped seq for *sid* (0 when unknown)."""
+    with _replay_lock:
+        return _replay_next_seq.get(sid or "", 0)
+
+
+def reset_replay_state() -> None:
+    """Test hook."""
+    with _replay_lock:
+        _replay_buffers.clear()
+        _replay_next_seq.clear()
+
+
+def replay_stats() -> dict:
+    """Telemetry: buffer occupancy for the ops/debug surface."""
+    with _replay_lock:
+        return {
+            "sessions": len(_replay_buffers),
+            "events": sum(len(b) for b in _replay_buffers.values()),
+            "max_per_session": _REPLAY_BUFFER_MAX,
+        }