spb/crypto-trade-bot

1"""Kraken WebSocket v2 public market-data client (wss://ws.kraken.com/v2).23Channels: ticker, book (depth 10/25/100/500/1000, CRC32 checksum over top 10), ohlc (interval minutes),4trade. Heartbeats arrive ~1/s when idle; we treat >`stale_after` seconds of silence as a dead link and5reconnect with exponential backoff. Every message is handed to `on_event(channel, type, data, received_at)`6with the exchange timestamp preserved (event-time semantics). Book state is maintained locally and7checksum-verified; a mismatch triggers resubscription (never trust a drifting book).8"""910from __future__ import annotations1112import asyncio13import json14import random15import time16from collections.abc import Awaitable, Callable17from decimal import Decimal18from typing import Any1920import websockets2122from ktrade.kraken.checksum import book_checksum, fmt_fixed23from ktrade.logging import get_logger2425log = get_logger("kraken.ws")2627EventHandler = Callable[[str, str, list[dict[str, Any]], float], Awaitable[None]]282930class LocalBook:31    def __init__(self, symbol: str, depth: int, price_decimals: int, qty_decimals: int) -> None:32        self.symbol = symbol33        self.depth = depth34        self.pd = price_decimals35        self.qd = qty_decimals36        self.bids: dict[Decimal, Decimal] = {}37        self.asks: dict[Decimal, Decimal] = {}38        self.ts: str | None = None39        self.valid = False4041    def apply(self, msg_type: str, d: dict[str, Any]) -> bool:42        if msg_type == "snapshot":43            self.bids = {Decimal(str(x["price"])): Decimal(str(x["qty"])) for x in d.get("bids", [])}44            self.asks = {Decimal(str(x["price"])): Decimal(str(x["qty"])) for x in d.get("asks", [])}45        else:46            for side, levels in (("bids", self.bids), ("asks", self.asks)):47                for x in d.get(side, []):48                    p, q = Decimal(str(x["price"])), Decimal(str(x["qty"]))49                    if q == 0:50                        levels.pop(p, None)51                    else:52                        levels[p] = q53            # truncate to subscribed depth after processing all updates54            self.bids = dict(sorted(self.bids.items(), key=lambda kv: kv[0], reverse=True)[: self.depth])55            self.asks = dict(sorted(self.asks.items(), key=lambda kv: kv[0])[: self.depth])56        self.ts = d.get("timestamp")57        expected = d.get("checksum")58        if expected is None:59            self.valid = True60            return True61        got = book_checksum(62            [(fmt_fixed(p, self.pd), fmt_fixed(q, self.qd)) for p, q in self.asks.items()],63            [(fmt_fixed(p, self.pd), fmt_fixed(q, self.qd)) for p, q in self.bids.items()],64        )65        self.valid = got == int(expected)66        return self.valid6768    def top(self, n: int = 10) -> dict[str, Any]:69        bids = sorted(self.bids.items(), key=lambda kv: kv[0], reverse=True)[:n]70        asks = sorted(self.asks.items(), key=lambda kv: kv[0])[:n]71        best_bid = float(bids[0][0]) if bids else None72        best_ask = float(asks[0][0]) if asks else None73        bid_depth = float(sum(q for _, q in bids))74        ask_depth = float(sum(q for _, q in asks))75        imb = (bid_depth - ask_depth) / (bid_depth + ask_depth) if (bid_depth + ask_depth) > 0 else 0.076        return {77            "symbol": self.symbol,78            "bids": [[float(p), float(q)] for p, q in bids],79            "asks": [[float(p), float(q)] for p, q in asks],80            "best_bid": best_bid,81            "best_ask": best_ask,82            "spread": (best_ask - best_bid) if best_bid and best_ask else None,83            "spread_bps": ((best_ask - best_bid) / ((best_ask + best_bid) / 2) * 1e4) if best_bid and best_ask else None,84            "imbalance": imb,85            "timestamp": self.ts,86            "valid": self.valid,87        }888990class KrakenWS:91    def __init__(92        self,93        url: str,94        symbols: list[str],95        on_event: EventHandler,96        *,97        book_depth: int = 25,98        ohlc_intervals: list[int] | None = None,99        precision: dict[str, tuple[int, int]] | None = None,100        stale_after: float = 10.0,101    ) -> None:102        self.url = url103        self.symbols = symbols104        self.on_event = on_event105        self.book_depth = book_depth106        self.ohlc_intervals = ohlc_intervals or [1]107        self.precision = precision or {}108        self.stale_after = stale_after109        self.books: dict[str, LocalBook] = {}110        self._stop = asyncio.Event()111        self.connected = False112        self.last_msg_at = 0.0113        self.reconnects = 0114        self.checksum_failures = 0115        self.messages = 0116117    def stop(self) -> None:118        self._stop.set()119120    def _subscriptions(self) -> list[dict[str, Any]]:121        subs: list[dict[str, Any]] = [122            {"method": "subscribe", "params": {"channel": "ticker", "symbol": self.symbols, "event_trigger": "trades"}, "req_id": 1},123            {"method": "subscribe", "params": {"channel": "book", "symbol": self.symbols, "depth": self.book_depth, "snapshot": True}, "req_id": 2},124            {"method": "subscribe", "params": {"channel": "trade", "symbol": self.symbols, "snapshot": False}, "req_id": 3},125        ]126        for i, iv in enumerate(self.ohlc_intervals):127            subs.append({"method": "subscribe", "params": {"channel": "ohlc", "symbol": self.symbols, "interval": iv, "snapshot": False}, "req_id": 10 + i})128        return subs129130    async def run(self) -> None:131        backoff = 1.0132        while not self._stop.is_set():133            try:134                await self._session()135                backoff = 1.0136            except asyncio.CancelledError:137                raise138            except Exception as e:  # noqa: BLE001139                self.connected = False140                self.reconnects += 1141                log.warning("kraken.ws.disconnected", error=str(e), reconnect_in=round(backoff, 1))142                await asyncio.sleep(backoff + random.random())143                backoff = min(backoff * 2, 60.0)144145    async def _session(self) -> None:146        async with websockets.connect(self.url, ping_interval=20, ping_timeout=20, max_size=8 * 1024 * 1024) as ws:147            self.connected = True148            self.last_msg_at = time.monotonic()149            for sub in self._subscriptions():150                await ws.send(json.dumps(sub))151            log.info("kraken.ws.connected", symbols=self.symbols, depth=self.book_depth)152            for s in self.symbols:153                pd, qd = self.precision.get(s, (8, 8))154                self.books[s] = LocalBook(s, self.book_depth, pd, qd)155            while not self._stop.is_set():156                try:157                    raw = await asyncio.wait_for(ws.recv(), timeout=self.stale_after)158                except TimeoutError:159                    raise ConnectionError(f"no message for {self.stale_after}s (heartbeat missing)") from None160                self.last_msg_at = time.monotonic()161                self.messages += 1162                msg = json.loads(raw)163                received_at = time.time()164                if "method" in msg:  # subscribe ack / error165                    if not msg.get("success", True):166                        log.error("kraken.ws.subscribe_error", msg=msg)167                    continue168                ch = msg.get("channel")169                if ch in (None, "heartbeat", "status"):170                    continue171                mtype = msg.get("type", "update")172                data = msg.get("data", [])173                if ch == "book":174                    for d in data:175                        b = self.books.get(d["symbol"])176                        if b is None:177                            continue178                        if not b.apply(mtype, d):179                            self.checksum_failures += 1180                            log.warning("kraken.ws.book_checksum_mismatch", symbol=d["symbol"])181                            await ws.send(json.dumps({"method": "unsubscribe", "params": {"channel": "book", "symbol": [d["symbol"]], "depth": self.book_depth}}))182                            await ws.send(json.dumps({"method": "subscribe", "params": {"channel": "book", "symbol": [d["symbol"]], "depth": self.book_depth, "snapshot": True}}))183                            continue184                    await self.on_event("book", mtype, [self.books[d["symbol"]].top(self.book_depth) for d in data if d["symbol"] in self.books], received_at)185                else:186                    await self.on_event(ch, mtype, data, received_at)187