first commit

2026-07-22 18:53:44 +08:00
commit 9259325d8a
60 changed files with 11081 additions and 0 deletions
@@ -0,0 +1,282 @@
+"""Polymarket CLOB market-channel WebSocket subscriber.
+
+Wire format reference (April 2026):
+  - endpoint: wss://ws-subscriptions-clob.polymarket.com/ws/market
+  - subscribe:   {"type":"market","assets_ids":[...],"custom_feature_enabled":true}
+  - snapshot:    event_type="book", asks/bids are lists of {price,size} strings
+  - delta:       event_type="price_change", wrapper="price_changes",
+                 side="BUY"|"SELL", size="0" means remove
+  - keepalive:   send the literal text frame "PING" every ~10s
+  - desync:      no sequence numbers; on hash mismatch or missed heartbeat,
+                 drop state and re-subscribe (server re-pushes snapshot)
+  - sharding:    cap tokens/socket around 100-200; open more sockets as needed
+"""
+from __future__ import annotations
+
+import asyncio
+import itertools
+from datetime import UTC, datetime
+from decimal import Decimal, InvalidOperation
+from typing import Any
+
+import orjson
+import websockets
+from loguru import logger
+from websockets.asyncio.client import ClientConnection
+from websockets.exceptions import ConnectionClosed
+
+from ..book.l2 import BookRegistry, LevelChange, Side
+from ..config import settings
+import socket
+
+DEFAULT_SHARD_SIZE = 100
+PING_INTERVAL_S = 10.0
+RECONNECT_BACKOFF_S = (1.0, 2.0, 4.0, 8.0, 15.0)
+
+
+async def _tunnel_through_proxy(uri: str) -> socket.socket | None:
+    """Open an HTTP CONNECT tunnel to the WS host via the configured proxy.
+
+    Returns a connected socket ready for the WS handshake, or None if no proxy.
+    """
+    if not settings.proxy:
+        return None
+    from urllib.parse import urlparse
+
+    parsed = urlparse(uri)
+    proxy_parsed = urlparse(settings.proxy)
+    host, port = parsed.hostname, parsed.port or (443 if parsed.scheme == "wss" else 80)
+    sock = socket.create_connection((proxy_parsed.hostname, proxy_parsed.port), timeout=10)
+    connect_req = (
+        f"CONNECT {host}:{port} HTTP/1.1\r\n"
+        f"Host: {host}:{port}\r\n"
+        f"Proxy-Connection: keep-alive\r\n\r\n"
+    )
+    sock.sendall(connect_req.encode())
+    resp = b""
+    while b"\r\n\r\n" not in resp:
+        chunk = sock.recv(4096)
+        if not chunk:
+            raise OSError("proxy closed before CONNECT response")
+        resp += chunk
+    status = resp.split(b"\r\n", 1)[0]
+    if b" 200 " not in status:
+        sock.close()
+        raise OSError(f"proxy CONNECT failed: {status!r}")
+    sock.settimeout(None)
+    return sock
+
+
+def _to_decimal(v: object, default: Decimal = Decimal(0)) -> Decimal:
+    try:
+        return Decimal(str(v))
+    except (InvalidOperation, TypeError, ValueError):
+        return default
+
+
+def _parse_ts_ms(v: Any) -> datetime:
+    try:
+        ms = int(v)
+    except (TypeError, ValueError):
+        return datetime.now(UTC)
+    return datetime.fromtimestamp(ms / 1000.0, tz=UTC)
+
+
+def _parse_level_list(raw: Any) -> list[tuple[Decimal, Decimal]]:
+    out: list[tuple[Decimal, Decimal]] = []
+    if not isinstance(raw, list):
+        return out
+    for lvl in raw:
+        if not isinstance(lvl, dict):
+            continue
+        price = _to_decimal(lvl.get("price"))
+        size = _to_decimal(lvl.get("size"))
+        if price <= 0:
+            continue
+        out.append((price, size))
+    return out
+
+
+def _side_from_buy_sell(raw: Any) -> Side | None:
+    if not isinstance(raw, str):
+        return None
+    s = raw.strip().upper()
+    if s == "BUY":
+        return Side.BID
+    if s == "SELL":
+        return Side.ASK
+    return None
+
+
+class MarketChannel:
+    """Maintains a single WS connection for one shard of token ids."""
+
+    def __init__(
+        self,
+        token_ids: list[str],
+        *,
+        registry: BookRegistry,
+        url: str | None = None,
+    ) -> None:
+        if not token_ids:
+            raise ValueError("MarketChannel needs at least one token id")
+        self._token_ids = list(token_ids)
+        self._registry = registry
+        self._url = url or settings.ws_url
+        self._stop = asyncio.Event()
+
+    async def run(self) -> None:
+        backoff_cycle = itertools.cycle(RECONNECT_BACKOFF_S)
+        while not self._stop.is_set():
+            try:
+                await self._connect_and_consume()
+                # Clean exit (e.g. stop requested); break loop.
+                if self._stop.is_set():
+                    return
+                delay = 1.0
+            except (TimeoutError, ConnectionClosed, OSError) as exc:
+                delay = next(backoff_cycle)
+                logger.warning(
+                    "ws disconnect ({}); reconnecting in {}s", exc.__class__.__name__, delay
+                )
+            except Exception as exc:
+                delay = next(backoff_cycle)
+                logger.exception("ws fatal ({}); reconnecting in {}s", exc, delay)
+            await asyncio.sleep(delay)
+
+    def stop(self) -> None:
+        self._stop.set()
+
+    async def _connect_and_consume(self) -> None:
+        sock = await _tunnel_through_proxy(self._url)
+        connect_kwargs = {"max_size": 2**22}
+        if sock is not None:
+            connect_kwargs["sock"] = sock
+        async with websockets.connect(self._url, **connect_kwargs) as ws:
+            await self._subscribe(ws)
+            heartbeat = asyncio.create_task(self._heartbeat(ws))
+            try:
+                async for raw in ws:
+                    if self._stop.is_set():
+                        break
+                    self._handle_message(raw)
+            finally:
+                heartbeat.cancel()
+                try:
+                    await heartbeat
+                except asyncio.CancelledError:
+                    pass
+
+    async def _subscribe(self, ws: ClientConnection) -> None:
+        msg = orjson.dumps(
+            {
+                "type": "market",
+                "assets_ids": self._token_ids,
+                "custom_feature_enabled": True,
+            }
+        )
+        await ws.send(msg)
+        logger.info("ws subscribed to {} tokens", len(self._token_ids))
+
+    async def _heartbeat(self, ws: ClientConnection) -> None:
+        while not self._stop.is_set():
+            await asyncio.sleep(PING_INTERVAL_S)
+            try:
+                await ws.send("PING")
+            except ConnectionClosed:
+                return
+
+    def _handle_message(self, raw: str | bytes) -> None:
+        if isinstance(raw, str):
+            if raw.strip() in ("PONG", "PING"):
+                return
+            payload = orjson.loads(raw)
+        else:
+            # Binary frames shouldn't normally arrive from this channel.
+            try:
+                payload = orjson.loads(raw)
+            except orjson.JSONDecodeError:
+                return
+        if isinstance(payload, list):
+            for item in payload:
+                self._dispatch(item)
+        elif isinstance(payload, dict):
+            self._dispatch(payload)
+
+    def _dispatch(self, msg: dict[str, Any]) -> None:
+        event_type = msg.get("event_type")
+        if event_type == "book":
+            self._apply_book(msg)
+        elif event_type == "price_change":
+            self._apply_price_change(msg)
+        elif event_type in ("tick_size_change", "last_trade_price", "best_bid_ask"):
+            # Not used by the arb math today; logged at debug.
+            logger.debug("ws {}: {}", event_type, msg)
+        elif event_type in ("new_market", "market_resolved"):
+            logger.info("ws {}: {}", event_type, msg)
+        else:
+            logger.debug("ws unknown event_type={}: {}", event_type, msg)
+
+    def _apply_book(self, msg: dict[str, Any]) -> None:
+        asset_id = msg.get("asset_id")
+        if not isinstance(asset_id, str):
+            return
+        bids = _parse_level_list(msg.get("bids"))
+        asks = _parse_level_list(msg.get("asks"))
+        ts = _parse_ts_ms(msg.get("timestamp"))
+        self._registry.apply_snapshot(
+            asset_id, bids=bids, asks=asks, timestamp=ts, book_hash=msg.get("hash")
+        )
+
+    def _apply_price_change(self, msg: dict[str, Any]) -> None:
+        ts = _parse_ts_ms(msg.get("timestamp"))
+        changes_raw = msg.get("price_changes")
+        if not isinstance(changes_raw, list):
+            return
+        by_asset: dict[str, list[LevelChange]] = {}
+        latest_hash: dict[str, str] = {}
+        for c in changes_raw:
+            if not isinstance(c, dict):
+                continue
+            asset_id = c.get("asset_id")
+            side = _side_from_buy_sell(c.get("side"))
+            if not isinstance(asset_id, str) or side is None:
+                continue
+            price = _to_decimal(c.get("price"))
+            size = _to_decimal(c.get("size"))
+            if price <= 0:
+                continue
+            by_asset.setdefault(asset_id, []).append(
+                LevelChange(price=price, size=size, side=side)
+            )
+            h = c.get("hash")
+            if isinstance(h, str):
+                latest_hash[asset_id] = h
+        for asset_id, changes in by_asset.items():
+            self._registry.apply_delta(
+                asset_id,
+                changes,
+                timestamp=ts,
+                book_hash=latest_hash.get(asset_id),
+            )
+
+
+def shard_tokens(token_ids: list[str], shard_size: int = DEFAULT_SHARD_SIZE) -> list[list[str]]:
+    return [token_ids[i : i + shard_size] for i in range(0, len(token_ids), shard_size)]
+
+
+async def run_market_channels(
+    token_ids: list[str],
+    *,
+    registry: BookRegistry,
+    shard_size: int = DEFAULT_SHARD_SIZE,
+) -> None:
+    """Run one MarketChannel per shard concurrently. Returns when all exit."""
+    if not token_ids:
+        return
+    channels = [MarketChannel(shard, registry=registry) for shard in shard_tokens(token_ids, shard_size)]
+    try:
+        await asyncio.gather(*(ch.run() for ch in channels))
+    finally:
+        for ch in channels:
+            ch.stop()