"""MarketDataService: owns the market WS, maintains a book per token. Subscribes to every YES+NO token of the markets we quote and routes each frame to that token's OrderBook. We do NOT set `custom_feature_enabled` (verified to broaden the feed beyond our assets); resolution is detected via catalog flags. On every book mutation it wakes the owning market's quoter via `on_dirty`, and feeds trade prints to `on_trade` for the flow estimator. Reconnects re-snapshot automatically because the server sends a fresh `book` on (re)subscribe. """ from __future__ import annotations import asyncio import json from collections.abc import Callable from typing import Any import websockets from polymaker.journal import Journal from polymaker.logging import get_logger from polymaker.marketdata.orderbook import BookView, OrderBook from polymaker.marketdata.parse import ( TradePrint, parse_book, parse_last_trade, parse_price_changes, parse_tick_size_change, ) log = get_logger("marketdata.service") DirtyCb = Callable[[str, str], None] # (condition_id, token_id) TradeCb = Callable[[TradePrint], None] class MarketDataService: def __init__( self, url: str = "wss://ws-subscriptions-clob.polymarket.com/ws/market", *, on_dirty: DirtyCb | None = None, on_trade: TradeCb | None = None, journal: Journal | None = None, proxy: str | None = None, ) -> None: self._url = url self._on_dirty = on_dirty or (lambda _c, _t: None) self._on_trade = on_trade or (lambda _tp: None) self._journal = journal self._proxy = proxy self.books: dict[str, OrderBook] = {} self._token_condition: dict[str, str] = {} self._subs: list[str] = [] self._ws: Any = None self._stop = asyncio.Event() # ── subscription management ───────────────────────────────────────── def set_markets(self, markets: list[tuple[str, list[str]]]) -> None: """markets = [(condition_id, [token_ids...])]. Rebuilds the desired set.""" subs: list[str] = [] for cond, tokens in markets: for tok in tokens: self._token_condition[tok] = cond self.books.setdefault(tok, OrderBook()) subs.append(tok) self._subs = subs def view(self, token_id: str) -> BookView: book = self.books.get(token_id) return book.view() if book else _empty_view() def book(self, token_id: str) -> OrderBook | None: return self.books.get(token_id) def last_update_ts(self, token_id: str) -> float: b = self.books.get(token_id) return b.last_update_ts if b else 0.0 # ── run loop ──────────────────────────────────────────────────────── async def run(self) -> None: backoff = 1.0 while not self._stop.is_set(): try: await self._connect_and_listen() backoff = 1.0 except (websockets.ConnectionClosed, OSError) as exc: log.warning("market_ws_dropped", err=str(exc), backoff=backoff) except Exception as exc: # noqa: BLE001 log.error("market_ws_error", err=str(exc)) if self._stop.is_set(): break await asyncio.sleep(backoff) backoff = min(backoff * 2, 30.0) async def _connect_and_listen(self) -> None: if not self._subs: await asyncio.sleep(1.0) return kwargs: dict[str, Any] = {"ping_interval": 5, "ping_timeout": None} if self._proxy: kwargs["proxy"] = self._proxy async with websockets.connect(self._url, **kwargs) as ws: self._ws = ws await ws.send(json.dumps({"assets_ids": self._subs, "type": "market"})) log.info("market_ws_subscribed", n=len(self._subs)) async for raw in ws: self._handle(raw) def stop(self) -> None: self._stop.set() # ── message handling ──────────────────────────────────────────────── def _handle(self, raw: str | bytes) -> None: try: data = json.loads(raw) except (json.JSONDecodeError, TypeError): return for msg in data if isinstance(data, list) else [data]: if not isinstance(msg, dict): continue self._dispatch(msg) def _dispatch(self, msg: dict[str, Any]) -> None: et = msg.get("event_type") if et == "book": self._on_book(msg) elif et == "price_change": self._on_price_change(msg) elif et == "last_trade_price": self._on_last_trade(msg) elif et == "tick_size_change": self._on_tick_change(msg) def _on_book(self, msg: dict[str, Any]) -> None: upd = parse_book(msg) if upd is None or upd.asset_id not in self.books: return book = self.books[upd.asset_id] if upd.tick_size: book.set_tick_size(upd.tick_size) book.apply_snapshot(upd.bids, upd.asks, upd.ts, upd.book_hash) self._journal_write("book", msg, upd.ts) self._wake(upd.asset_id) def _on_price_change(self, msg: dict[str, Any]) -> None: changes = parse_price_changes(msg) touched: set[str] = set() for ch in changes: book = self.books.get(ch.asset_id) if book is None: continue book.apply_delta(ch.side, ch.price, ch.size, ch.ts) touched.add(ch.asset_id) if changes: self._journal_write("price_change", msg, changes[0].ts) for tok in touched: self._wake(tok) def _on_last_trade(self, msg: dict[str, Any]) -> None: tp = parse_last_trade(msg) if tp is None or tp.asset_id not in self.books: return self._journal_write("last_trade_price", msg, tp.ts) self._on_trade(tp) def _on_tick_change(self, msg: dict[str, Any]) -> None: tc = parse_tick_size_change(msg) if tc and tc.asset_id in self.books: self.books[tc.asset_id].set_tick_size(tc.tick_size) log.info("tick_size_change", token=tc.asset_id[:12], tick=tc.tick_size) def _wake(self, token_id: str) -> None: cond = self._token_condition.get(token_id) if cond: self._on_dirty(cond, token_id) def _journal_write(self, kind: str, payload: dict[str, Any], ts: float) -> None: if self._journal is not None: self._journal.write(kind, payload, ts) def _empty_view() -> BookView: return BookView(None, 0.0, None, 0.0, None, None, 0.0, 0.0)