fix(ws): orderbook resync; seq reset detection; permanent close codes; error context; clean unsubscribe

kalshi/ws/backpressure.py

@@ -4,6 +4,7 @@
 import asyncio
 import collections
 from collections.abc import AsyncIterator
+from dataclasses import dataclass
 from enum import Enum
 from typing import Generic, TypeVar
 
@@ -14,6 +15,22 @@
 _SENTINEL = object()
 
 
+@dataclass(frozen=True)
+class _ErrorSentinel:
+    """Wraps a terminal exception so the iterator can raise it from ``__anext__``.
+
+    Used by :meth:`MessageQueue.put_error` to surface real, programmatic
+    error signals (e.g. ``KalshiBackpressureError`` from the recv loop)
+    to ``async for`` consumers — previously such errors were swallowed
+    and the iterator exited indistinguishably from a clean shutdown
+    (#207). The exception is buffered like any other item and raised
+    only when the consumer reaches it, preserving any prior in-flight
+    messages that landed before the error fired.
+    """
+
+    exc: BaseException
+
+
 class OverflowStrategy(Enum):
     """What to do when the message queue is full."""
 
@@ -28,7 +45,9 @@ class MessageQueue(Generic[T]):
     """Bounded async queue with configurable overflow behavior.
 
     Implements AsyncIterator so consumers can ``async for msg in queue``.
-    Iteration stops when a sentinel is pushed (graceful shutdown).
+    Iteration stops when a sentinel is pushed (graceful shutdown), and
+    raises the wrapped exception when an error sentinel from
+    :meth:`put_error` is reached (#207).
     """
 
     def __init__(
@@ -71,11 +90,37 @@ async def put(self, item: T) -> 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]

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).