fix(ws): roll back seq on backpressure + reset all tickers on multi-ticker gap

kalshi/ws/client.py

@@ -229,13 +229,25 @@ async def _recv_loop(self) -> None:
                 raise
 
     async def _process_frame(self, raw: str) -> None:
-        """Parse, track seq, update orderbook, and dispatch a single frame."""
+        """Parse, track seq, update orderbook, and dispatch a single frame.
+
+        Seq tracking + orderbook-apply are provisional with respect to the
+        downstream dispatch. If dispatch raises ``KalshiBackpressureError``
+        (ERROR-strategy queue overflow), the seq watermark is rolled back
+        so the dropped message stays visible as a future gap rather than
+        being silently treated as already-seen. Orderbook state isn't
+        rolled back — the recv loop tears down via sentinel broadcast on
+        this error, so the local book becomes orphaned but is never
+        desynced-and-still-consumed.
+        """
         assert self._dispatcher is not None
         data = json.loads(raw)
         sid = data.get("sid")
         seq = data.get("seq")
         msg_type = data.get("type", "")
 
+        prev_seq: int | None = None
+        tracked = False
         if sid is not None and seq is not None and self._seq_tracker:
             channel = ""
             sub = (
@@ -245,13 +257,16 @@ async def _process_frame(self, raw: str) -> None:
             )
             if sub:
                 channel = sub.channel
+            prev_seq = self._seq_tracker.peek(sid)
             ok = await self._seq_tracker.track(
                 sid, seq, msg_type if msg_type else channel
             )
             if not ok:
                 # Gap detected — skip dispatching this message.
                 # The gap handler will trigger resync.
                 return
+            # Only meaningful once we know we'll dispatch (and might roll back).
+            tracked = True
 
         # Check for orderbook messages
         if msg_type == "orderbook_snapshot" and self._orderbook_mgr:
@@ -261,7 +276,17 @@ async def _process_frame(self, raw: str) -> None:
             delta = OrderbookDeltaMessage.model_validate(data)
             self._orderbook_mgr.apply_delta(delta)
 
-        await self._dispatcher.dispatch(raw)
+        try:
+            await self._dispatcher.dispatch(raw)
+        except KalshiBackpressureError:
+            # 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.
+            if tracked and sid is not None and self._seq_tracker:
+                self._seq_tracker.rollback(sid, prev_seq)
+            raise
 
     async def _handle_reconnect(self) -> None:
         """Reconnect after ConnectionClosed and re-establish subscriptions.
@@ -295,18 +320,27 @@ async def _handle_reconnect(self) -> None:
                 self._running = False
 
     async def _handle_seq_gap(self, gap: SequenceGap) -> None:
-        """Handle a sequence gap by logging and triggering resync."""
+        """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.
+        """
         logger.warning(
             "Sequence gap on sid %d: expected %d, got %d. Triggering resync.",
             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 this ticker and reset sequence tracking
+                # 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 tickers and self._orderbook_mgr:
-                    self._orderbook_mgr.remove(tickers[0])
+                if self._orderbook_mgr:
+                    for ticker in tickers:
+                        self._orderbook_mgr.remove(ticker)
                 if self._seq_tracker:
                     self._seq_tracker.reset(gap.sid)
 

kalshi/ws/sequence.py

@@ -74,6 +74,30 @@ async def track(self, sid: int, seq: int | None, channel: str) -> bool:
 
         return False
 
+    def peek(self, sid: int) -> int | None:
+        """Return the current last-seen seq for ``sid``, or None if untracked.
+
+        Capture this before calling :meth:`track` so the watermark can be
+        restored via :meth:`rollback` if downstream dispatch fails — an
+        already-advanced watermark would silently treat the dropped message
+        as already-seen on the next gap check.
+        """
+        return self._last_seq.get(sid)
+
+    def rollback(self, sid: int, prev: int | None) -> None:
+        """Restore the last-seen seq for ``sid`` to ``prev``.
+
+        If ``prev`` is None, the entry is removed entirely (the message was
+        the first one seen for this sid and never landed). This is the
+        compensation for a failed downstream dispatch — pair every successful
+        :meth:`track` whose dispatch may raise with a captured :meth:`peek`
+        and call this on failure.
+        """
+        if prev is None:
+            self._last_seq.pop(sid, None)
+        else:
+            self._last_seq[sid] = prev
+
     def reset(self, sid: int) -> None:
         """Reset tracking for a subscription (after resync/resubscribe)."""
         self._last_seq.pop(sid, None)

pyproject.toml

@@ -81,6 +81,10 @@ dev = [
     "respx>=0.21,<1",
     "ruff>=0.8,<1",
     "mypy>=1.13,<3",
+    # Transitive of mypy 2.1.x. Pin <0.5 because 0.5.0 only ships cp314
+    # wheels (cp39-abi3 dropped), breaking CI on Python 3.12 + 3.13 runners.
+    # Drop once upstream republishes broader wheels.
+    "ast-serialize<0.5",
     "datamodel-code-generator>=0.26",
     "pyyaml>=6",
     "python-dotenv>=1",

tests/ws/test_multi_ticker_gap.py

@@ -0,0 +1,168 @@
+"""Regression tests for multi-ticker orderbook_delta gap clear (#79).
+
+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.
+
+Determinism: scenarios drive the gap handler directly via the public
+API (``session._handle_seq_gap``) so we don't race the recv loop forcing
+a real gap on the wire.
+"""
+from __future__ import annotations
+
+import asyncio
+
+import pytest
+from cryptography.hazmat.primitives.asymmetric import rsa
+
+from kalshi.auth import KalshiAuth
+from kalshi.config import KalshiConfig
+from kalshi.ws.backpressure import MessageQueue
+from kalshi.ws.channels import Subscription
+from kalshi.ws.client import KalshiWebSocket
+from kalshi.ws.orderbook import OrderbookManager
+from kalshi.ws.sequence import SequenceGap, SequenceTracker
+
+
+@pytest.mark.asyncio
+class TestMultiTickerGapClear:
+    """#79: gap on a multi-ticker orderbook_delta sub clears EVERY ticker."""
+
+    async def test_gap_clears_all_tickers_in_sub(
+        self,
+        fake_ws,  # type: ignore[no-untyped-def]
+        test_auth,  # type: ignore[no-untyped-def]
+    ) -> 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).
+        """
+        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=["A", "B", "C"],
+            )
+
+            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
+
+            # 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,
+                    "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]
+                    for t in ("A", "B", "C")
+                )
+
+            async with asyncio.timeout(2.0):
+                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.
+            await session._handle_seq_gap(
+                SequenceGap(sid=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
+
+    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."""
+        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=["SOLO"])
+
+            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": "SOLO", "market_id": "x",
+                    "yes": [["0.50", "100"]], "no": [],
+                },
+            })
+
+            async with asyncio.timeout(2.0):
+                while session._orderbook_mgr.get("SOLO") is None:
+                    await asyncio.sleep(0.01)
+
+            await session._handle_seq_gap(
+                SequenceGap(sid=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
+    (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.
+        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()
+
+        # Hand-build a SubscriptionManager-shaped object with one sub.
+        class _StubMgr:
+            def __init__(self) -> None:
+                queue: MessageQueue[object] = MessageQueue()
+                self._sub = Subscription(
+                    client_id=1, channel="orderbook_delta",
+                    params={},  # no market_tickers
+                    queue=queue,
+                )
+                self._sub.server_sid = 42
+
+            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]
+
+        # Should not raise — no tickers to iterate, just reset seq.
+        await ws._handle_seq_gap(SequenceGap(sid=42, expected=1, received=5))