From f5db1a1eca8d8b85da2f66b4241468e4db7a2a95 Mon Sep 17 00:00:00 2001 From: audit-w1c Date: Wed, 20 May 2026 21:23:36 -0500 Subject: [PATCH] fix(ws): orderbook resync; seq reset detection; permanent close codes; error context; clean unsubscribe - _handle_seq_gap drives a real unsubscribe+resubscribe for orderbook_delta AND order_group_updates (was orderbook-only, was just clear-and-pray). All-markets subs supported via per-sid ticker tracking in OrderbookManager (apply_snapshot records sid; remove_by_sid/tickers_for_sid expose teardown). - SequenceTracker: three-way distinction (seq==last drops, seq None: self._event.set() async def put_sentinel(self) -> None: - """Push shutdown sentinel. Causes async iteration to stop.""" + """Push shutdown sentinel. Causes async iteration to stop. + + Idempotent: after the queue is closed (sentinel or error sentinel + already pushed), subsequent calls are no-ops. Without this the + deque's ``maxlen=maxsize+1`` would evict a real buffered item to + make room for a redundant sentinel — losing the last in-flight + message between e.g. an ERROR-overflow ``put_error`` and the + recv-loop's broadcast ``put_sentinel`` fan-out. + """ + if self._closed: + return self._closed = True self._buffer.append(_SENTINEL) self._event.set() + async def put_error(self, exc: BaseException) -> None: + """Push a terminal error sentinel. + + The iterator yields any items already in the buffer first, then + raises ``exc`` when it reaches the sentinel. Subsequent ``put`` / + ``put_error`` / ``put_sentinel`` calls are silently dropped (the + queue is now closed). Used by the recv loop on + ``KalshiBackpressureError`` so consumers see the failure rather + than a silent ``StopAsyncIteration`` (#207). + """ + if self._closed: + return + self._closed = True + self._buffer.append(_ErrorSentinel(exc=exc)) + self._event.set() + async def get(self) -> T: """Get next item, waiting if empty.""" while not self._buffer: @@ -85,6 +130,8 @@ async def get(self) -> T: item = self._buffer.popleft() if item is _SENTINEL: raise StopAsyncIteration + if isinstance(item, _ErrorSentinel): + raise item.exc self._size -= 1 return item # type: ignore[return-value] @@ -103,5 +150,7 @@ async def __anext__(self) -> T: item = self._buffer.popleft() if item is _SENTINEL: raise StopAsyncIteration + if isinstance(item, _ErrorSentinel): + raise item.exc self._size -= 1 return item # type: ignore[return-value] diff --git a/kalshi/ws/channels.py b/kalshi/ws/channels.py index 0ae6db21..1072e9ad 100644 --- a/kalshi/ws/channels.py +++ b/kalshi/ws/channels.py @@ -16,7 +16,9 @@ logger = logging.getLogger("kalshi.ws") -# Keys forwarded from user params to the subscribe command. +# Keys accepted (and forwarded) on the subscribe command across every channel. +# Used both as a fallback allow-set for channels not in ``_CHANNEL_PARAMS`` +# and as the forwarding set in ``to_subscribe_params``. _SUBSCRIBE_FORWARD_KEYS = ( "market_ticker", "market_tickers", @@ -28,6 +30,31 @@ "skip_ticker_ack", ) +# Per-channel allow-set used by ``Subscription.to_subscribe_params`` to reject +# unknown keys at submission (#195). Each set is the union of params the typed +# subscribe helper accepts plus any AsyncAPI-documented extras for that +# channel. Channels not listed here fall back to ``_SUBSCRIBE_FORWARD_KEYS`` +# (the public ``subscribe(channel=...)`` escape hatch for ``root`` and +# ``control_frames`` and any future server-added channel). +_CHANNEL_PARAMS: dict[str, frozenset[str]] = { + "ticker": frozenset({"market_ticker", "market_tickers", "market_id", "market_ids"}), + "trade": frozenset({"market_ticker", "market_tickers", "market_id", "market_ids"}), + "orderbook_delta": frozenset({ + "market_ticker", "market_tickers", "market_id", "market_ids", + "send_initial_snapshot", + }), + "fill": frozenset(), + "market_positions": frozenset(), + "user_orders": frozenset(), + "order_group_updates": frozenset(), + "market_lifecycle_v2": frozenset({ + "market_ticker", "market_tickers", "market_id", "market_ids", + }), + "multivariate": frozenset(), + "multivariate_market_lifecycle": frozenset(), + "communications": frozenset({"shard_factor", "shard_key"}), +} + class Subscription: """A single channel subscription with durable identity.""" @@ -46,7 +73,29 @@ def __init__( self.server_sid: int | None = None def to_subscribe_params(self) -> dict[str, Any]: - """Build the params dict for the subscribe command.""" + """Build the params dict for the subscribe command. + + Validates ``self.params`` against the channel's allow-set (#195). + Unknown keys raise :class:`KalshiSubscriptionError` so typos like + ``tickerz`` or ``send_snapshot`` fail loudly at submit time rather + than being silently dropped on the wire — the dropped key would + previously cause the consumer to subscribe to a much broader + stream than intended, or to never receive the requested snapshot. + """ + allowed = _CHANNEL_PARAMS.get(self.channel) + if allowed is None: + # Unmodeled channel (e.g. the generic ``subscribe(channel=...)`` + # escape hatch for ``root`` / ``control_frames``). Fall back to + # the union of forwarded keys so users can still drive those + # channels without us hardcoding every server-side parameter. + allowed = frozenset(_SUBSCRIBE_FORWARD_KEYS) + for key in self.params: + if key not in allowed: + raise KalshiSubscriptionError( + f"Unknown param key {key!r} for channel {self.channel!r}", + channel=self.channel, + op="subscribe", + ) result: dict[str, Any] = {"channels": [self.channel]} for key in _SUBSCRIBE_FORWARD_KEYS: if key in self.params: @@ -100,7 +149,13 @@ def _get_msg_id(self) -> int: return mid async def _wait_for_response( - self, msg_id: int, timeout: float = 5.0 + self, + msg_id: int, + timeout: float = 5.0, + *, + channel: str | None = None, + client_id: int | None = None, + op: str | None = None, ) -> dict[str, Any]: """Read frames until we get the response matching our command id. @@ -114,13 +169,20 @@ async def _wait_for_response( between ``_sid_to_client.clear()`` and the new sid mapping are silently dropped — a real signal-loss bug under high-volume reconnect bursts. + + ``channel`` / ``client_id`` / ``op`` are populated on the + ``KalshiSubscriptionError`` so callers can write op-specific + recovery handlers without parsing the error string (#213). """ deadline = asyncio.get_running_loop().time() + timeout while True: remaining = deadline - asyncio.get_running_loop().time() if remaining <= 0: raise KalshiSubscriptionError( - f"Timed out waiting for response to command {msg_id}" + f"Timed out waiting for response to command {msg_id}", + channel=channel, + client_id=client_id, + op=op, # type: ignore[arg-type] ) try: raw = await asyncio.wait_for( @@ -213,12 +275,17 @@ async def subscribe( await self._connection.send(cmd) # Read frames until we get our subscribe ack (by matching id) - data = await self._wait_for_response(msg_id) + data = await self._wait_for_response( + msg_id, channel=channel, client_id=client_id, op="subscribe", + ) if data.get("type") == "error": error_msg = data.get("msg", {}) raise KalshiSubscriptionError( str(error_msg.get("msg", "Subscribe failed")), error_code=error_msg.get("code"), + channel=channel, + client_id=client_id, + op="subscribe", ) server_sid = data.get("msg", {}).get("sid") @@ -245,7 +312,9 @@ async def unsubscribe(self, client_id: int) -> None: cmd = {"id": msg_id, "cmd": "unsubscribe", "params": {"sids": [sub.server_sid]}} await self._connection.send(cmd) - await self._wait_for_response(msg_id) + await self._wait_for_response( + msg_id, channel=sub.channel, client_id=client_id, op="unsubscribe", + ) # F-P-08: push sentinel before deleting so any held iterator exits # cleanly via StopAsyncIteration instead of hanging on queue.get(). await sub.queue.put_sentinel() @@ -267,7 +336,12 @@ async def update_subscription( """Add or remove markets from an existing subscription.""" sub = self._subscriptions.get(client_id) if not sub or sub.server_sid is None: - raise KalshiSubscriptionError("Subscription not found or not active") + raise KalshiSubscriptionError( + "Subscription not found or not active", + channel=sub.channel if sub else None, + client_id=client_id, + op="update_subscription", + ) msg_id = self._get_msg_id() params: dict[str, Any] = {"sids": [sub.server_sid], "action": action} @@ -280,7 +354,9 @@ async def update_subscription( cmd = {"id": msg_id, "cmd": "update_subscription", "params": params} await self._connection.send(cmd) - await self._wait_for_response(msg_id) + await self._wait_for_response( + msg_id, channel=sub.channel, client_id=client_id, op="update_subscription", + ) logger.debug("Updated subscription client_id=%d action=%s", client_id, action) async def resubscribe_all(self) -> None: @@ -326,12 +402,18 @@ async def resubscribe_all(self) -> None: cmd = {"id": msg_id, "cmd": "subscribe", "params": params} await self._connection.send(cmd) - data = await self._wait_for_response(msg_id) + data = await self._wait_for_response( + msg_id, channel=sub.channel, + client_id=client_id, op="subscribe", + ) if data.get("type") == "error": error_msg = data.get("msg", {}) raise KalshiSubscriptionError( str(error_msg.get("msg", "Resubscribe failed")), error_code=error_msg.get("code"), + channel=sub.channel, + client_id=client_id, + op="subscribe", ) new_sid = data.get("msg", {}).get("sid") if new_sid is not None: @@ -360,6 +442,107 @@ async def resubscribe_all(self) -> None: finally: self._stashing = False + async def resubscribe_one(self, client_id: int) -> int | None: + """Tear down and re-establish a single subscription. + + Used by the gap-recovery path so a single bad-seq frame on one + sid doesn't require dropping every other live subscription. + Sends an unsubscribe for the current ``server_sid`` and then a + subscribe with the same channel/params. Orderbook subs are + forced to request a fresh snapshot via + ``send_initial_snapshot=True``. + + Returns the new ``server_sid`` (or ``None`` if the server didn't + echo one). The old sid mapping is cleared before the new sid is + installed so stale in-flight frames cannot mis-route through the + old client_id. + + Raises :class:`KalshiSubscriptionError` if either the + unsubscribe or subscribe is rejected; callers are responsible + for surfacing the failure to the consumer (typically via + :meth:`broadcast_error`). The resubscribe-window stash (#176) + is enabled around the unsubscribe+subscribe pair so frames the + server emits between the two acks are captured for later + replay. + """ + sub = self._subscriptions.get(client_id) + if sub is None: + return None + old_sid = sub.server_sid + + # Enable the stash so data frames arriving between unsubscribe-ack + # and subscribe-ack survive the sid handoff (#176). + prev_stashing = self._stashing + self._stashing = True + try: + if old_sid is not None: + # Best-effort unsubscribe. Server may have already torn down + # the sid (the very gap that triggered resubscribe may be + # part of a server-initiated reset); ignore unsubscribe + # failures so the fresh subscribe still has a chance. + try: + msg_id = self._get_msg_id() + cmd = { + "id": msg_id, "cmd": "unsubscribe", + "params": {"sids": [old_sid]}, + } + await self._connection.send(cmd) + await self._wait_for_response( + msg_id, channel=sub.channel, + client_id=client_id, op="unsubscribe", + ) + except KalshiSubscriptionError: + logger.debug( + "Unsubscribe during resubscribe_one failed for " + "client_id=%d sid=%d; continuing with fresh subscribe", + client_id, old_sid, exc_info=True, + ) + self._sid_to_client.pop(old_sid, None) + sub.server_sid = None + + msg_id = self._get_msg_id() + params = sub.to_subscribe_params() + if sub.channel == "orderbook_delta": + params["send_initial_snapshot"] = True + cmd = {"id": msg_id, "cmd": "subscribe", "params": params} + await self._connection.send(cmd) + data = await self._wait_for_response( + msg_id, channel=sub.channel, + client_id=client_id, op="subscribe", + ) + if data.get("type") == "error": + error_msg = data.get("msg", {}) + raise KalshiSubscriptionError( + str(error_msg.get("msg", "Resubscribe failed")), + error_code=error_msg.get("code"), + channel=sub.channel, + client_id=client_id, + op="subscribe", + ) + new_sid = data.get("msg", {}).get("sid") + if isinstance(new_sid, int): + sub.server_sid = new_sid + self._sid_to_client[new_sid] = client_id + return new_sid if isinstance(new_sid, int) else None + finally: + self._stashing = prev_stashing + + async def broadcast_error( + self, client_id: int, exc: BaseException + ) -> None: + """Surface ``exc`` to the subscription's iterator (#207, #189). + + Puts an error sentinel on the queue so any active ``async for`` + on the subscription raises ``exc`` rather than seeing a silent + ``StopAsyncIteration``. Subsequent puts on the queue become + no-ops (the queue is now closed). Safe to call for unknown + ``client_id`` — it becomes a no-op. + """ + sub = self._subscriptions.get(client_id) + if sub is None: + return + await sub.queue.put_error(exc) + def take_stash(self) -> dict[int, collections.deque[str]]: """Return and clear the resubscribe-window stash atomically (#176). diff --git a/kalshi/ws/client.py b/kalshi/ws/client.py index d364f2ef..6f2b8785 100644 --- a/kalshi/ws/client.py +++ b/kalshi/ws/client.py @@ -16,6 +16,8 @@ from kalshi.config import KalshiConfig from kalshi.errors import ( KalshiBackpressureError, + KalshiConnectionError, + KalshiSequenceGapError, KalshiSubscriptionError, ) from kalshi.models.markets import Orderbook @@ -45,6 +47,14 @@ [Callable[..., Awaitable[None]]], Callable[..., Awaitable[None]] ] +# #197: close codes that signal a permanent server-side rejection. Reconnect +# will fail every attempt for these (auth/protocol/policy/payload violations), +# so fast-fail instead of burning the 10-retry budget. App-specific 4xxx +# codes (notably 4001 = auth) are included wholesale per the Kalshi docs. +_PERMANENT_CLOSE_CODES: frozenset[int] = frozenset( + {1002, 1003, 1007, 1008, 1009, 1010} +) | frozenset(range(4000, 5000)) + class KalshiWebSocket: """WebSocket client for real-time Kalshi market data. @@ -102,6 +112,7 @@ async def _start(self) -> None: sub_mgr=self._sub_mgr, on_error=self._on_error, seq_tracker=self._seq_tracker, + orderbook_mgr=self._orderbook_mgr, ) self._running = True @@ -171,7 +182,29 @@ async def _recv_loop(self) -> None: # F-P-04: cancellation while awaiting recv = no frame read, # no data lost. Safe to exit. break - except ConnectionClosed: + except ConnectionClosed as e: + # #197: classify the close code. The server returns codes + # in the 4xxx range for app-specific failures (4001 auth, + # ...) and RFC 6455 reserves 1002/3/7-10 for protocol + # violations. None of these will recover on retry — fail + # fast instead of burning the 10-retry budget under load. + code = None + rcvd = getattr(e, "rcvd", None) + sent = getattr(e, "sent", None) + if rcvd is not None: + code = rcvd.code + if code is None and sent is not None: + code = sent.code + if code in _PERMANENT_CLOSE_CODES: + reason = (rcvd.reason if rcvd is not None else "") or "" + logger.error( + "WS closed with permanent code %d (%s); not reconnecting", + code, reason or "", + ) + await self._broadcast_sentinels() + raise KalshiConnectionError( + f"WebSocket closed with permanent code {code}: {reason}" + ) from e if not self._running: break await self._handle_reconnect() @@ -284,7 +317,9 @@ async def _process_frame(self, raw: str) -> None: pre_validated: BaseModel | None = None if msg_type == "orderbook_snapshot" and self._orderbook_mgr: snapshot = OrderbookSnapshotMessage.model_validate(data) - self._orderbook_mgr.apply_snapshot(snapshot) + # Record per-sid ticker ownership so all-markets subscriptions + # can be torn down on gap recovery / unsubscribe (#189, #206). + self._orderbook_mgr.apply_snapshot(snapshot, sid=snapshot.sid) pre_validated = snapshot elif msg_type == "orderbook_delta" and self._orderbook_mgr: delta = OrderbookDeltaMessage.model_validate(data) @@ -293,14 +328,22 @@ async def _process_frame(self, raw: str) -> None: try: await self._dispatcher.dispatch(data, pre_validated=pre_validated) - except KalshiBackpressureError: + except KalshiBackpressureError as exc: # Dispatch failed -> the consumer never saw this message. Roll # the seq watermark back so the dropped seq is treated as a - # future gap (not silently as already-seen). Re-raise so the - # recv loop's existing handler broadcasts sentinels and tears - # the loop down. + # future gap (not silently as already-seen). Surface the error + # to the affected sub's iterator (#207) so the consumer can + # halt rather than see an indistinguishable-from-close + # ``StopAsyncIteration``. Re-raise so the recv loop's existing + # handler broadcasts sentinels to the OTHER subs and tears the + # loop down — the affected queue's error sentinel was put + # first, so it surfaces before the close sentinel. if tracked and sid is not None and self._seq_tracker: self._seq_tracker.rollback(sid, prev_seq) + if sid is not None and self._sub_mgr is not None: + affected = self._sub_mgr.get_subscription_by_sid(sid) + if affected is not None: + await self._sub_mgr.broadcast_error(affected.client_id, exc) raise async def _handle_reconnect(self) -> None: @@ -378,29 +421,72 @@ async def _drain_resubscribe_stash(self) -> None: ) async def _handle_seq_gap(self, gap: SequenceGap) -> None: - """Handle a sequence gap by logging and triggering resync. - - A single ``orderbook_delta`` subscription can cover multiple tickers - under one sid. The gap envelope doesn't identify which ticker - missed an update, so we clear EVERY ticker on the affected - subscription — clearing only ``tickers[0]`` would leave the other - books silently diverged from server truth. + """Handle a sequence gap or reset by driving a real resubscribe. + + Replaces the previous clear-and-pray behavior (#189). Works for + any sequenced channel — orderbook_delta covers the canonical + case, order_group_updates also reaches this path (#205). + + Steps: + + 1. Tear down per-sid local state (orderbook books seeded by this + sid via :meth:`OrderbookManager.remove_by_sid`; seq watermark + via :meth:`SequenceTracker.reset`). All-markets orderbook + subscriptions are handled correctly because the manager + tracks ticker ownership by sid, not by the caller's params. + 2. Drive :meth:`SubscriptionManager.resubscribe_one`. The server + assigns a new sid; for orderbook subs the resubscribe forces + ``send_initial_snapshot=True`` so the book repopulates. + 3. If the resubscribe fails, surface ``KalshiSequenceGapError`` + to the consumer iterator via + :meth:`SubscriptionManager.broadcast_error` (#207-style) so + safety-critical strategies can halt rather than consume from + a permanently-empty local book. """ logger.warning( - "Sequence gap on sid %d: expected %d, got %d. Triggering resync.", - gap.sid, gap.expected, gap.received, + "Sequence %s on sid %d: expected %d, got %d. Triggering resync.", + gap.kind, gap.sid, gap.expected, gap.received, ) - if self._sub_mgr: - sub = self._sub_mgr.get_subscription_by_sid(gap.sid) - if sub and sub.channel == "orderbook_delta": - # Clear orderbook state for ALL tickers in the sub and reset - # sequence tracking so the next snapshot rebootstraps cleanly. - tickers = sub.params.get("market_tickers", []) - if self._orderbook_mgr: - for ticker in tickers: - self._orderbook_mgr.remove(ticker) - if self._seq_tracker: - self._seq_tracker.reset(gap.sid) + if self._sub_mgr is None: + return + sub = self._sub_mgr.get_subscription_by_sid(gap.sid) + if sub is None: + return + + # 1. Tear down per-sid local state. + if self._orderbook_mgr is not None: + cleared = self._orderbook_mgr.remove_by_sid(gap.sid) + if cleared: + logger.warning( + "Cleared %d orderbook entries for sid %d on %s", + len(cleared), gap.sid, gap.kind, + ) + if self._seq_tracker is not None: + self._seq_tracker.reset(gap.sid) + + client_id = sub.client_id + channel = sub.channel + + # 2. Drive a real resubscribe so the server emits a fresh snapshot + # (and a new sid). Surface failures to the consumer iterator. + try: + await self._sub_mgr.resubscribe_one(client_id) + except KalshiSubscriptionError as e: + logger.error( + "Resubscribe after gap failed for client_id=%d channel=%s", + client_id, channel, exc_info=True, + ) + await self._sub_mgr.broadcast_error( + client_id, + KalshiSequenceGapError( + f"Resubscribe failed for {channel} after {gap.kind}: {e}", + channel=channel, + sid=gap.sid, + client_id=client_id, + last_seq=gap.expected - 1, + next_seq=gap.received, + ), + ) # ------------------------------------------------------------------ # Internal subscribe helper @@ -557,6 +643,40 @@ async def subscribe( channel, params=params, overflow=overflow, maxsize=maxsize, ) + # ------------------------------------------------------------------ + # Unsubscribe + # ------------------------------------------------------------------ + + async def unsubscribe(self, client_id: int) -> None: + """Tear down a subscription and its associated local state (#206). + + Removes any local :class:`OrderbookManager` books seeded by the + sub's current ``server_sid`` (works for both explicit-tickers + and all-markets orderbook subscriptions via the manager's per-sid + index), then delegates to + :meth:`SubscriptionManager.unsubscribe` to send the wire command, + push the consumer-iterator sentinel, and drop the bookkeeping. + Resets the seq watermark too so any sid reuse after server-side + renumber doesn't replay against a stale floor. + """ + if self._sub_mgr is None: + return + sub = self._sub_mgr.get_subscription(client_id) + if sub is None: + return + old_sid = sub.server_sid + if old_sid is not None: + if self._orderbook_mgr is not None: + cleared = self._orderbook_mgr.remove_by_sid(old_sid) + if cleared: + logger.debug( + "Cleared %d orderbook entries for sid %d on unsubscribe", + len(cleared), old_sid, + ) + if self._seq_tracker is not None: + self._seq_tracker.reset(old_sid) + await self._sub_mgr.unsubscribe(client_id) + # ------------------------------------------------------------------ # Callback API # ------------------------------------------------------------------ diff --git a/kalshi/ws/dispatch.py b/kalshi/ws/dispatch.py index 00aff54c..2fc5eaa1 100644 --- a/kalshi/ws/dispatch.py +++ b/kalshi/ws/dispatch.py @@ -19,6 +19,7 @@ from kalshi.ws.models.ticker import TickerMessage from kalshi.ws.models.trade import TradeMessage from kalshi.ws.models.user_orders import UserOrdersMessage +from kalshi.ws.orderbook import OrderbookManager from kalshi.ws.sequence import SequenceTracker logger = logging.getLogger("kalshi.ws") @@ -51,10 +52,12 @@ def __init__( sub_mgr: SubscriptionManager, on_error: Callable[[ErrorMessage], Awaitable[None]] | None = None, seq_tracker: SequenceTracker | None = None, + orderbook_mgr: OrderbookManager | None = None, ) -> None: self._sub_mgr = sub_mgr self._on_error = on_error self._seq_tracker = seq_tracker + self._orderbook_mgr = orderbook_mgr self._callbacks: dict[str, Callable[[Any], Awaitable[None]]] = {} def register_callback( @@ -106,6 +109,12 @@ async def _handle_server_unsubscribe(self, data: dict[str, Any]) -> None: client_id = self._sub_mgr._sid_to_client.pop(sid, None) if self._seq_tracker is not None: self._seq_tracker.reset(sid) + # #206: server-initiated unsubscribe must also tear down local + # orderbook state seeded by this sid; otherwise stale books are + # served by ``OrderbookManager.get`` indefinitely after the + # server reaps the sub. + if self._orderbook_mgr is not None: + self._orderbook_mgr.remove_by_sid(sid) if client_id is None: logger.debug( diff --git a/kalshi/ws/orderbook.py b/kalshi/ws/orderbook.py index 7b186319..95056b26 100644 --- a/kalshi/ws/orderbook.py +++ b/kalshi/ws/orderbook.py @@ -26,11 +26,19 @@ class _BookState: Kept separate from the public :class:`Orderbook` model so the manager can mutate freely without leaking changes to previously-handed-out snapshots. + + ``sid`` records the server-assigned subscription id that produced this + book's most recent snapshot. Used by + :meth:`OrderbookManager.remove_by_sid` so all-markets subscriptions + (no ``market_tickers`` param) can be cleanly torn down on gap recovery + or unsubscribe without the caller having to know which tickers the + server happened to forward through that sid. """ ticker: str yes: dict[Decimal, Decimal] = field(default_factory=dict) no: dict[Decimal, Decimal] = field(default_factory=dict) + sid: int | None = None def to_orderbook(self) -> Orderbook: """Build a fresh ``Orderbook`` snapshot from the current state. @@ -63,6 +71,11 @@ class OrderbookManager: returns a fresh :class:`Orderbook` instance. Consumers may safely retain references; subsequent updates will not mutate previously-returned books. + Per-sid tracking (:meth:`tickers_for_sid`, :meth:`remove_by_sid`) lets + callers tear down all books owned by a subscription without enumerating + tickers — critical for all-markets ``subscribe_orderbook_delta`` where + the server picks which tickers to forward. + Usage: mgr = OrderbookManager() book = mgr.apply_snapshot(snapshot_msg) # Initialize @@ -72,19 +85,51 @@ class OrderbookManager: def __init__(self) -> None: self._books: dict[str, _BookState] = {} - - def apply_snapshot(self, msg: OrderbookSnapshotMessage) -> Orderbook: - """Initialize (or reset) a book from a full snapshot.""" + # sid -> set of tickers seeded by snapshots on that sid. Maintained + # in lockstep with ``_books`` so ``remove_by_sid`` is O(k) in the + # number of tickers owned by the sid, not O(n) over every book. + self._sid_tickers: dict[int, set[str]] = {} + + def apply_snapshot( + self, msg: OrderbookSnapshotMessage, sid: int | None = None + ) -> Orderbook: + """Initialize (or reset) a book from a full snapshot. + + ``sid`` records which subscription produced this snapshot so the + ticker can be torn down later via :meth:`remove_by_sid` — required + for all-markets subscriptions that don't pin tickers up front. + If omitted, defaults to ``msg.sid`` from the envelope. + """ ticker = msg.msg.market_ticker + sid_val = sid if sid is not None else msg.sid yes_levels = dict(msg.msg.yes) no_levels = dict(msg.msg.no) - state = _BookState(ticker=ticker, yes=yes_levels, no=no_levels) + + # If a prior book existed under a different sid (e.g. server moved + # the ticker between subs), unindex it from the old sid first so + # remove_by_sid on the stale id doesn't drop the live book. + prior = self._books.get(ticker) + if prior is not None and prior.sid is not None and prior.sid != sid_val: + old_bucket = self._sid_tickers.get(prior.sid) + if old_bucket is not None: + old_bucket.discard(ticker) + if not old_bucket: + self._sid_tickers.pop(prior.sid, None) + + state = _BookState(ticker=ticker, yes=yes_levels, no=no_levels, sid=sid_val) self._books[ticker] = state + if sid_val is not None: + bucket = self._sid_tickers.get(sid_val) + if bucket is None: + bucket = set() + self._sid_tickers[sid_val] = bucket + bucket.add(ticker) logger.debug( - "Orderbook snapshot: %s (%d yes, %d no levels)", + "Orderbook snapshot: %s (%d yes, %d no levels, sid=%s)", ticker, len(yes_levels), len(no_levels), + sid_val, ) return state.to_orderbook() @@ -135,8 +180,42 @@ def get(self, ticker: str) -> Orderbook | None: def remove(self, ticker: str) -> None: """Remove a book (e.g., on unsubscribe).""" - self._books.pop(ticker, None) + state = self._books.pop(ticker, None) + if state is not None and state.sid is not None: + bucket = self._sid_tickers.get(state.sid) + if bucket is not None: + bucket.discard(ticker) + if not bucket: + self._sid_tickers.pop(state.sid, None) + + def tickers_for_sid(self, sid: int) -> set[str]: + """Return tickers currently tracked under ``sid``. + + Returns a fresh set; mutating the returned value does not affect + the manager's internal state. + """ + bucket = self._sid_tickers.get(sid) + return set(bucket) if bucket is not None else set() + + def remove_by_sid(self, sid: int) -> list[str]: + """Remove every book associated with ``sid``. + + Returns the tickers that were removed (empty if the sid was + unknown). Used by the gap-recovery and unsubscribe paths so a + single call tears down all per-sid local state, including + all-markets subscriptions where the caller didn't know which + tickers the server would forward. + """ + bucket = self._sid_tickers.pop(sid, None) + if not bucket: + return [] + removed: list[str] = [] + for ticker in bucket: + if self._books.pop(ticker, None) is not None: + removed.append(ticker) + return removed def clear(self) -> None: """Remove all books (e.g., on reconnect before resubscribe).""" self._books.clear() + self._sid_tickers.clear() diff --git a/kalshi/ws/sequence.py b/kalshi/ws/sequence.py index 863b297c..4d58c5d9 100644 --- a/kalshi/ws/sequence.py +++ b/kalshi/ws/sequence.py @@ -13,11 +13,19 @@ @dataclass class SequenceGap: - """Represents a detected sequence gap.""" + """Represents a detected sequence gap. + + ``kind`` is ``"gap"`` for the normal forward-gap case (``received > + expected``) and ``"reset"`` for the seq-went-backwards case + (``received <= last``, treated as a server-side reset that preserved + the sid). Both cases require the same recovery action (clear local + state, resubscribe), so they share the ``on_gap`` callback. + """ sid: int expected: int received: int + kind: str = "gap" class SequenceTracker: @@ -39,10 +47,29 @@ def should_track(self, channel: str) -> bool: return channel in SEQUENCED_CHANNELS async def track(self, sid: int, seq: int | None, channel: str) -> bool: - """Track a message's sequence number. Returns True if OK, False if gap detected. - - For non-sequenced channels, always returns True. - For seq=None on a sequenced channel, returns True (first message or snapshot). + """Track a message's sequence number. + + Returns ``True`` if the caller should dispatch the message, + ``False`` if it should be dropped. + + Three-way distinction for sequenced channels (#196): + + 1. ``seq == expected`` — happy path. Advance watermark, dispatch. + 2. ``seq > expected`` — forward gap. Fire ``on_gap`` with + ``kind="gap"``, advance watermark to ``seq``, drop the + message (caller's resync handler will resubscribe). + 3. ``seq == last`` — exact duplicate. **Drop** without + dispatch. Previously dispatched-without-tracking, which + let a re-delivered delta be applied twice to the orderbook. + 4. ``seq < last`` — server-side reset that reused the sid (or + seq=1 after a high watermark). Fire ``on_gap`` with + ``kind="reset"``, rewind watermark to ``seq``, drop the + message so the resync handler can resubscribe from a clean + snapshot. Without this, every subsequent post-reset message + would be treated as another duplicate until the old + watermark was re-crossed, silently corrupting state. + + Non-sequenced channels and ``seq=None`` always return ``True``. """ if not self.should_track(channel) or seq is None: return True @@ -59,13 +86,29 @@ async def track(self, sid: int, seq: int | None, channel: str) -> bool: self._last_seq[sid] = seq return True - if seq <= last: - # Duplicate or old message, skip - logger.debug("Duplicate seq %d for sid %d (last=%d)", seq, sid, last) - return True + if seq == last: + # Exact duplicate. Drop instead of dispatching, so a redelivered + # delta cannot be applied twice to the orderbook. + logger.debug("Duplicate seq %d for sid %d (last=%d); dropping", seq, sid, last) + return False + + if seq < last: + # Server-side reset that preserved the sid (or seq=1 after a + # high watermark). Treat as a gap with kind="reset" so the + # gap handler clears local state and resubscribes from a + # fresh snapshot. Rewind the watermark so subsequent + # post-reset messages aren't silently dropped as "still old". + gap = SequenceGap(sid=sid, expected=expected, received=seq, kind="reset") + logger.warning( + "Sequence reset on sid=%d: last=%d got=%d", sid, last, seq, + ) + self._last_seq[sid] = seq + if self._on_gap is not None: + await self._on_gap(gap) + return False - # Gap detected - gap = SequenceGap(sid=sid, expected=expected, received=seq) + # Forward gap (seq > expected). + gap = SequenceGap(sid=sid, expected=expected, received=seq, kind="gap") logger.warning("Sequence gap: sid=%d expected=%d got=%d", sid, expected, seq) self._last_seq[sid] = seq # Accept the new seq to continue tracking diff --git a/tests/ws/test_channels.py b/tests/ws/test_channels.py index 3857c723..27fb0c25 100644 --- a/tests/ws/test_channels.py +++ b/tests/ws/test_channels.py @@ -53,26 +53,63 @@ def test_to_subscribe_params_no_extra(self) -> None: result = sub.to_subscribe_params() assert result == {"channels": ["fill"]} - def test_to_subscribe_params_all_keys(self) -> None: + def test_to_subscribe_params_orderbook_keys(self) -> None: queue: MessageQueue[object] = MessageQueue(maxsize=10) params = { "market_ticker": "T1", "market_tickers": ["T1", "T2"], "market_id": "id1", "market_ids": ["id1", "id2"], - "shard_factor": 2, - "shard_key": "k", "send_initial_snapshot": True, - "skip_ticker_ack": True, } sub = Subscription(client_id=1, channel="orderbook_delta", params=params, queue=queue) result = sub.to_subscribe_params() assert result["channels"] == ["orderbook_delta"] assert result["market_ticker"] == "T1" assert result["market_tickers"] == ["T1", "T2"] - assert result["shard_factor"] == 2 assert result["send_initial_snapshot"] is True + def test_to_subscribe_params_rejects_unknown_key(self) -> None: + """#195: typos / shard params on a non-sharded channel must fail loudly.""" + queue: MessageQueue[object] = MessageQueue(maxsize=10) + sub = Subscription( + client_id=1, channel="orderbook_delta", + params={"tickerz": ["T1"]}, queue=queue, + ) + with pytest.raises(KalshiSubscriptionError) as excinfo: + sub.to_subscribe_params() + assert excinfo.value.channel == "orderbook_delta" + assert excinfo.value.op == "subscribe" + + def test_to_subscribe_params_rejects_shard_on_orderbook(self) -> None: + """shard_factor/shard_key are communications-only; reject on orderbook_delta.""" + queue: MessageQueue[object] = MessageQueue(maxsize=10) + sub = Subscription( + client_id=1, channel="orderbook_delta", + params={"shard_factor": 2}, queue=queue, + ) + with pytest.raises(KalshiSubscriptionError): + sub.to_subscribe_params() + + def test_to_subscribe_params_communications_shard_keys(self) -> None: + queue: MessageQueue[object] = MessageQueue(maxsize=10) + sub = Subscription( + client_id=1, channel="communications", + params={"shard_factor": 2, "shard_key": 1}, queue=queue, + ) + result = sub.to_subscribe_params() + assert result == {"channels": ["communications"], "shard_factor": 2, "shard_key": 1} + + def test_to_subscribe_params_unmodeled_channel_falls_back(self) -> None: + """Generic subscribe escape hatch: unmodeled channels accept any forward key.""" + queue: MessageQueue[object] = MessageQueue(maxsize=10) + sub = Subscription( + client_id=1, channel="root", + params={"skip_ticker_ack": True}, queue=queue, + ) + result = sub.to_subscribe_params() + assert result == {"channels": ["root"], "skip_ticker_ack": True} + def test_initial_server_sid_is_none(self) -> None: queue: MessageQueue[object] = MessageQueue(maxsize=10) sub = Subscription(client_id=1, channel="ticker", params={}, queue=queue) diff --git a/tests/ws/test_multi_ticker_gap.py b/tests/ws/test_multi_ticker_gap.py index a0a8fe49..0214334f 100644 --- a/tests/ws/test_multi_ticker_gap.py +++ b/tests/ws/test_multi_ticker_gap.py @@ -1,10 +1,13 @@ -"""Regression tests for multi-ticker orderbook_delta gap clear (#79). +"""Regression tests for multi-ticker orderbook_delta gap clear (#79, #189). A single ``orderbook_delta`` subscription can cover multiple tickers under one sid. The gap envelope does not identify which ticker missed an update, so the safe option is to clear EVERY ticker for the affected subscription. Previously only ``tickers[0]`` was cleared, leaving the -other books diverged from server truth. +other books diverged from server truth (#79). Since #189 the handler +also drives a real resubscribe; these tests pause the recv loop before +invoking the handler so its inner ``_wait_for_response`` doesn't race +the live recv task on ``connection.recv()``. Determinism: scenarios drive the gap handler directly via the public API (``session._handle_seq_gap``) so we don't race the recv loop forcing @@ -19,6 +22,7 @@ from kalshi.auth import KalshiAuth from kalshi.config import KalshiConfig +from kalshi.errors import KalshiSubscriptionError from kalshi.ws.backpressure import MessageQueue from kalshi.ws.channels import Subscription from kalshi.ws.client import KalshiWebSocket @@ -28,7 +32,9 @@ @pytest.mark.asyncio class TestMultiTickerGapClear: - """#79: gap on a multi-ticker orderbook_delta sub clears EVERY ticker.""" + """#79/#189: gap on a multi-ticker orderbook_delta sub clears every ticker + AND drives a resubscribe to refetch the snapshot. + """ async def test_gap_clears_all_tickers_in_sub( self, @@ -37,7 +43,7 @@ async def test_gap_clears_all_tickers_in_sub( ) -> None: """Subscribe to tickers=["A","B","C"], seed all three books via snapshots, then fire a gap on the shared sid -> all three books - must be cleared (previously only "A" was cleared). + must be cleared and the subscription resubscribed. """ config = KalshiConfig(ws_base_url=fake_ws.url, timeout=5.0) ws = KalshiWebSocket(auth=test_auth, config=config) @@ -49,21 +55,21 @@ async def test_gap_clears_all_tickers_in_sub( assert session._sub_mgr is not None assert session._orderbook_mgr is not None assert session._seq_tracker is not None - sid = next(iter(session._sub_mgr.active_subscriptions.values())).server_sid - assert sid is not None + old_sid = next( + iter(session._sub_mgr.active_subscriptions.values()) + ).server_sid + assert old_sid is not None # Seed all three books via individual snapshots on the same sid. for i, ticker in enumerate(["A", "B", "C"], start=1): await fake_ws.send_to_all({ - "type": "orderbook_snapshot", "sid": sid, "seq": i, + "type": "orderbook_snapshot", "sid": old_sid, "seq": i, "msg": { "market_ticker": ticker, "market_id": "x", "yes": [["0.50", "100"]], "no": [], }, }) - # Wait for all three to land in the manager via an async-event - # poll (no fixed sleep timing). async def all_seeded() -> bool: return all( session._orderbook_mgr.get(t) is not None # type: ignore[union-attr] @@ -74,35 +80,35 @@ async def all_seeded() -> bool: while not await all_seeded(): await asyncio.sleep(0.01) - # Directly invoke the gap handler — equivalent to the seq - # tracker firing on_gap on its own. Doing it directly removes - # any flakiness around forcing a real gap through the wire. + # Pause the recv loop so resubscribe_one can own connection.recv() + # without racing the background task. In production the gap + # handler fires from inside the recv loop's frame-processing path + # so there is no concurrency issue; the direct-call test must + # simulate that quiescence explicitly. + await session._pause_recv_loop() await session._handle_seq_gap( - SequenceGap(sid=sid, expected=4, received=10) + SequenceGap(sid=old_sid, expected=4, received=10) ) - # #79: every ticker in the sub must be cleared. - assert session._orderbook_mgr.get("A") is None, ( - "#79: A was cleared (already worked pre-fix)" - ) - assert session._orderbook_mgr.get("B") is None, ( - "#79 regression: B's book was not cleared; only tickers[0] " - "was cleared, leaving B diverged from server truth." - ) - assert session._orderbook_mgr.get("C") is None, ( - "#79 regression: C's book was not cleared; only tickers[0] " - "was cleared, leaving C diverged from server truth." - ) - - # Seq tracker is also reset for the sid. - assert session._seq_tracker.peek(sid) is None + # All three books cleared via remove_by_sid. + assert session._orderbook_mgr.get("A") is None + assert session._orderbook_mgr.get("B") is None + assert session._orderbook_mgr.get("C") is None + # Seq tracker is also reset for the old sid. + assert session._seq_tracker.peek(old_sid) is None + # #189: resubscribe issued; subscription has a new sid (fake_ws + # increments sids monotonically). + sub = session._sub_mgr.get_subscription(1) + assert sub is not None + assert sub.server_sid is not None + assert sub.server_sid != old_sid async def test_single_ticker_sub_still_cleared( self, fake_ws, # type: ignore[no-untyped-def] test_auth, # type: ignore[no-untyped-def] ) -> None: - """Sanity check: single-ticker case still clears the (only) ticker.""" + """Single-ticker case still clears the (only) ticker and resubscribes.""" config = KalshiConfig(ws_base_url=fake_ws.url, timeout=5.0) ws = KalshiWebSocket(auth=test_auth, config=config) async with ws.connect() as session: @@ -110,11 +116,13 @@ async def test_single_ticker_sub_still_cleared( assert session._sub_mgr is not None assert session._orderbook_mgr is not None - sid = next(iter(session._sub_mgr.active_subscriptions.values())).server_sid - assert sid is not None + old_sid = next( + iter(session._sub_mgr.active_subscriptions.values()) + ).server_sid + assert old_sid is not None await fake_ws.send_to_all({ - "type": "orderbook_snapshot", "sid": sid, "seq": 1, + "type": "orderbook_snapshot", "sid": old_sid, "seq": 1, "msg": { "market_ticker": "SOLO", "market_id": "x", "yes": [["0.50", "100"]], "no": [], @@ -125,21 +133,23 @@ async def test_single_ticker_sub_still_cleared( while session._orderbook_mgr.get("SOLO") is None: await asyncio.sleep(0.01) + await session._pause_recv_loop() await session._handle_seq_gap( - SequenceGap(sid=sid, expected=2, received=5) + SequenceGap(sid=old_sid, expected=2, received=5) ) assert session._orderbook_mgr.get("SOLO") is None @pytest.mark.asyncio class TestGapHandlerWithoutTickers: - """#79 edge case: an orderbook_delta sub with no market_tickers param + """#79/#189 edge case: an orderbook_delta sub with no market_tickers param (subscribe-to-all) shouldn't crash the gap handler. """ async def test_no_tickers_param_does_not_raise(self) -> None: - # We just need a KalshiWebSocket instance with the managers wired - # up directly — no real connection. + # KalshiWebSocket instance with the managers wired up directly — no + # real connection. Stubbed _sub_mgr exposes only the subset of + # SubscriptionManager methods the gap handler uses. key = rsa.generate_private_key(public_exponent=65537, key_size=2048) auth = KalshiAuth(key_id="test", private_key=key) config = KalshiConfig(ws_base_url="ws://localhost:1") @@ -148,7 +158,6 @@ async def test_no_tickers_param_does_not_raise(self) -> None: ws._orderbook_mgr = OrderbookManager() ws._seq_tracker = SequenceTracker() - # Hand-build a SubscriptionManager-shaped object with one sub. class _StubMgr: def __init__(self) -> None: queue: MessageQueue[object] = MessageQueue() @@ -158,11 +167,81 @@ def __init__(self) -> None: queue=queue, ) self._sub.server_sid = 42 + self.resubscribe_called = False + self.broadcast_called = False def get_subscription_by_sid(self, sid: int) -> Subscription | None: return self._sub if sid == 42 else None - ws._sub_mgr = _StubMgr() # type: ignore[assignment] + def get_subscription(self, client_id: int) -> Subscription | None: + return self._sub if client_id == self._sub.client_id else None + + async def resubscribe_one(self, client_id: int) -> int | None: + self.resubscribe_called = True + # Simulate a successful resubscribe with a fresh sid. + self._sub.server_sid = 99 + return 99 + + async def broadcast_error( + self, client_id: int, exc: BaseException + ) -> None: + self.broadcast_called = True + + stub = _StubMgr() + ws._sub_mgr = stub # type: ignore[assignment] - # Should not raise — no tickers to iterate, just reset seq. await ws._handle_seq_gap(SequenceGap(sid=42, expected=1, received=5)) + assert stub.resubscribe_called + assert not stub.broadcast_called + + async def test_resubscribe_failure_broadcasts_error(self) -> None: + """If resubscribe_one raises KalshiSubscriptionError, the gap handler + surfaces it to the consumer iterator via broadcast_error (#189/#207).""" + key = rsa.generate_private_key(public_exponent=65537, key_size=2048) + auth = KalshiAuth(key_id="test", private_key=key) + config = KalshiConfig(ws_base_url="ws://localhost:1") + ws = KalshiWebSocket(auth=auth, config=config) + ws._orderbook_mgr = OrderbookManager() + ws._seq_tracker = SequenceTracker() + + class _FailingMgr: + def __init__(self) -> None: + queue: MessageQueue[object] = MessageQueue() + self._sub = Subscription( + client_id=7, channel="orderbook_delta", + params={}, queue=queue, + ) + self._sub.server_sid = 11 + self.broadcast_args: tuple[int, BaseException] | None = None + + def get_subscription_by_sid(self, sid: int) -> Subscription | None: + return self._sub if sid == 11 else None + + def get_subscription(self, client_id: int) -> Subscription | None: + return self._sub if client_id == 7 else None + + async def resubscribe_one(self, client_id: int) -> int | None: + raise KalshiSubscriptionError( + "boom", channel="orderbook_delta", + client_id=client_id, op="subscribe", + ) + + async def broadcast_error( + self, client_id: int, exc: BaseException + ) -> None: + self.broadcast_args = (client_id, exc) + + stub = _FailingMgr() + ws._sub_mgr = stub # type: ignore[assignment] + + await ws._handle_seq_gap(SequenceGap(sid=11, expected=5, received=10)) + assert stub.broadcast_args is not None + cid, exc = stub.broadcast_args + assert cid == 7 + from kalshi.errors import KalshiSequenceGapError + assert isinstance(exc, KalshiSequenceGapError) + assert exc.channel == "orderbook_delta" + assert exc.sid == 11 + assert exc.client_id == 7 + assert exc.last_seq == 4 + assert exc.next_seq == 10 diff --git a/tests/ws/test_recv_loop_hardening.py b/tests/ws/test_recv_loop_hardening.py index 4067edbe..e45a1a10 100644 --- a/tests/ws/test_recv_loop_hardening.py +++ b/tests/ws/test_recv_loop_hardening.py @@ -350,17 +350,21 @@ async def test_backpressure_error_breaks_loop_and_sentinels( assert recv_task is not None await asyncio.wait_for(recv_task, timeout=2.0) - # Iterator must terminate. Buffer holds [snapshot, SENTINEL] - # (sentinel is appended unconditionally by put_sentinel). + # #207: iterator yields the snapshot, then raises + # KalshiBackpressureError so the consumer can distinguish data + # loss from a clean shutdown. Previously the error was + # swallowed and the iterator exited with StopAsyncIteration. + from kalshi.errors import KalshiBackpressureError collected: list[object] = [] try: async with asyncio.timeout(3.0): - async for msg in stream: - collected.append(msg) + with pytest.raises(KalshiBackpressureError): + async for msg in stream: + collected.append(msg) except TimeoutError: pytest.fail( - "#83 regression: BackpressureError was swallowed; " - "iterator did not receive sentinel." + "#83 regression: iterator did not terminate after " + "BackpressureError." ) assert len(collected) == 1 # only the snapshot was queued diff --git a/tests/ws/test_seq_correctness.py b/tests/ws/test_seq_correctness.py index f26922d8..6bff3204 100644 --- a/tests/ws/test_seq_correctness.py +++ b/tests/ws/test_seq_correctness.py @@ -15,6 +15,7 @@ import pytest from kalshi.config import KalshiConfig +from kalshi.errors import KalshiBackpressureError from kalshi.ws.client import KalshiWebSocket from kalshi.ws.sequence import SequenceTracker @@ -122,8 +123,11 @@ async def snapshot_landed() -> None: "dropped message; orderbook is now silently desynced." ) - # Drain the iterator so the session can shut down cleanly. + # #207: iterator now raises KalshiBackpressureError after the + # snapshot lands. Previously the iterator exited silently with + # StopAsyncIteration — indistinguishable from a clean shutdown. collected: list[object] = [] - async for msg in stream: - collected.append(msg) + with pytest.raises(KalshiBackpressureError): + async for msg in stream: + collected.append(msg) assert len(collected) == 1 # only the snapshot landed diff --git a/tests/ws/test_sequence.py b/tests/ws/test_sequence.py index 03fc4d97..bc88c17d 100644 --- a/tests/ws/test_sequence.py +++ b/tests/ws/test_sequence.py @@ -28,12 +28,15 @@ async def on_gap(gap: SequenceGap) -> None: assert gaps[0].expected == 2 assert gaps[0].received == 4 - async def test_duplicate_seq_accepted(self) -> None: + async def test_duplicate_seq_dropped(self) -> None: + """#196: exact-duplicate seq is dropped (return False), not dispatched.""" tracker = SequenceTracker() await tracker.track(1, 1, "orderbook_delta") await tracker.track(1, 2, "orderbook_delta") result = await tracker.track(1, 2, "orderbook_delta") # duplicate - assert result is True # duplicates are OK + assert result is False + # Watermark unchanged. + assert tracker.peek(1) == 2 async def test_non_sequenced_channel_always_ok(self) -> None: tracker = SequenceTracker() diff --git a/tests/ws/test_w1c_correctness.py b/tests/ws/test_w1c_correctness.py new file mode 100644 index 00000000..3111e313 --- /dev/null +++ b/tests/ws/test_w1c_correctness.py @@ -0,0 +1,652 @@ +"""Regression tests for the W1-C correctness sweep. + +Covers: +- #189 orderbook resync after gap (CRITICAL) +- #195 generic subscribe param validation +- #196 seq reset detection (three-way) +- #197 reconnect on permanent close codes +- #205 order_group_updates gap recovery +- #206 unsubscribe orderbook leak +- #207 backpressure ERROR strategy raises through iterator +- #213 raise-site kwarg population + +Determinism: where possible, drive state directly via public hooks and +poll on observable conditions. Avoid arbitrary sleeps that would race +the recv loop. +""" +from __future__ import annotations + +import asyncio + +import pytest + +from kalshi.config import KalshiConfig +from kalshi.errors import ( + KalshiBackpressureError, + KalshiConnectionError, + KalshiSequenceGapError, + KalshiSubscriptionError, +) +from kalshi.ws.backpressure import MessageQueue, OverflowStrategy +from kalshi.ws.channels import SubscriptionManager +from kalshi.ws.client import _PERMANENT_CLOSE_CODES, KalshiWebSocket +from kalshi.ws.connection import ConnectionManager +from kalshi.ws.sequence import SequenceGap, SequenceTracker + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +async def _wait_for(predicate, timeout: float = 2.0) -> None: + async with asyncio.timeout(timeout): + while not predicate(): + await asyncio.sleep(0.01) + + +# --------------------------------------------------------------------------- +# #195 — generic subscribe param validation +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +class TestGenericSubscribeParamValidation: + async def test_generic_subscribe_rejects_unknown_param_keys( + self, fake_ws, test_auth, # type: ignore[no-untyped-def] + ) -> None: + """#195: a typo'd or unknown key fails at submission with context.""" + config = KalshiConfig(ws_base_url=fake_ws.url, timeout=5.0) + ws = KalshiWebSocket(auth=test_auth, config=config) + async with ws.connect(): + with pytest.raises(KalshiSubscriptionError) as excinfo: + await ws.subscribe( + "orderbook_delta", + params={"tickerz": ["X"]}, # typo'd key + ) + assert excinfo.value.channel == "orderbook_delta" + assert excinfo.value.op == "subscribe" + + +# --------------------------------------------------------------------------- +# #196 — seq three-way distinction +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +class TestSeqResetDistinction: + async def test_seq_equal_to_last_dropped_not_dispatched(self) -> None: + gaps: list[SequenceGap] = [] + + async def on_gap(gap: SequenceGap) -> None: + gaps.append(gap) + + tracker = SequenceTracker(on_gap=on_gap) + await tracker.track(1, 5, "orderbook_delta") + # Exact duplicate: drop, do NOT fire on_gap. + result = await tracker.track(1, 5, "orderbook_delta") + assert result is False + assert gaps == [] + assert tracker.peek(1) == 5 + + async def test_seq_below_last_triggers_reset_recovery(self) -> None: + gaps: list[SequenceGap] = [] + + async def on_gap(gap: SequenceGap) -> None: + gaps.append(gap) + + tracker = SequenceTracker(on_gap=on_gap) + await tracker.track(1, 10, "orderbook_delta") + # seq < last: server-side reset. Fire on_gap with kind="reset". + result = await tracker.track(1, 3, "orderbook_delta") + assert result is False + assert len(gaps) == 1 + assert gaps[0].kind == "reset" + assert gaps[0].received == 3 + # Watermark rewound to 3 so subsequent post-reset frames aren't + # silently dropped as "still old". + assert tracker.peek(1) == 3 + + async def test_seq_1_after_high_watermark_treated_as_reset(self) -> None: + gaps: list[SequenceGap] = [] + + async def on_gap(gap: SequenceGap) -> None: + gaps.append(gap) + + tracker = SequenceTracker(on_gap=on_gap) + await tracker.track(1, 42, "orderbook_delta") + result = await tracker.track(1, 1, "orderbook_delta") + assert result is False + assert len(gaps) == 1 + assert gaps[0].kind == "reset" + assert tracker.peek(1) == 1 + + +# --------------------------------------------------------------------------- +# #197 — permanent close code classification +# --------------------------------------------------------------------------- + + +class TestPermanentCloseCodes: + def test_permanent_set_includes_4xxx(self) -> None: + assert 4001 in _PERMANENT_CLOSE_CODES + assert 4999 in _PERMANENT_CLOSE_CODES + assert 5000 not in _PERMANENT_CLOSE_CODES + + def test_permanent_set_includes_rfc6455_codes(self) -> None: + for code in (1002, 1003, 1007, 1008, 1009, 1010): + assert code in _PERMANENT_CLOSE_CODES + + def test_recoverable_codes_excluded(self) -> None: + # 1006 abnormal close, 1011 server error, 1001 going away — all recoverable. + for code in (1001, 1006, 1011): + assert code not in _PERMANENT_CLOSE_CODES + + +@pytest.mark.asyncio +class TestCloseCodeBehavior: + async def test_close_4001_does_not_reconnect_and_raises( + self, fake_ws, test_auth, # type: ignore[no-untyped-def] + ) -> None: + """#197: close code 4001 (auth) raises KalshiConnectionError, no retry.""" + config = KalshiConfig( + ws_base_url=fake_ws.url, timeout=5.0, + retry_base_delay=0.01, retry_max_delay=0.05, ws_max_retries=3, + ) + ws = KalshiWebSocket(auth=test_auth, config=config) + async with ws.connect() as session: + stream = await session.subscribe_ticker(tickers=["T1"]) + # Server closes with 4001. + for sub in list(fake_ws.connections): + await sub.close(code=4001, reason="auth failure") + # Recv task must raise KalshiConnectionError. + recv_task = session._recv_task + assert recv_task is not None + with pytest.raises(KalshiConnectionError) as excinfo: + await asyncio.wait_for(recv_task, timeout=2.0) + assert "4001" in str(excinfo.value) + # Iterator terminated cleanly via the broadcast sentinel. + with pytest.raises(StopAsyncIteration): + await stream.__anext__() + + async def test_close_1008_does_not_reconnect( + self, fake_ws, test_auth, # type: ignore[no-untyped-def] + ) -> None: + """#197: policy violation (1008) raises, doesn't retry.""" + config = KalshiConfig( + ws_base_url=fake_ws.url, timeout=5.0, + retry_base_delay=0.01, retry_max_delay=0.05, ws_max_retries=3, + ) + ws = KalshiWebSocket(auth=test_auth, config=config) + async with ws.connect() as session: + await session.subscribe_ticker(tickers=["T1"]) + for sub in list(fake_ws.connections): + await sub.close(code=1008, reason="policy") + recv_task = session._recv_task + assert recv_task is not None + with pytest.raises(KalshiConnectionError): + await asyncio.wait_for(recv_task, timeout=2.0) + + async def test_close_1006_reconnects( + self, fake_ws, test_auth, # type: ignore[no-untyped-def] + ) -> None: + """#197: abnormal closure (1006) is recoverable — reconnect path runs. + + We assert by checking the connection state ends in STREAMING after + the reconnect completes; if 1006 had been classified permanent the + recv loop would raise KalshiConnectionError instead. + """ + config = KalshiConfig( + ws_base_url=fake_ws.url, timeout=5.0, + retry_base_delay=0.01, retry_max_delay=0.05, ws_max_retries=3, + ) + ws = KalshiWebSocket(auth=test_auth, config=config) + async with ws.connect() as session: + await session.subscribe_ticker(tickers=["T1"]) + # Drop the underlying socket without sending a close frame — + # websockets surfaces this as code 1006. + for sub in list(fake_ws.connections): + # transport-level abrupt close + sub.transport.close() + # Give the recv loop a chance to reconnect. + from kalshi.ws.connection import ConnectionState + await _wait_for( + lambda: ws._connection is not None + and ws._connection.state in ( + ConnectionState.STREAMING, + ConnectionState.CONNECTED, + ), + timeout=3.0, + ) + assert ws._connection is not None + assert ws._connection.state in ( + ConnectionState.STREAMING, ConnectionState.CONNECTED, + ) + + async def test_close_4001_broadcasts_sentinels( + self, fake_ws, test_auth, # type: ignore[no-untyped-def] + ) -> None: + """#197: permanent close fans out sentinels so iterators terminate.""" + config = KalshiConfig(ws_base_url=fake_ws.url, timeout=5.0) + ws = KalshiWebSocket(auth=test_auth, config=config) + async with ws.connect() as session: + stream = await session.subscribe_ticker(tickers=["T1"]) + for sub in list(fake_ws.connections): + await sub.close(code=4001, reason="auth") + # Iterator terminates within bounded time via sentinel. + collected: list[object] = [] + with pytest.raises(StopAsyncIteration): + async with asyncio.timeout(2.0): + while True: + collected.append(await stream.__anext__()) + # Retrieve the recv task's exception so the test doesn't leak + # an unretrieved-task warning at GC time. + recv_task = session._recv_task + assert recv_task is not None + with pytest.raises(KalshiConnectionError): + await asyncio.wait_for(recv_task, timeout=1.0) + + +# --------------------------------------------------------------------------- +# #189 — orderbook resync after gap (CRITICAL) +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +class TestOrderbookResync: + async def test_orderbook_resync_after_gap_repopulates_book( + self, fake_ws, test_auth, # type: ignore[no-untyped-def] + ) -> None: + """#189: forced gap on per-ticker sub -> book is repopulated after resync. + + Subscribe to ["T1"], seed via snapshot, force a gap, verify the + book is repopulated by a fresh snapshot on the new sid. + """ + config = KalshiConfig(ws_base_url=fake_ws.url, timeout=5.0) + ws = KalshiWebSocket(auth=test_auth, config=config) + async with ws.connect() as session: + await session.subscribe_orderbook_delta(tickers=["T1"]) + assert session._sub_mgr is not None + assert session._orderbook_mgr is not None + assert session._seq_tracker is not None + old_sid = next( + iter(session._sub_mgr.active_subscriptions.values()) + ).server_sid + assert old_sid is not None + + # Seed snapshot via the wire so seq tracker sees seq=1. + await fake_ws.send_to_all({ + "type": "orderbook_snapshot", "sid": old_sid, "seq": 1, + "msg": { + "market_ticker": "T1", "market_id": "x", + "yes": [["0.50", "100"]], "no": [], + }, + }) + await _wait_for( + lambda: session._orderbook_mgr.get("T1") is not None # type: ignore[union-attr] + ) + + # Pre-arrange: when the resubscribe creates a new sid, the server + # should send a fresh snapshot. fake_ws doesn't auto-snapshot, so + # we'll send it once we observe the new sid. + # Force a gap by sending seq=N+2 (expected 2, got 5). + await fake_ws.send_to_all({ + "type": "orderbook_delta", "sid": old_sid, "seq": 5, + "msg": { + "market_ticker": "T1", "market_id": "x", + "price": "0.51", "delta": "10", "side": "yes", + }, + }) + + # Wait for the resubscribe to land: subscription gets a new sid. + def resubscribed() -> bool: + sub = session._sub_mgr.get_subscription(1) # type: ignore[union-attr] + return ( + sub is not None + and sub.server_sid is not None + and sub.server_sid != old_sid + ) + + await _wait_for(resubscribed, timeout=3.0) + + # The book was cleared on gap (remove_by_sid). + assert session._orderbook_mgr.get("T1") is None + + # Send a fresh snapshot on the new sid; book repopulates. + sub = session._sub_mgr.get_subscription(1) + assert sub is not None + new_sid = sub.server_sid + assert new_sid is not None + await fake_ws.send_to_all({ + "type": "orderbook_snapshot", "sid": new_sid, "seq": 1, + "msg": { + "market_ticker": "T1", "market_id": "x", + "yes": [["0.55", "200"]], "no": [], + }, + }) + await _wait_for( + lambda: session._orderbook_mgr.get("T1") is not None # type: ignore[union-attr] + ) + book = session._orderbook_mgr.get("T1") + assert book is not None + assert len(book.yes) == 1 + + async def test_orderbook_resync_works_for_all_markets_sub( + self, fake_ws, test_auth, # type: ignore[no-untyped-def] + ) -> None: + """#189: all-markets sub (no tickers=) — gap clears all per-sid books.""" + config = KalshiConfig(ws_base_url=fake_ws.url, timeout=5.0) + ws = KalshiWebSocket(auth=test_auth, config=config) + async with ws.connect() as session: + await session.subscribe_orderbook_delta() # no tickers + assert session._sub_mgr is not None + assert session._orderbook_mgr is not None + old_sid = next( + iter(session._sub_mgr.active_subscriptions.values()) + ).server_sid + assert old_sid is not None + + # Seed two tickers under the same sid. + for i, ticker in enumerate(["A", "B"], start=1): + await fake_ws.send_to_all({ + "type": "orderbook_snapshot", "sid": old_sid, "seq": i, + "msg": { + "market_ticker": ticker, "market_id": "x", + "yes": [["0.50", "100"]], "no": [], + }, + }) + await _wait_for( + lambda: session._orderbook_mgr.get("A") is not None # type: ignore[union-attr] + and session._orderbook_mgr.get("B") is not None # type: ignore[union-attr] + ) + # Confirm per-sid tracking is populated. + assert session._orderbook_mgr.tickers_for_sid(old_sid) == {"A", "B"} + + # Force a gap. + await fake_ws.send_to_all({ + "type": "orderbook_delta", "sid": old_sid, "seq": 99, + "msg": { + "market_ticker": "A", "market_id": "x", + "price": "0.51", "delta": "10", "side": "yes", + }, + }) + + def resubscribed() -> bool: + sub = session._sub_mgr.get_subscription(1) # type: ignore[union-attr] + return ( + sub is not None + and sub.server_sid is not None + and sub.server_sid != old_sid + ) + + await _wait_for(resubscribed, timeout=3.0) + + # Both books cleared via remove_by_sid even though the gap + # envelope didn't identify any specific ticker. + assert session._orderbook_mgr.get("A") is None + assert session._orderbook_mgr.get("B") is None + + async def test_order_group_updates_gap_triggers_resubscribe( + self, fake_ws, test_auth, # type: ignore[no-untyped-def] + ) -> None: + """#205: order_group_updates gap routes through generic resubscribe.""" + config = KalshiConfig(ws_base_url=fake_ws.url, timeout=5.0) + ws = KalshiWebSocket(auth=test_auth, config=config) + async with ws.connect() as session: + await session.subscribe_order_group() + assert session._sub_mgr is not None + assert session._seq_tracker is not None + old_sid = next( + iter(session._sub_mgr.active_subscriptions.values()) + ).server_sid + assert old_sid is not None + + # Pause recv loop so we can manipulate seq tracker and re-inject + # without racing. + await session._pause_recv_loop() + session._seq_tracker._last_seq[old_sid] = 5 + + # Drive the gap handler directly (the gap-detected path inside + # _process_frame is exercised by sequence-tracker tests). + await session._handle_seq_gap( + SequenceGap(sid=old_sid, expected=6, received=10) + ) + + # Subscription was resubscribed: new server_sid. + sub = session._sub_mgr.get_subscription(1) + assert sub is not None + assert sub.server_sid is not None + assert sub.server_sid != old_sid + + +# --------------------------------------------------------------------------- +# #206 — unsubscribe leak (orderbook teardown) +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +class TestUnsubscribeOrderbookTeardown: + async def test_unsubscribe_removes_orderbook_state_for_subscribed_tickers( + self, fake_ws, test_auth, # type: ignore[no-untyped-def] + ) -> None: + config = KalshiConfig(ws_base_url=fake_ws.url, timeout=5.0) + ws = KalshiWebSocket(auth=test_auth, config=config) + async with ws.connect() as session: + await session.subscribe_orderbook_delta(tickers=["X"]) + assert session._sub_mgr is not None + assert session._orderbook_mgr is not None + sid = next( + iter(session._sub_mgr.active_subscriptions.values()) + ).server_sid + assert sid is not None + await fake_ws.send_to_all({ + "type": "orderbook_snapshot", "sid": sid, "seq": 1, + "msg": { + "market_ticker": "X", "market_id": "x", + "yes": [["0.50", "100"]], "no": [], + }, + }) + await _wait_for( + lambda: session._orderbook_mgr.get("X") is not None # type: ignore[union-attr] + ) + + # Unsubscribe via the public API. + await session._pause_recv_loop() + await session.unsubscribe(client_id=1) + assert session._orderbook_mgr.get("X") is None + + async def test_unsubscribe_works_for_all_markets_sub( + self, fake_ws, test_auth, # type: ignore[no-untyped-def] + ) -> None: + config = KalshiConfig(ws_base_url=fake_ws.url, timeout=5.0) + ws = KalshiWebSocket(auth=test_auth, config=config) + async with ws.connect() as session: + await session.subscribe_orderbook_delta() + assert session._sub_mgr is not None + assert session._orderbook_mgr is not None + sid = next( + iter(session._sub_mgr.active_subscriptions.values()) + ).server_sid + assert sid is not None + for i, ticker in enumerate(["P", "Q", "R"], start=1): + await fake_ws.send_to_all({ + "type": "orderbook_snapshot", "sid": sid, "seq": i, + "msg": { + "market_ticker": ticker, "market_id": "x", + "yes": [["0.5", "1"]], "no": [], + }, + }) + await _wait_for( + lambda: all( + session._orderbook_mgr.get(t) is not None # type: ignore[union-attr] + for t in ("P", "Q", "R") + ) + ) + + await session._pause_recv_loop() + await session.unsubscribe(client_id=1) + for t in ("P", "Q", "R"): + assert session._orderbook_mgr.get(t) is None + + async def test_get_after_unsubscribe_returns_none_not_stale_book( + self, fake_ws, test_auth, # type: ignore[no-untyped-def] + ) -> None: + config = KalshiConfig(ws_base_url=fake_ws.url, timeout=5.0) + ws = KalshiWebSocket(auth=test_auth, config=config) + async with ws.connect() as session: + await session.subscribe_orderbook_delta(tickers=["Y"]) + assert session._sub_mgr is not None + assert session._orderbook_mgr is not None + sid = next( + iter(session._sub_mgr.active_subscriptions.values()) + ).server_sid + assert sid is not None + await fake_ws.send_to_all({ + "type": "orderbook_snapshot", "sid": sid, "seq": 1, + "msg": { + "market_ticker": "Y", "market_id": "x", + "yes": [["0.50", "100"]], "no": [], + }, + }) + await _wait_for( + lambda: session._orderbook_mgr.get("Y") is not None # type: ignore[union-attr] + ) + await session._pause_recv_loop() + await session.unsubscribe(client_id=1) + assert session._orderbook_mgr.get("Y") is None + + +# --------------------------------------------------------------------------- +# #207 — backpressure ERROR raises through iterator +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +class TestBackpressureErrorThroughIterator: + async def test_error_strategy_overflow_raises_kalshi_backpressure_error( + self, fake_ws, test_auth, # type: ignore[no-untyped-def] + ) -> None: + """#207: iterator raises KalshiBackpressureError, not silent close.""" + config = KalshiConfig(ws_base_url=fake_ws.url, timeout=5.0) + ws = KalshiWebSocket(auth=test_auth, config=config) + async with ws.connect() as session: + stream = await session.subscribe_orderbook_delta( + tickers=["T1"], maxsize=1, + ) + assert session._sub_mgr is not None + sid = next( + iter(session._sub_mgr.active_subscriptions.values()) + ).server_sid + assert sid is not None + + await fake_ws.send_to_all({ + "type": "orderbook_snapshot", "sid": sid, "seq": 1, + "msg": { + "market_ticker": "T1", "market_id": "x", + "yes": [["0.50", "100"]], "no": [], + }, + }) + await asyncio.sleep(0.1) + await fake_ws.send_to_all({ + "type": "orderbook_delta", "sid": sid, "seq": 2, + "msg": { + "market_ticker": "T1", "market_id": "x", + "price": "0.51", "delta": "10", "side": "yes", + }, + }) + + recv_task = session._recv_task + assert recv_task is not None + await asyncio.wait_for(recv_task, timeout=2.0) + + collected: list[object] = [] + with pytest.raises(KalshiBackpressureError): + async for msg in stream: + collected.append(msg) + assert len(collected) == 1 + + async def test_clean_close_still_raises_stop_async_iteration(self) -> None: + """Regression: a clean shutdown via put_sentinel still yields + StopAsyncIteration (not BackpressureError).""" + q: MessageQueue[str] = MessageQueue(maxsize=4, overflow=OverflowStrategy.ERROR) + await q.put("a") + await q.put_sentinel() + collected: list[str] = [] + async for msg in q: + collected.append(msg) + assert collected == ["a"] + + async def test_put_error_is_idempotent_after_close(self) -> None: + """Second put_error is a no-op; iterator still raises the first error.""" + q: MessageQueue[int] = MessageQueue(maxsize=4) + await q.put(1) + first = KalshiBackpressureError("first") + second = KalshiBackpressureError("second") + await q.put_error(first) + await q.put_error(second) # ignored + await q.put_sentinel() # also ignored (already closed) + collected: list[int] = [] + raised: BaseException | None = None + try: + async for msg in q: + collected.append(msg) + except KalshiBackpressureError as e: + raised = e + assert collected == [1] + assert raised is first # not `second` + + +# --------------------------------------------------------------------------- +# #213 — raise-site kwarg population +# --------------------------------------------------------------------------- + + +class TestErrorContextKwargs: + def test_seq_gap_error_includes_channel_sid_seq_kwargs(self) -> None: + err = KalshiSequenceGapError( + "boom", channel="orderbook_delta", + sid=42, client_id=7, last_seq=5, next_seq=10, + ) + assert err.channel == "orderbook_delta" + assert err.sid == 42 + assert err.client_id == 7 + assert err.last_seq == 5 + assert err.next_seq == 10 + + def test_subscription_error_includes_channel_op_kwargs(self) -> None: + err = KalshiSubscriptionError( + "boom", channel="ticker", client_id=3, op="subscribe", + ) + assert err.channel == "ticker" + assert err.client_id == 3 + assert err.op == "subscribe" + + +@pytest.mark.asyncio +class TestRaiseSiteKwargs: + async def test_subscribe_error_populates_kwargs( + self, fake_ws, test_auth, # type: ignore[no-untyped-def] + ) -> None: + """#213: subscribe-ack error carries channel/op context.""" + fake_ws._force_error = True + config = KalshiConfig(ws_base_url=fake_ws.url, timeout=5.0) + ws = KalshiWebSocket(auth=test_auth, config=config) + async with ws.connect(): + with pytest.raises(KalshiSubscriptionError) as excinfo: + await ws.subscribe_ticker(tickers=["T1"]) + assert excinfo.value.channel == "ticker" + assert excinfo.value.op == "subscribe" + assert excinfo.value.client_id is not None + + async def test_update_subscription_error_populates_kwargs( + self, test_auth, # type: ignore[no-untyped-def] + ) -> None: + """update_subscription on missing client_id raises with op kwarg.""" + # Stand up a connection manager but feed update_subscription a bad id. + # We don't need the fake server for this -- only the lookup miss path. + config = KalshiConfig(ws_base_url="ws://localhost:1") + conn = ConnectionManager(auth=test_auth, config=config) + mgr = SubscriptionManager(conn) + with pytest.raises(KalshiSubscriptionError) as excinfo: + await mgr.update_subscription(99, action="add_markets") + assert excinfo.value.op == "update_subscription" + assert excinfo.value.client_id == 99