diff --git a/kalshi/ws/client.py b/kalshi/ws/client.py index b499888..4fe3e2c 100644 --- a/kalshi/ws/client.py +++ b/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,6 +257,7 @@ 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 ) @@ -252,6 +265,8 @@ async def _process_frame(self, raw: str) -> None: # 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,7 +320,14 @@ 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, @@ -303,10 +335,12 @@ async def _handle_seq_gap(self, gap: SequenceGap) -> None: 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) diff --git a/kalshi/ws/sequence.py b/kalshi/ws/sequence.py index c97f453..863b297 100644 --- a/kalshi/ws/sequence.py +++ b/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) diff --git a/pyproject.toml b/pyproject.toml index 44aedc9..ec5e2a1 100644 --- a/pyproject.toml +++ b/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", diff --git a/tests/ws/test_multi_ticker_gap.py b/tests/ws/test_multi_ticker_gap.py new file mode 100644 index 0000000..a0a8fe4 --- /dev/null +++ b/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)) diff --git a/tests/ws/test_seq_correctness.py b/tests/ws/test_seq_correctness.py new file mode 100644 index 0000000..f26922d --- /dev/null +++ b/tests/ws/test_seq_correctness.py @@ -0,0 +1,129 @@ +"""Regression tests for WS seq-tracker correctness (#78). + +#78: ERROR-overflow backpressure must not silently advance the seq +watermark past a dropped message. + +Determinism: scenarios drive the seq tracker via the public API so we +don't rely on asyncio.sleep to race the recv loop. For the end-to-end +backpressure test, we await ``session._recv_task`` — the recv loop +completes once dispatch raises BackpressureError, no sleep required. +""" +from __future__ import annotations + +import asyncio + +import pytest + +from kalshi.config import KalshiConfig +from kalshi.ws.client import KalshiWebSocket +from kalshi.ws.sequence import SequenceTracker + + +@pytest.mark.asyncio +class TestSequenceTrackerRollback: + """#78: peek + rollback let callers undo a track() that downstream rejected.""" + + async def test_peek_returns_none_for_unknown_sid(self) -> None: + tracker = SequenceTracker() + assert tracker.peek(42) is None + + async def test_peek_returns_last_seq(self) -> None: + tracker = SequenceTracker() + await tracker.track(1, 5, "orderbook_delta") + assert tracker.peek(1) == 5 + + async def test_rollback_restores_prior_seq(self) -> None: + tracker = SequenceTracker() + await tracker.track(1, 5, "orderbook_delta") + prev = tracker.peek(1) + await tracker.track(1, 6, "orderbook_delta") + assert tracker.peek(1) == 6 + tracker.rollback(1, prev) + assert tracker.peek(1) == 5 + + async def test_rollback_to_none_removes_entry(self) -> None: + tracker = SequenceTracker() + # First message ever; prev is None. + prev = tracker.peek(1) + await tracker.track(1, 1, "orderbook_delta") + assert tracker.peek(1) == 1 + tracker.rollback(1, prev) + assert tracker.peek(1) is None + # And the next track is treated as the first message again. + assert await tracker.track(1, 1, "orderbook_delta") is True + + +@pytest.mark.asyncio +class TestBackpressureDoesNotAdvanceSeq: + """#78 end-to-end: ERROR-overflow on a full queue must not advance _last_seq.""" + + async def test_dropped_delta_keeps_seq_watermark_unchanged( + self, + fake_ws, # type: ignore[no-untyped-def] + test_auth, # type: ignore[no-untyped-def] + ) -> None: + """ERROR-strategy queue at maxsize -> send N+1 messages -> tracker's + last_seq reflects only successfully-landed messages. + + Setup: subscribe with maxsize=1, send snapshot (seq=1) — lands in + queue. Send delta (seq=2) — backpressure raises. The recv loop + exits via sentinel broadcast. Tracker's last_seq for the sid must + still be 1, not 2. + """ + 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 + assert session._seq_tracker is not None + sid = next(iter(session._sub_mgr.active_subscriptions.values())).server_sid + assert sid is not None + + # Snapshot fills the queue (maxsize=1). + 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": [], + }, + }) + # Deterministic wait: poll the seq tracker until it observes the + # snapshot, bounded by wait_for. Avoids CI flakes under load. + async def snapshot_landed() -> None: + while session._seq_tracker.peek(sid) != 1: + await asyncio.sleep(0.01) + await asyncio.wait_for(snapshot_landed(), timeout=2.0) + + assert session._seq_tracker.peek(sid) == 1 + + # This delta overflows -> BackpressureError -> recv loop tears + # down via sentinel broadcast. + 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) + + # #78: the seq watermark MUST still be 1, not 2. If it were 2, + # a post-reconnect message with seq=2 would be treated as a + # duplicate (silent desync); a message with seq=3 would not + # be detected as a gap. + assert session._seq_tracker.peek(sid) == 1, ( + "#78 regression: ERROR-overflow advanced last_seq past a " + "dropped message; orderbook is now silently desynced." + ) + + # Drain the iterator so the session can shut down cleanly. + collected: list[object] = [] + async for msg in stream: + collected.append(msg) + assert len(collected) == 1 # only the snapshot landed