""" WebSocket Event Bus — replaces sse_bus.py for real-time communication. Drop-in replacement: all routers continue calling broadcast_sync() unchanged. The WS endpoint uses connect()/disconnect()/replay_missed() to manage clients. Protocol (JSON frames): Server → Client: { "seq": 123, "type": "order_updated", "data": { ... } } Client → Server: { "cursor": 120 } (sent immediately after connect) Server → Client: { "type": "ping" } (every 25s keepalive) Client → Server: { "type": "pong" } (optional, ignored if not sent) On connect the client sends its last known seq. The server replays everything it has stored since that seq, then switches to live streaming. Events are stored in the sync_events SQLite table (written here, pruned on startup). """ import asyncio import json import logging from typing import Dict, Set from datetime import datetime, timedelta, timezone from fastapi import WebSocket logger = logging.getLogger(__name__) # ── Event loop (captured once at startup) ──────────────────────────────────── _main_loop: asyncio.AbstractEventLoop | None = None def init_loop(loop: asyncio.AbstractEventLoop) -> None: global _main_loop _main_loop = loop # ── Connected clients: user_id → set of asyncio.Queue ──────────────────────── _connections: Dict[int, Set[asyncio.Queue]] = {} async def connect(user_id: int) -> asyncio.Queue: q: asyncio.Queue = asyncio.Queue(maxsize=512) if user_id not in _connections: _connections[user_id] = set() _connections[user_id].add(q) return q async def disconnect(user_id: int, q: asyncio.Queue) -> None: if user_id in _connections: _connections[user_id].discard(q) if not _connections[user_id]: del _connections[user_id] # ── Persistence helpers ─────────────────────────────────────────────────────── def _get_db(): from database import SessionLocal return SessionLocal() def _persist_event(event_type: str, data: dict, user_ids: list[int] | None) -> int: """Write event to sync_events table, return the new seq_id.""" from sqlalchemy import text db = _get_db() try: result = db.execute( text( "INSERT INTO sync_events (event_type, payload, target_user_ids, created_at) " "VALUES (:et, :payload, :uids, :now)" ), { "et": event_type, "payload": json.dumps(data), "uids": json.dumps(user_ids) if user_ids is not None else None, "now": datetime.now(timezone.utc).isoformat(), }, ) db.commit() return result.lastrowid finally: db.close() def get_events_since(cursor: int, user_id: int) -> list[dict]: """Return all events with seq_id > cursor that are visible to user_id.""" from sqlalchemy import text db = _get_db() try: rows = db.execute( text( "SELECT id, event_type, payload, target_user_ids FROM sync_events " "WHERE id > :cursor ORDER BY id ASC LIMIT 500" ), {"cursor": cursor}, ).fetchall() finally: db.close() events = [] for row in rows: target_user_ids = json.loads(row[3]) if row[3] else None if target_user_ids is None or user_id in target_user_ids: events.append({ "seq": row[0], "type": row[1], "data": json.loads(row[2]), }) return events def prune_old_events(hours: int = 24) -> int: """Delete events older than `hours`. Called on startup.""" from sqlalchemy import text db = _get_db() try: cutoff = (datetime.now(timezone.utc) - timedelta(hours=hours)).isoformat() result = db.execute( text("DELETE FROM sync_events WHERE created_at < :cutoff"), {"cutoff": cutoff}, ) db.commit() return result.rowcount finally: db.close() # ── Broadcast ───────────────────────────────────────────────────────────────── def broadcast_sync(event_type: str, data: dict, *, user_ids: list[int] | None = None) -> None: """ Called from sync route thread-pool workers. Persists the event to DB (so reconnecting clients can replay it), then schedules an async push to all currently connected sockets. """ try: seq_id = _persist_event(event_type, data, user_ids) except Exception: logger.exception("ws_bus: failed to persist event %s", event_type) seq_id = 0 if _main_loop is None: return _main_loop.call_soon_threadsafe( _main_loop.create_task, _broadcast_live(seq_id, event_type, data, user_ids), ) async def _broadcast_live( seq_id: int, event_type: str, data: dict, user_ids: list[int] | None, ) -> None: frame = json.dumps({"seq": seq_id, "type": event_type, "data": data}) targets = ( {uid: qs for uid, qs in _connections.items() if uid in user_ids} if user_ids is not None else dict(_connections) ) for qs in targets.values(): for q in list(qs): try: q.put_nowait(frame) except asyncio.QueueFull: pass # slow client — drop live frame; they'll replay on reconnect