"""Kraken WebSocket v2 public market-data client (wss://ws.kraken.com/v2). Channels: ticker, book (depth 10/25/100/500/1000, CRC32 checksum over top 10), ohlc (interval minutes), trade. Heartbeats arrive ~1/s when idle; we treat >`stale_after` seconds of silence as a dead link and reconnect with exponential backoff. Every message is handed to `on_event(channel, type, data, received_at)` with the exchange timestamp preserved (event-time semantics). Book state is maintained locally and checksum-verified; a mismatch triggers resubscription (never trust a drifting book). """ from __future__ import annotations import asyncio import json import random import time from collections.abc import Awaitable, Callable from decimal import Decimal from typing import Any import websockets from ktrade.kraken.checksum import book_checksum, fmt_fixed from ktrade.logging import get_logger log = get_logger("kraken.ws") EventHandler = Callable[[str, str, list[dict[str, Any]], float], Awaitable[None]] class LocalBook: def __init__(self, symbol: str, depth: int, price_decimals: int, qty_decimals: int) -> None: self.symbol = symbol self.depth = depth self.pd = price_decimals self.qd = qty_decimals self.bids: dict[Decimal, Decimal] = {} self.asks: dict[Decimal, Decimal] = {} self.ts: str | None = None self.valid = False def apply(self, msg_type: str, d: dict[str, Any]) -> bool: if msg_type == "snapshot": self.bids = {Decimal(str(x["price"])): Decimal(str(x["qty"])) for x in d.get("bids", [])} self.asks = {Decimal(str(x["price"])): Decimal(str(x["qty"])) for x in d.get("asks", [])} else: for side, levels in (("bids", self.bids), ("asks", self.asks)): for x in d.get(side, []): p, q = Decimal(str(x["price"])), Decimal(str(x["qty"])) if q == 0: levels.pop(p, None) else: levels[p] = q # truncate to subscribed depth after processing all updates self.bids = dict(sorted(self.bids.items(), key=lambda kv: kv[0], reverse=True)[: self.depth]) self.asks = dict(sorted(self.asks.items(), key=lambda kv: kv[0])[: self.depth]) self.ts = d.get("timestamp") expected = d.get("checksum") if expected is None: self.valid = True return True got = book_checksum( [(fmt_fixed(p, self.pd), fmt_fixed(q, self.qd)) for p, q in self.asks.items()], [(fmt_fixed(p, self.pd), fmt_fixed(q, self.qd)) for p, q in self.bids.items()], ) self.valid = got == int(expected) return self.valid def top(self, n: int = 10) -> dict[str, Any]: bids = sorted(self.bids.items(), key=lambda kv: kv[0], reverse=True)[:n] asks = sorted(self.asks.items(), key=lambda kv: kv[0])[:n] best_bid = float(bids[0][0]) if bids else None best_ask = float(asks[0][0]) if asks else None bid_depth = float(sum(q for _, q in bids)) ask_depth = float(sum(q for _, q in asks)) imb = (bid_depth - ask_depth) / (bid_depth + ask_depth) if (bid_depth + ask_depth) > 0 else 0.0 return { "symbol": self.symbol, "bids": [[float(p), float(q)] for p, q in bids], "asks": [[float(p), float(q)] for p, q in asks], "best_bid": best_bid, "best_ask": best_ask, "spread": (best_ask - best_bid) if best_bid and best_ask else None, "spread_bps": ((best_ask - best_bid) / ((best_ask + best_bid) / 2) * 1e4) if best_bid and best_ask else None, "imbalance": imb, "timestamp": self.ts, "valid": self.valid, } class KrakenWS: def __init__( self, url: str, symbols: list[str], on_event: EventHandler, *, book_depth: int = 25, ohlc_intervals: list[int] | None = None, precision: dict[str, tuple[int, int]] | None = None, stale_after: float = 10.0, ) -> None: self.url = url self.symbols = symbols self.on_event = on_event self.book_depth = book_depth self.ohlc_intervals = ohlc_intervals or [1] self.precision = precision or {} self.stale_after = stale_after self.books: dict[str, LocalBook] = {} self._stop = asyncio.Event() self.connected = False self.last_msg_at = 0.0 self.reconnects = 0 self.checksum_failures = 0 self.messages = 0 def stop(self) -> None: self._stop.set() def _subscriptions(self) -> list[dict[str, Any]]: subs: list[dict[str, Any]] = [ {"method": "subscribe", "params": {"channel": "ticker", "symbol": self.symbols, "event_trigger": "trades"}, "req_id": 1}, {"method": "subscribe", "params": {"channel": "book", "symbol": self.symbols, "depth": self.book_depth, "snapshot": True}, "req_id": 2}, {"method": "subscribe", "params": {"channel": "trade", "symbol": self.symbols, "snapshot": False}, "req_id": 3}, ] for i, iv in enumerate(self.ohlc_intervals): subs.append({"method": "subscribe", "params": {"channel": "ohlc", "symbol": self.symbols, "interval": iv, "snapshot": False}, "req_id": 10 + i}) return subs async def run(self) -> None: backoff = 1.0 while not self._stop.is_set(): try: await self._session() backoff = 1.0 except asyncio.CancelledError: raise except Exception as e: # noqa: BLE001 self.connected = False self.reconnects += 1 log.warning("kraken.ws.disconnected", error=str(e), reconnect_in=round(backoff, 1)) await asyncio.sleep(backoff + random.random()) backoff = min(backoff * 2, 60.0) async def _session(self) -> None: async with websockets.connect(self.url, ping_interval=20, ping_timeout=20, max_size=8 * 1024 * 1024) as ws: self.connected = True self.last_msg_at = time.monotonic() for sub in self._subscriptions(): await ws.send(json.dumps(sub)) log.info("kraken.ws.connected", symbols=self.symbols, depth=self.book_depth) for s in self.symbols: pd, qd = self.precision.get(s, (8, 8)) self.books[s] = LocalBook(s, self.book_depth, pd, qd) while not self._stop.is_set(): try: raw = await asyncio.wait_for(ws.recv(), timeout=self.stale_after) except TimeoutError: raise ConnectionError(f"no message for {self.stale_after}s (heartbeat missing)") from None self.last_msg_at = time.monotonic() self.messages += 1 msg = json.loads(raw) received_at = time.time() if "method" in msg: # subscribe ack / error if not msg.get("success", True): log.error("kraken.ws.subscribe_error", msg=msg) continue ch = msg.get("channel") if ch in (None, "heartbeat", "status"): continue mtype = msg.get("type", "update") data = msg.get("data", []) if ch == "book": for d in data: b = self.books.get(d["symbol"]) if b is None: continue if not b.apply(mtype, d): self.checksum_failures += 1 log.warning("kraken.ws.book_checksum_mismatch", symbol=d["symbol"]) await ws.send(json.dumps({"method": "unsubscribe", "params": {"channel": "book", "symbol": [d["symbol"]], "depth": self.book_depth}})) await ws.send(json.dumps({"method": "subscribe", "params": {"channel": "book", "symbol": [d["symbol"]], "depth": self.book_depth, "snapshot": True}})) continue 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) else: await self.on_event(ch, mtype, data, received_at)