pykalshi 0.1.0__py3-none-any.whl → 0.2.0__py3-none-any.whl

pykalshi/feed.py

@@ -0,0 +1,592 @@
+"""Real-time data feed via WebSocket.
+
+This module provides streaming market data through Kalshi's WebSocket API.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import itertools
+import json
+import logging
+import threading
+import time
+from typing import Any, Callable, Optional, TYPE_CHECKING, Union
+
+from pydantic import BaseModel, ConfigDict
+
+if TYPE_CHECKING:
+    from .client import KalshiClient
+
+logger = logging.getLogger(__name__)
+
+# WebSocket endpoints
+DEFAULT_WS_BASE = "wss://api.elections.kalshi.com/trade-api/ws/v2"
+DEMO_WS_BASE = "wss://demo-api.kalshi.co/trade-api/ws/v2"
+_WS_SIGN_PATH = "/trade-api/ws/v2"
+
+
+# --- WebSocket Message Models ---
+
+
+class TickerMessage(BaseModel):
+    """Real-time market ticker update.
+
+    Sent when price, volume, or open interest changes for a subscribed market.
+    """
+
+    market_ticker: str
+    price: Optional[int] = None
+    yes_bid: Optional[int] = None
+    yes_ask: Optional[int] = None
+    volume: Optional[int] = None
+    open_interest: Optional[int] = None
+    dollar_volume: Optional[int] = None
+    dollar_open_interest: Optional[int] = None
+    ts: Optional[int] = None
+
+    model_config = ConfigDict(extra="ignore")
+
+
+class OrderbookSnapshotMessage(BaseModel):
+    """Full orderbook state received on initial subscription.
+
+    Contains all current price levels. After this, you'll receive
+    OrderbookDeltaMessage for incremental updates.
+    """
+
+    market_ticker: str
+    yes: Optional[list[tuple[int, int]]] = None  # [(price, quantity), ...]
+    no: Optional[list[tuple[int, int]]] = None
+
+    model_config = ConfigDict(extra="ignore")
+
+
+class OrderbookDeltaMessage(BaseModel):
+    """Incremental orderbook update.
+
+    Represents a change at a single price level. Apply to local orderbook state.
+    """
+
+    market_ticker: str
+    price: int
+    delta: int  # Positive = added, negative = removed
+    side: str  # "yes" or "no"
+
+    model_config = ConfigDict(extra="ignore")
+
+
+class TradeMessage(BaseModel):
+    """Public trade execution.
+
+    Sent when any trade occurs on subscribed markets.
+    """
+
+    market_ticker: Optional[str] = None
+    ticker: Optional[str] = None
+    trade_id: Optional[str] = None
+    count: Optional[int] = None
+    yes_price: Optional[int] = None
+    no_price: Optional[int] = None
+    taker_side: Optional[str] = None
+    ts: Optional[int] = None
+
+    model_config = ConfigDict(extra="ignore")
+
+
+class FillMessage(BaseModel):
+    """User fill notification (private channel).
+
+    Sent when your orders are filled.
+    """
+
+    trade_id: Optional[str] = None
+    ticker: Optional[str] = None
+    order_id: Optional[str] = None
+    side: Optional[str] = None
+    action: Optional[str] = None
+    count: Optional[int] = None
+    yes_price: Optional[int] = None
+    no_price: Optional[int] = None
+    is_taker: Optional[bool] = None
+    ts: Optional[int] = None
+
+    model_config = ConfigDict(extra="ignore")
+
+
+class PositionMessage(BaseModel):
+    """Real-time position update (private channel).
+
+    Sent when your position in a market changes (after fills settle).
+    Includes realized P&L and current exposure.
+    """
+
+    ticker: str
+    position: Optional[int] = None  # Net contracts (positive = yes, negative = no)
+    market_exposure: Optional[int] = None  # Current exposure in cents
+    realized_pnl: Optional[int] = None  # Realized P&L in cents
+    total_traded: Optional[int] = None  # Total contracts traded
+    resting_orders_count: Optional[int] = None  # Open orders count
+    fees_paid: Optional[int] = None  # Fees paid in cents
+    ts: Optional[int] = None
+
+    model_config = ConfigDict(extra="ignore")
+
+
+class MarketLifecycleMessage(BaseModel):
+    """Market lifecycle state change (public channel).
+
+    Sent when a market's status changes (open, closed, settled, etc.).
+    """
+
+    market_ticker: str
+    status: Optional[str] = None
+    result: Optional[str] = None  # Settlement result ("yes" or "no")
+    ts: Optional[int] = None
+
+    model_config = ConfigDict(extra="ignore")
+
+
+class OrderGroupUpdateMessage(BaseModel):
+    """Order group lifecycle update (private channel).
+
+    Sent when an order group's status changes (triggered, canceled, etc.).
+    """
+
+    order_group_id: str
+    status: Optional[str] = None  # "active", "triggered", "canceled"
+    ts: Optional[int] = None
+
+    model_config = ConfigDict(extra="ignore")
+
+
+# Type alias for orderbook messages (handlers receive either type)
+OrderbookMessage = Union[OrderbookSnapshotMessage, OrderbookDeltaMessage]
+
+# Maps message "type" field to model class
+_MESSAGE_MODELS: dict[str, type[BaseModel]] = {
+    "ticker": TickerMessage,
+    "orderbook_snapshot": OrderbookSnapshotMessage,
+    "orderbook_delta": OrderbookDeltaMessage,
+    "trade": TradeMessage,
+    "fill": FillMessage,
+    "market_position": PositionMessage,
+    "market_lifecycle": MarketLifecycleMessage,
+    "order_group_update": OrderGroupUpdateMessage,
+}
+
+# Maps message types to channel name for handler lookup
+_TYPE_TO_CHANNEL: dict[str, str] = {
+    "orderbook_snapshot": "orderbook_delta",
+    "orderbook_delta": "orderbook_delta",
+    "ticker": "ticker",
+    "trade": "trade",
+    "fill": "fill",
+    "market_position": "market_positions",
+    "market_lifecycle": "market_lifecycle",
+    "order_group_update": "order_group_updates",
+}
+
+
+class Feed:
+    """Real-time streaming data feed via WebSocket.
+
+    Provides a clean interface to Kalshi's WebSocket API with automatic
+    reconnection, typed message models, and callback-based handling.
+
+    Usage:
+        feed = client.feed()
+
+        @feed.on("ticker")
+        def handle_ticker(msg: TickerMessage):
+            print(f"{msg.market_ticker}: {msg.yes_bid}/{msg.yes_ask}")
+
+        @feed.on("orderbook_delta")
+        def handle_orderbook(msg: OrderbookMessage):
+            if isinstance(msg, OrderbookSnapshotMessage):
+                # Initialize local orderbook
+                pass
+            else:
+                # Apply delta
+                pass
+
+        feed.subscribe("ticker", market_ticker="KXBTC-26JAN")
+        feed.subscribe("orderbook_delta", market_ticker="KXBTC-26JAN")
+
+        feed.start()  # Runs in background thread
+        # ... do other work ...
+        feed.stop()
+
+        # Or use as context manager:
+        with client.feed() as feed:
+            feed.on("ticker", my_handler)
+            feed.subscribe("ticker", market_ticker="KXBTC-26JAN")
+            time.sleep(60)
+
+    Available channels:
+        - "ticker": Market price/volume updates (public)
+        - "trade": Public trade executions (public)
+        - "orderbook_delta": Orderbook snapshots and deltas (requires auth)
+        - "fill": Your order fills (requires auth, no market filter)
+        - "market_positions": Real-time position updates with P&L (requires auth, no market filter)
+        - "market_lifecycle": Market state changes (public)
+        - "order_group_updates": Order group lifecycle changes (requires auth)
+    """
+
+    def __init__(self, client: KalshiClient) -> None:
+        """Initialize the feed.
+
+        Args:
+            client: Authenticated KalshiClient instance.
+        """
+        self._client = client
+        self._handlers: dict[str, list[Callable]] = {}
+        self._active_subs: list[dict] = []
+        self._ws: Any = None
+        self._loop: Optional[asyncio.AbstractEventLoop] = None
+        self._thread: Optional[threading.Thread] = None
+        self._running = False
+        self._cmd_id_counter = itertools.count(1)  # Thread-safe counter
+        self._connected = threading.Event()
+        self._lock = threading.Lock()
+        self._metrics_lock = threading.Lock()
+
+        # Latency and health tracking (protected by _metrics_lock)
+        self._connected_at: Optional[float] = None
+        self._last_message_at: Optional[float] = None
+        self._last_server_ts: Optional[int] = None  # Server timestamp in ms
+        self._message_count: int = 0
+        self._reconnect_count: int = 0
+
+        # Determine WS URL from client's API base
+        self._ws_url = DEMO_WS_BASE if "demo" in client.api_base else DEFAULT_WS_BASE
+
+    def on(
+        self, channel: str, handler: Optional[Callable] = None
+    ) -> Callable:
+        """Register a handler for a channel.
+
+        Can be used as a decorator or called directly:
+
+            @feed.on("ticker")
+            def handle(msg: TickerMessage):
+                ...
+
+            # or
+            feed.on("ticker", my_handler)
+
+        Args:
+            channel: Channel name ("ticker", "orderbook_delta", "trade", "fill", "market_positions").
+            handler: Optional handler function. If None, returns a decorator.
+
+        Returns:
+            The handler function (for decorator chaining).
+        """
+        if handler is not None:
+            self._handlers.setdefault(channel, []).append(handler)
+            return handler
+
+        def decorator(fn: Callable) -> Callable:
+            self._handlers.setdefault(channel, []).append(fn)
+            return fn
+
+        return decorator
+
+    def subscribe(
+        self,
+        channel: str,
+        *,
+        market_ticker: Optional[str] = None,
+        market_tickers: Optional[list[str]] = None,
+    ) -> None:
+        """Subscribe to a channel.
+
+        Args:
+            channel: Channel name ("ticker", "orderbook_delta", "trade", "fill", "market_positions").
+            market_ticker: Filter to a single market.
+            market_tickers: Filter to multiple markets.
+
+        Note:
+            - For "fill" and "market_positions" channels, market filters are ignored
+              (you get all your fills/positions).
+            - Can be called before or after start(). If called after, subscription
+              is sent immediately.
+        """
+        params: dict[str, Any] = {"channels": [channel]}
+        if market_ticker is not None:
+            params["market_ticker"] = market_ticker
+        if market_tickers is not None:
+            params["market_tickers"] = market_tickers
+
+        with self._lock:
+            self._active_subs.append(params)
+
+        # Send immediately if connected
+        if self._loop and self._connected.is_set():
+            asyncio.run_coroutine_threadsafe(
+                self._send_cmd("subscribe", params), self._loop
+            )
+
+    def unsubscribe(
+        self,
+        channel: str,
+        *,
+        market_ticker: Optional[str] = None,
+    ) -> None:
+        """Unsubscribe from a channel.
+
+        Args:
+            channel: Channel name.
+            market_ticker: Market to unsubscribe from (must match subscribe call).
+        """
+        params: dict[str, Any] = {"channels": [channel]}
+        if market_ticker is not None:
+            params["market_ticker"] = market_ticker
+
+        # Remove from active subs
+        with self._lock:
+            self._active_subs = [
+                s
+                for s in self._active_subs
+                if not (
+                    s.get("channels") == [channel]
+                    and s.get("market_ticker") == market_ticker
+                )
+            ]
+
+        if self._loop and self._connected.is_set():
+            asyncio.run_coroutine_threadsafe(
+                self._send_cmd("unsubscribe", params), self._loop
+            )
+
+    def start(self) -> None:
+        """Start the feed in a background thread.
+
+        Blocks briefly (up to 10s) until the initial connection is established.
+        If connection fails, the feed continues retrying in the background.
+        """
+        with self._lock:
+            if self._running:
+                return
+            self._running = True
+            self._connected.clear()
+            self._thread = threading.Thread(
+                target=self._run, name="kalshi-feed", daemon=True
+            )
+            self._thread.start()
+        self._connected.wait(timeout=10)
+
+    def stop(self) -> None:
+        """Stop the feed and disconnect."""
+        with self._lock:
+            if not self._running:
+                return
+            self._running = False
+        if self._loop and self._loop.is_running():
+            self._loop.call_soon_threadsafe(self._loop.stop)
+        if self._thread:
+            self._thread.join(timeout=5)
+            self._thread = None
+        self._connected.clear()
+        self._connected_at = None
+
+    @property
+    def is_connected(self) -> bool:
+        """Whether the WebSocket is currently connected."""
+        return self._connected.is_set()
+
+    @property
+    def latency_ms(self) -> Optional[float]:
+        """Estimated latency in milliseconds based on last message timestamp.
+
+        Returns None if no messages with timestamps have been received.
+        This measures the difference between the server's timestamp and
+        when we received the message locally. Assumes clocks are synchronized.
+        """
+        with self._metrics_lock:
+            if self._last_server_ts is None or self._last_message_at is None:
+                return None
+            local_ms = self._last_message_at * 1000
+            return local_ms - self._last_server_ts
+
+    @property
+    def messages_received(self) -> int:
+        """Total number of messages received since feed started."""
+        with self._metrics_lock:
+            return self._message_count
+
+    @property
+    def uptime_seconds(self) -> Optional[float]:
+        """Seconds since connection was established. None if not connected."""
+        with self._metrics_lock:
+            if self._connected_at is None or not self.is_connected:
+                return None
+            return time.time() - self._connected_at
+
+    @property
+    def seconds_since_last_message(self) -> Optional[float]:
+        """Seconds since last message was received. None if no messages yet."""
+        with self._metrics_lock:
+            if self._last_message_at is None:
+                return None
+            return time.time() - self._last_message_at
+
+    @property
+    def reconnect_count(self) -> int:
+        """Number of times the feed has reconnected (0 on first connection)."""
+        with self._metrics_lock:
+            return self._reconnect_count
+
+    def _run(self) -> None:
+        """Background thread entry point."""
+        self._loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(self._loop)
+        try:
+            self._loop.run_until_complete(self._connect_loop())
+        except Exception as e:
+            logger.error("Feed loop crashed: %s", e)
+        finally:
+            self._loop.close()
+            self._loop = None
+
+    async def _connect_loop(self) -> None:
+        """Main connection loop with auto-reconnect."""
+        try:
+            import websockets
+        except ImportError:
+            raise ImportError(
+                "websockets is required for Feed. Install with: pip install websockets"
+            )
+
+        backoff = 0.5
+        max_backoff = 30
+
+        while self._running:
+            try:
+                headers = self._auth_headers()
+                async with websockets.connect(
+                    self._ws_url,
+                    additional_headers=headers,
+                    ping_interval=20,
+                    ping_timeout=10,
+                ) as ws:
+                    self._ws = ws
+                    backoff = 0.5  # Reset on successful connect
+
+                    # Track connection time
+                    with self._metrics_lock:
+                        if self._connected_at is not None:
+                            self._reconnect_count += 1
+                        self._connected_at = time.time()
+
+                    # Replay all active subscriptions
+                    with self._lock:
+                        subs = list(self._active_subs)
+                    for params in subs:
+                        await self._send_cmd("subscribe", params)
+
+                    self._connected.set()
+                    logger.info("Feed connected to %s", self._ws_url)
+
+                    async for raw_msg in ws:
+                        self._dispatch(raw_msg)
+
+            except asyncio.CancelledError:
+                break
+            except Exception as e:
+                self._connected.clear()
+                self._ws = None
+                if not self._running:
+                    break
+                logger.warning(
+                    "Feed disconnected (%s), reconnecting in %.1fs",
+                    type(e).__name__,
+                    backoff,
+                )
+                await asyncio.sleep(backoff)
+                backoff = min(backoff * 2, max_backoff)
+
+        self._connected.clear()
+        self._ws = None
+
+    def _auth_headers(self) -> dict[str, str]:
+        """Generate authentication headers for WebSocket handshake."""
+        timestamp, signature = self._client._sign_request("GET", _WS_SIGN_PATH)
+        return {
+            "KALSHI-ACCESS-KEY": self._client.api_key_id,
+            "KALSHI-ACCESS-SIGNATURE": signature,
+            "KALSHI-ACCESS-TIMESTAMP": timestamp,
+        }
+
+    def _next_id(self) -> int:
+        """Get next command ID (thread-safe)."""
+        return next(self._cmd_id_counter)
+
+    async def _send_cmd(self, cmd: str, params: dict) -> None:
+        """Send a command over the WebSocket."""
+        if self._ws:
+            msg = json.dumps({"id": self._next_id(), "cmd": cmd, "params": params})
+            await self._ws.send(msg)
+            logger.debug("Sent %s: %s", cmd, msg)
+
+    def _dispatch(self, raw: str | bytes) -> None:
+        """Parse incoming message and dispatch to handlers."""
+        receive_time = time.time()
+        with self._metrics_lock:
+            self._last_message_at = receive_time
+            self._message_count += 1
+
+        try:
+            data = json.loads(raw)
+        except (json.JSONDecodeError, TypeError):
+            logger.warning("Malformed message: %.200s", raw)
+            return
+
+        msg_type = data.get("type")
+        if not msg_type:
+            return
+
+        # Extract server timestamp if present (in milliseconds)
+        payload = data.get("msg", data)
+        if isinstance(payload, dict):
+            ts = payload.get("ts")
+            if ts is not None:
+                with self._metrics_lock:
+                    self._last_server_ts = ts
+
+        # Resolve channel for handler lookup
+        channel = _TYPE_TO_CHANNEL.get(msg_type, msg_type)
+        handlers = self._handlers.get(channel)
+        if not handlers:
+            return
+
+        # Parse payload into typed model (payload already extracted above)
+        model_cls = _MESSAGE_MODELS.get(msg_type)
+        if model_cls:
+            try:
+                parsed = model_cls.model_validate(payload)
+            except Exception:
+                logger.debug("Failed to parse %s, passing raw dict", msg_type)
+                parsed = payload
+        else:
+            parsed = payload
+
+        for handler in handlers:
+            try:
+                handler(parsed)
+            except Exception:
+                logger.exception("Handler error on channel %s", channel)
+
+    def __enter__(self):
+        self.start()
+        return self
+
+    def __exit__(self, *args):
+        self.stop()
+
+    def __repr__(self) -> str:
+        status = "connected" if self.is_connected else "disconnected"
+        n = len(self._active_subs)
+        latency = self.latency_ms
+        latency_str = f" latency={latency:.1f}ms" if latency is not None else ""
+        return f"<Feed {status} subs={n} msgs={self._message_count}{latency_str}>"