From 2d6afd9c255e9b2d46aad4035f7441ec38c5409b Mon Sep 17 00:00:00 2001 From: Jeff West Date: Wed, 20 May 2026 17:51:15 -0500 Subject: [PATCH 1/5] fix(ws): stash + replay data frames during resubscribe window (#176) Fixes silent message loss when the server sends data frames on a freshly-assigned sid before the corresponding subscribe ack reaches `_wait_for_response`. Previously these frames were non-matching from the wait's perspective and discarded with a debug log; under market-burst reconnects on ticker / trade / fill channels, the SDK could drop tens of messages per reconnect. SubscriptionManager now stashes non-matching data frames in a per-sid bounded deque (stash_maxlen=1000 default) for the duration of resubscribe_all. _handle_reconnect drains the stash through _process_frame after resubscribe completes so the frames flow through the normal dispatch path -- seq tracker advances correctly, orderbook manager applies, iterators receive them in arrival order. Coordinates with #139's seq-gap tracking: replayed frames go through seq_tracker.track exactly once, so the first live frame after resubscribe sees the right watermark and doesn't trip a spurious gap. Stash is bounded; overflow evicts oldest via deque maxlen with a single WARNING per fill event. Frames whose sid wasn't re-mapped (per-sub resubscribe failure, F-P-01 path from #77) are dropped on drain with a debug log. 5 unit tests cover the stash mechanics directly (collect-by-sid, no-sid drop, maxlen eviction, take_stash atomicity, default-off flag). 3 integration tests cover the dispatcher integration: - drain replays through dispatch into the iterator queue - drain advances seq_tracker on sequenced channels (orderbook) - drain skips frames whose sid didn't get re-mapped The race-engineered end-to-end test (server sends pre-ack data during a reconnect) is intentionally not included: race timing is hardware-dependent and the value is the post-drain dispatch path, which the direct tests cover deterministically. Existing reconnect tests in this file exercise the close -> reconnect -> resubscribe pipeline. --- CHANGELOG.md | 32 +++++ kalshi/ws/channels.py | 178 ++++++++++++++++++++------- kalshi/ws/client.py | 45 +++++++ tests/ws/test_channels.py | 73 +++++++++++ tests/ws/test_recv_loop_hardening.py | 140 +++++++++++++++++++++ 5 files changed, 421 insertions(+), 47 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index cbb9f2d..7047eea 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,38 @@ All notable changes to kalshi-sdk will be documented in this file. ## Unreleased +### WS resubscribe-window frame stashing (#176) + +Fixes silent message loss during reconnect bursts on high-volume channels. +Previously, between ``SubscriptionManager._sid_to_client.clear()`` and the +new sid mapping landing in ``_wait_for_response``, any data frame the +server sent on the freshly-assigned sid was non-matching from the wait's +perspective and **discarded** with a debug log. Under market-burst +reconnects on ``ticker`` / ``trade`` / ``fill``, the SDK could drop tens +of messages per reconnect. + +``SubscriptionManager`` now stashes those non-matching data frames in a +per-sid bounded deque (``stash_maxlen=1000`` per sid by default) for the +duration of ``resubscribe_all``. After resubscribe completes, +``KalshiWebSocket._handle_reconnect`` drains the stash through +``_process_frame`` so the frames flow through the normal dispatch path +— seq tracker advances, orderbook manager applies, iterator consumers +receive them in arrival order. + +The drain coordinates with #139's seq-gap tracking: replayed frames +go through ``seq_tracker.track`` exactly once, so the first live frame +after resubscribe sees the right watermark and doesn't trip a spurious +gap on what would otherwise look like a seq 0 → N jump. + +Stash bound: per-sid deque uses ``collections.deque(maxlen=stash_maxlen)``. +On overflow, oldest evicts (deque semantics) and a single WARNING per +fill event is logged so callers notice congestion. Memory is bounded at +``stash_maxlen * len(active_subs) * avg_frame_size`` worst-case. + +Frames whose sid did not get re-mapped during ``resubscribe_all`` (a +per-sub failure that #77's F-P-01 isolates) are dropped on drain with a +debug log — there's no consumer to deliver them to. + ### WS `run_forever(stop_event=...)` cooperative shutdown (#177) `KalshiWebSocket.run_forever()` now accepts an optional diff --git a/kalshi/ws/channels.py b/kalshi/ws/channels.py index b43c1ef..073bd0e 100644 --- a/kalshi/ws/channels.py +++ b/kalshi/ws/channels.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio +import collections import json import logging from typing import Any @@ -61,12 +62,30 @@ class SubscriptionManager: - server_sid -> client_id mapping (rebuilt on reconnect) """ - def __init__(self, connection: ConnectionManager) -> None: + def __init__( + self, + connection: ConnectionManager, + *, + stash_maxlen: int = 1000, + ) -> None: self._connection = connection self._subscriptions: dict[int, Subscription] = {} # client_id -> Subscription self._sid_to_client: dict[int, int] = {} # server_sid -> client_id self._next_client_id = 1 self._next_msg_id = 1 + # Stash for #176: data frames arriving during `_wait_for_response` + # while `_stashing` is True (i.e., during `resubscribe_all`) are + # captured here keyed by their server sid, then replayed through + # the dispatcher once resubscribe completes and all new sids are + # wired. Without this, the recv loop misses frames sent between + # `_sid_to_client.clear()` and the ack landing in + # `_wait_for_response`. Each per-sid deque is bounded by + # `stash_maxlen` (1000 default — generous enough for normal + # market-burst reconnects, low enough to bound memory if + # resubscribe stalls). + self._stash: dict[int, collections.deque[str]] = {} + self._stashing: bool = False + self._stash_maxlen: int = stash_maxlen def _get_msg_id(self) -> int: mid = self._next_msg_id @@ -78,9 +97,16 @@ async def _wait_for_response( ) -> dict[str, Any]: """Read frames until we get the response matching our command id. - Non-matching frames (e.g. data messages queued before the ack) are - logged and discarded. The recv loop is paused during subscribe so - these frames would not have been consumed anyway. + Non-matching frames during a normal subscribe/unsubscribe are + logged and discarded; the recv loop is paused so these frames + would not have been consumed anyway. + + During ``resubscribe_all`` (``self._stashing is True``) the + non-matching frames are STASHED by sid for replay after the + resubscribe completes (#176). Without this, frames arriving + between ``_sid_to_client.clear()`` and the new sid mapping are + silently dropped — a real signal-loss bug under high-volume + reconnect bursts. """ deadline = asyncio.get_event_loop().time() + timeout while True: @@ -103,11 +129,46 @@ async def _wait_for_response( data: dict[str, Any] = json.loads(raw) if data.get("id") == msg_id: return data - # Non-matching frame (data message that arrived before ack) + # Non-matching frame. During resubscribe, stash it by sid for + # post-resubscribe replay; otherwise discard. + if self._stashing: + self._maybe_stash(raw, data) + else: + logger.debug( + "Discarding non-matching frame during subscribe: type=%s", + data.get("type"), + ) + + def _maybe_stash(self, raw: str, data: dict[str, Any]) -> None: + """Save a raw data frame to the per-sid stash, if it carries a sid. + + Frames without a sid (control envelopes that fall through to the + wait loop) are still discarded — they have nowhere to replay to. + Per-sid deques are bounded by ``self._stash_maxlen``; on overflow, + the deque silently evicts oldest (deque's own ``maxlen`` semantics) + and a single WARNING per (sid, replay-cycle) is logged so callers + notice congestion. + """ + sid = data.get("sid") + if sid is None: logger.debug( - "Discarding non-matching frame during subscribe: type=%s", + "Stash mode: dropping non-matching frame with no sid: type=%s", data.get("type"), ) + return + bucket = self._stash.get(sid) + if bucket is None: + bucket = collections.deque(maxlen=self._stash_maxlen) + self._stash[sid] = bucket + elif len(bucket) == self._stash_maxlen: + # About to evict oldest. Log once per fill, not per frame. + logger.warning( + "Stash for sid %d is full (%d frames); oldest frame will be " + "evicted. Resubscribe may be stalled or the channel is too " + "high-volume for the configured stash_maxlen.", + sid, self._stash_maxlen, + ) + bucket.append(raw) async def subscribe( self, @@ -213,54 +274,77 @@ async def resubscribe_all(self) -> None: F-P-01: Each resubscribe is independent. If one fails, the failed subscription is removed and its iterator gets a sentinel, but other subscriptions continue working. The whole reconnect is not aborted. + + #176: ``_stashing`` is enabled for the duration of this method so + non-matching data frames received by ``_wait_for_response`` are + captured by sid for post-resubscribe replay. The caller + (``KalshiWebSocket._handle_reconnect``) is responsible for + draining the stash via ``take_stash()`` and routing the frames + back through the dispatcher. """ old_subs = dict(self._subscriptions) # F-P-03: clear sid->client mapping before sending any resubscribe # so stale in-flight frames with old sids can't mis-route. self._sid_to_client.clear() - for client_id, sub in old_subs.items(): - sub.server_sid = None # Clear old sid - try: - msg_id = self._get_msg_id() - # Re-subscribe with send_initial_snapshot for orderbook channels - 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) - 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"), + self._stashing = True + try: + for client_id, sub in old_subs.items(): + sub.server_sid = None # Clear old sid + try: + msg_id = self._get_msg_id() + # Re-subscribe with send_initial_snapshot for orderbook channels + 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) + 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"), + ) + new_sid = data.get("msg", {}).get("sid") + if new_sid is not None: + sub.server_sid = new_sid + self._sid_to_client[new_sid] = client_id + logger.debug( + "Resubscribed %s: client_id=%d, new_sid=%s", + sub.channel, + client_id, + new_sid, ) - new_sid = data.get("msg", {}).get("sid") - if new_sid is not None: - sub.server_sid = new_sid - self._sid_to_client[new_sid] = client_id - logger.debug( - "Resubscribed %s: client_id=%d, new_sid=%s", - sub.channel, - client_id, - new_sid, - ) - except Exception: - # F-P-01: per-sub failure is isolated. Push sentinel so the - # iterator exits cleanly, drop the subscription, continue with - # the rest. `exc_info=True` so an unexpected AttributeError / - # programming bug doesn't lose its traceback the way the - # earlier `: %s, e` formulation did. - logger.warning( - "Resubscribe failed for client_id=%d channel=%s", - client_id, - sub.channel, - exc_info=True, - ) - await sub.queue.put_sentinel() - self._subscriptions.pop(client_id, None) + except Exception: + # F-P-01: per-sub failure is isolated. Push sentinel so the + # iterator exits cleanly, drop the subscription, continue with + # the rest. `exc_info=True` so an unexpected AttributeError / + # programming bug doesn't lose its traceback the way the + # earlier `: %s, e` formulation did. + logger.warning( + "Resubscribe failed for client_id=%d channel=%s", + client_id, + sub.channel, + exc_info=True, + ) + await sub.queue.put_sentinel() + self._subscriptions.pop(client_id, None) + finally: + self._stashing = False + + def take_stash(self) -> dict[int, collections.deque[str]]: + """Return and clear the resubscribe-window stash atomically (#176). + + Returns a ``{sid: deque[raw_frame]}`` mapping. The caller is + responsible for replaying the raw frames through the dispatcher + for each sid (typically via ``KalshiWebSocket._process_frame``). + Frames whose sid did not get re-mapped during resubscribe should + be dropped by the caller with a log. + """ + stash, self._stash = self._stash, {} + return stash def get_subscription_by_sid(self, server_sid: int) -> Subscription | None: """Look up a subscription by current server sid.""" diff --git a/kalshi/ws/client.py b/kalshi/ws/client.py index 0a7a822..d364f2e 100644 --- a/kalshi/ws/client.py +++ b/kalshi/ws/client.py @@ -322,6 +322,13 @@ async def _handle_reconnect(self) -> None: if self._orderbook_mgr: self._orderbook_mgr.clear() await self._sub_mgr.resubscribe_all() + # #176: replay frames that the server sent on + # newly-assigned sids before the corresponding + # subscribe ack landed. SubscriptionManager stashed + # them keyed by sid during _wait_for_response; drain + # through _process_frame so seq tracking + orderbook + # state stay consistent with the natural arrival path. + await self._drain_resubscribe_stash() # #88: use the public transition; no reach-through to # ConnectionManager's name-mangled _set_state. await self._connection.mark_streaming() @@ -332,6 +339,44 @@ async def _handle_reconnect(self) -> None: await self._broadcast_sentinels() self._running = False + async def _drain_resubscribe_stash(self) -> None: + """Replay frames captured during ``resubscribe_all`` through dispatch. + + #176: ``SubscriptionManager._wait_for_response`` captures non-matching + data frames into a per-sid stash while a resubscribe is in flight, + because between ``_sid_to_client.clear()`` and the new sid mapping + the dispatcher has no route for those frames. After all subscribes + complete, this method drains the stash via ``_process_frame`` so + seq tracking and orderbook state stay consistent with the natural + arrival path. + + Frames whose sid did not get re-mapped (subscription failed during + ``resubscribe_all``) are dropped with a debug log — there's no + consumer to deliver them to. + """ + if self._sub_mgr is None: + return + stash = self._sub_mgr.take_stash() + for sid, raw_frames in stash.items(): + if self._sub_mgr.get_subscription_by_sid(sid) is None: + logger.debug( + "Dropping %d stashed frames for unmapped sid %d " + "after resubscribe (subscription likely failed)", + len(raw_frames), sid, + ) + continue + for raw in raw_frames: + try: + await self._process_frame(raw) + except Exception: + # Per-frame failure during replay: log and continue. + # The recv-loop's own error-handling is bypassed here + # because we're in the orchestration path, not the loop. + logger.warning( + "Failed to replay stashed frame for sid %d", + sid, exc_info=True, + ) + async def _handle_seq_gap(self, gap: SequenceGap) -> None: """Handle a sequence gap by logging and triggering resync. diff --git a/tests/ws/test_channels.py b/tests/ws/test_channels.py index 7a2bbec..0daa448 100644 --- a/tests/ws/test_channels.py +++ b/tests/ws/test_channels.py @@ -2,6 +2,8 @@ from __future__ import annotations +import json + import pytest from kalshi.config import KalshiConfig @@ -332,3 +334,74 @@ async def test_msg_ids_are_sequential( await sub_mgr.subscribe("fill") ids = [c["id"] for c in fake_ws.received_commands] assert ids == [1, 2] + + +# --------------------------------------------------------------------------- +# SubscriptionManager — resubscribe-window stash (#176) +# --------------------------------------------------------------------------- + + +class TestResubscribeStash: + def test_maybe_stash_collects_by_sid( + self, sub_mgr: SubscriptionManager + ) -> None: + """Frames carrying a sid land in the per-sid deque in arrival order.""" + sub_mgr._maybe_stash('{"sid": 1, "seq": 1}', {"sid": 1, "seq": 1}) + sub_mgr._maybe_stash('{"sid": 2, "seq": 1}', {"sid": 2, "seq": 1}) + sub_mgr._maybe_stash('{"sid": 1, "seq": 2}', {"sid": 1, "seq": 2}) + assert set(sub_mgr._stash.keys()) == {1, 2} + assert list(sub_mgr._stash[1]) == [ + '{"sid": 1, "seq": 1}', + '{"sid": 1, "seq": 2}', + ] + assert list(sub_mgr._stash[2]) == ['{"sid": 2, "seq": 1}'] + + def test_maybe_stash_drops_frame_without_sid( + self, sub_mgr: SubscriptionManager + ) -> None: + """Control envelopes without a sid have nowhere to replay to — drop.""" + sub_mgr._maybe_stash('{"type": "ok"}', {"type": "ok"}) + assert sub_mgr._stash == {} + + def test_maybe_stash_maxlen_evicts_oldest( + self, connected_mgr, caplog # type: ignore[no-untyped-def] + ) -> None: + """Per-sid deque is bounded; on overflow, oldest evicts and a WARNING fires.""" + import logging + mgr = SubscriptionManager(connected_mgr, stash_maxlen=3) + with caplog.at_level(logging.WARNING, logger="kalshi.ws"): + for i in range(5): + mgr._maybe_stash(f'{{"sid": 7, "seq": {i}}}', {"sid": 7, "seq": i}) + bucket = mgr._stash[7] + # maxlen=3 → only the last 3 survive (oldest 0,1 evicted) + assert len(bucket) == 3 + assert [json.loads(r)["seq"] for r in bucket] == [2, 3, 4] + # WARNING fired (once per fill event, not per frame) + assert any( + "is full (3 frames)" in rec.message + for rec in caplog.records + if rec.levelno == logging.WARNING + ) + + def test_take_stash_returns_and_clears( + self, sub_mgr: SubscriptionManager + ) -> None: + """take_stash atomically returns the current stash and resets to empty.""" + sub_mgr._maybe_stash('{"sid": 9, "seq": 1}', {"sid": 9, "seq": 1}) + sub_mgr._maybe_stash('{"sid": 9, "seq": 2}', {"sid": 9, "seq": 2}) + taken = sub_mgr.take_stash() + assert list(taken[9]) == [ + '{"sid": 9, "seq": 1}', + '{"sid": 9, "seq": 2}', + ] + assert sub_mgr._stash == {} + # Second take is empty + assert sub_mgr.take_stash() == {} + + def test_stashing_flag_off_by_default( + self, sub_mgr: SubscriptionManager + ) -> None: + """_wait_for_response only stashes when explicitly toggled by + resubscribe_all (or future opt-in callers). Default is off so + normal subscribe paths don't accumulate stale state.""" + assert sub_mgr._stashing is False diff --git a/tests/ws/test_recv_loop_hardening.py b/tests/ws/test_recv_loop_hardening.py index 25f6b5e..76b37b5 100644 --- a/tests/ws/test_recv_loop_hardening.py +++ b/tests/ws/test_recv_loop_hardening.py @@ -8,8 +8,10 @@ from __future__ import annotations import asyncio +import collections import json import logging +from decimal import Decimal from typing import Any import pytest @@ -448,3 +450,141 @@ async def boom(_data: dict[str, Any], **_kw: Any) -> None: assert recv_task is not None assert recv_task.done() assert isinstance(recv_task.exception(), AttributeError) + + +# --------------------------------------------------------------------------- +# #176 — resubscribe-window stash: data frames sent before subscribe ack +# are stashed by sid and replayed after the new sid mapping lands +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +class TestResubscribeStashIntegration: + async def test_stash_drain_replays_through_dispatch( + self, + fake_ws, + test_auth, # type: ignore[no-untyped-def] + ) -> None: + """#176: a frame stashed by sid during resubscribe must replay + through `_process_frame` → dispatch so the iterator receives it + and the seq tracker treats it as the natural first frame on the + new sid (no spurious gap on the next live frame). + + Directly exercises `_drain_resubscribe_stash` rather than racing + a real server: race timing is hardware-dependent and the value + is the post-drain dispatch path, not the race detector itself. + Other tests in this file already exercise the close→reconnect→ + resubscribe pipeline. + """ + 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"]) + assert session._sub_mgr is not None + sub = next(iter(session._sub_mgr.active_subscriptions.values())) + sid = sub.server_sid + assert sid is not None + + # Inject a stashed frame for the live sid, then drain. + stashed_frame = json.dumps({ + "type": "ticker", + "sid": sid, + "seq": 1, + "msg": ticker_payload_dict( + market_ticker="T1", + market_id="m1", + yes_bid_dollars="0.4500", + ), + }) + session._sub_mgr._stash[sid] = collections.deque([stashed_frame]) + + await session._drain_resubscribe_stash() + + # The replayed frame reached the iterator queue. + msg = await asyncio.wait_for(stream.__anext__(), timeout=2.0) + assert msg.msg.market_ticker == "T1" + assert msg.msg.yes_bid == Decimal("0.4500") + # Stash is empty after drain. + assert session._sub_mgr._stash == {} + + async def test_stash_drain_advances_seq_tracker_on_sequenced_channels( + self, + fake_ws, + test_auth, # type: ignore[no-untyped-def] + ) -> None: + """#176 + #139: a stashed orderbook_delta frame must advance the + seq tracker via the normal `_process_frame → seq_tracker.track` + path, so the next live frame on the same sid (seq+1) doesn't + trip a spurious gap. + + Without the drain going through `_process_frame`, the seq tracker + would never see the stashed seq, and the first live frame after + resubscribe would look like a gap from seq 0 → seq 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"]) + assert session._sub_mgr is not None + sub = next(iter(session._sub_mgr.active_subscriptions.values())) + sid = sub.server_sid + assert sid is not None + + # Inject a stashed snapshot (seq=1) for the live sid. + stashed_frame = json.dumps({ + "type": "orderbook_snapshot", + "sid": sid, + "seq": 1, + "msg": { + "market_ticker": "T1", + "market_id": "m1", + "yes": [["0.50", "100"]], + "no": [], + }, + }) + session._sub_mgr._stash[sid] = collections.deque([stashed_frame]) + + await session._drain_resubscribe_stash() + + # Replay drained the frame to the iterator… + msg = await asyncio.wait_for(stream.__anext__(), timeout=2.0) + assert msg.msg.market_ticker == "T1" + # …and advanced the seq tracker watermark, so seq=2 on the + # next live frame won't look like a gap from 0 -> 2. + assert session._seq_tracker is not None + assert session._seq_tracker.peek(sid) == 1 + + async def test_resubscribe_stash_skips_unmapped_sids( + self, + fake_ws, + test_auth, # type: ignore[no-untyped-def] + caplog, + ) -> None: + """If a sub fails during resubscribe (no new sid mapped), any + frames stashed under a sid that the server did emit but that's + no longer in `_sid_to_client` get dropped with a debug log + rather than crashing the drain or routing to nowhere.""" + import logging + + 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_ticker(tickers=["T1"]) + assert session._sub_mgr is not None + + # Pre-populate the stash with a frame for an unknown sid as + # if resubscribe had captured it but the sub failed. + session._sub_mgr._stash[999] = collections.deque( + ['{"type": "ticker", "sid": 999, "seq": 1, "msg": {}}'] + ) + + with caplog.at_level(logging.DEBUG, logger="kalshi.ws"): + await session._drain_resubscribe_stash() + + # Stash is cleared. + assert session._sub_mgr._stash == {} + # Drop was logged. + assert any( + "Dropping 1 stashed frames for unmapped sid 999" in r.message + for r in caplog.records + ) From 6eafeab507defc98ee6056da09c59c5e02b38aed Mon Sep 17 00:00:00 2001 From: Jeff West Date: Wed, 20 May 2026 18:02:51 -0500 Subject: [PATCH 2/5] fixup: gate overflow WARNING per (sid, resubscribe cycle), not per frame MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Bot caught a real bug: `_maybe_stash`'s `elif len(bucket) == self._stash_maxlen` check is True on EVERY append once the deque is full, so the WARNING fired on every frame after the fill, not once as the comment and PR description claimed. Under a prolonged stall on a high-volume sid, this is per-frame log spam. Fix: track warned sids in a `_stash_warned: set[int]` attribute. The warning fires once per (sid, resubscribe cycle); `take_stash()` clears the set so the next resubscribe gets fresh warnings. Tightened the test: - Renamed to test_maybe_stash_maxlen_evicts_oldest_with_one_warning_per_fill - Pushes 10 frames into a maxlen=3 deque, asserts EXACTLY one WARNING (not >= 1 via `any(...)` as before) - Verifies take_stash() reset by running a second 5-frame cycle and asserting the total warning count goes from 1 -> 2 (not 1+5=6) Also moved `import logging` from inside the test body to module level. asyncio.get_event_loop() deprecation (channels.py:111,113) noted by bot but left alone — pre-existing, out of scope here. --- kalshi/ws/channels.py | 18 ++++++++++++++++-- tests/ws/test_channels.py | 40 +++++++++++++++++++++++++++------------ 2 files changed, 44 insertions(+), 14 deletions(-) diff --git a/kalshi/ws/channels.py b/kalshi/ws/channels.py index 073bd0e..9c130ef 100644 --- a/kalshi/ws/channels.py +++ b/kalshi/ws/channels.py @@ -86,6 +86,13 @@ def __init__( self._stash: dict[int, collections.deque[str]] = {} self._stashing: bool = False self._stash_maxlen: int = stash_maxlen + # Sids that have already triggered an overflow WARNING in the current + # resubscribe cycle. Cleared by `take_stash()` so the next resubscribe + # gets fresh warnings. Without this gating, the `len(bucket) == + # maxlen` check below would be True on every append after the deque + # fills, producing one WARNING per frame on every high-volume + # channel during a prolonged stall. + self._stash_warned: set[int] = set() def _get_msg_id(self) -> int: mid = self._next_msg_id @@ -160,8 +167,12 @@ def _maybe_stash(self, raw: str, data: dict[str, Any]) -> None: if bucket is None: bucket = collections.deque(maxlen=self._stash_maxlen) self._stash[sid] = bucket - elif len(bucket) == self._stash_maxlen: - # About to evict oldest. Log once per fill, not per frame. + elif len(bucket) == self._stash_maxlen and sid not in self._stash_warned: + # About to evict oldest. Log once per (sid, resubscribe cycle): + # without the warned-set gate this fires on every subsequent + # append while the deque stays full, which under a prolonged + # stall becomes per-frame log spam on every high-volume sid. + self._stash_warned.add(sid) logger.warning( "Stash for sid %d is full (%d frames); oldest frame will be " "evicted. Resubscribe may be stalled or the channel is too " @@ -344,6 +355,9 @@ def take_stash(self) -> dict[int, collections.deque[str]]: be dropped by the caller with a log. """ stash, self._stash = self._stash, {} + # Reset warning suppression so the next resubscribe cycle's overflow + # warnings aren't accidentally muted by stale state. + self._stash_warned.clear() return stash def get_subscription_by_sid(self, server_sid: int) -> Subscription | None: diff --git a/tests/ws/test_channels.py b/tests/ws/test_channels.py index 0daa448..fc35076 100644 --- a/tests/ws/test_channels.py +++ b/tests/ws/test_channels.py @@ -3,6 +3,7 @@ from __future__ import annotations import json +import logging import pytest @@ -363,25 +364,40 @@ def test_maybe_stash_drops_frame_without_sid( sub_mgr._maybe_stash('{"type": "ok"}', {"type": "ok"}) assert sub_mgr._stash == {} - def test_maybe_stash_maxlen_evicts_oldest( + def test_maybe_stash_maxlen_evicts_oldest_with_one_warning_per_fill( self, connected_mgr, caplog # type: ignore[no-untyped-def] ) -> None: - """Per-sid deque is bounded; on overflow, oldest evicts and a WARNING fires.""" - import logging + """Per-sid deque is bounded; on overflow, oldest evicts and EXACTLY + ONE WARNING fires per (sid, resubscribe cycle) — not one per frame + (#187 review fix). Without the per-sid gate, every append after the + deque fills would log, producing per-frame spam on high-volume + channels during a prolonged stall.""" mgr = SubscriptionManager(connected_mgr, stash_maxlen=3) with caplog.at_level(logging.WARNING, logger="kalshi.ws"): - for i in range(5): + for i in range(10): mgr._maybe_stash(f'{{"sid": 7, "seq": {i}}}', {"sid": 7, "seq": i}) bucket = mgr._stash[7] - # maxlen=3 → only the last 3 survive (oldest 0,1 evicted) + # maxlen=3 → only the last 3 survive (oldest 0..6 evicted) assert len(bucket) == 3 - assert [json.loads(r)["seq"] for r in bucket] == [2, 3, 4] - # WARNING fired (once per fill event, not per frame) - assert any( - "is full (3 frames)" in rec.message - for rec in caplog.records - if rec.levelno == logging.WARNING - ) + assert [json.loads(r)["seq"] for r in bucket] == [7, 8, 9] + # EXACTLY one WARNING for sid 7 across 10 appends. + warns = [ + r for r in caplog.records + if r.levelno == logging.WARNING and "Stash for sid 7 is full" in r.message + ] + assert len(warns) == 1 + # take_stash() clears the warning suppression so the next cycle + # gets fresh warnings. + mgr.take_stash() + with caplog.at_level(logging.WARNING, logger="kalshi.ws"): + for i in range(10, 15): + mgr._maybe_stash(f'{{"sid": 7, "seq": {i}}}', {"sid": 7, "seq": i}) + warns2 = [ + r for r in caplog.records + if r.levelno == logging.WARNING and "Stash for sid 7 is full" in r.message + ] + # Two cycles → two warnings total (the original + one for the new fill). + assert len(warns2) == 2 def test_take_stash_returns_and_clears( self, sub_mgr: SubscriptionManager From 9767e57e0eca8119ed0729433634bb3caefbac49 Mon Sep 17 00:00:00 2001 From: Jeff West Date: Wed, 20 May 2026 18:11:09 -0500 Subject: [PATCH 3/5] fixup: defensive stash clear + drop deprecated get_event_loop() Round-2 bot review on PR #187: 1. asyncio.get_event_loop() in _wait_for_response was deprecated since 3.10. Pre-existing, but since this PR already touches the function body it's in-scope. Swapped both calls to asyncio.get_running_loop().time() which is the correct API inside an async def with a guaranteed-running loop. 2. Stash state could survive a resubscribe_all exception. Today the outer _handle_reconnect's except sets _running=False so no next reconnect happens on the same instance, but the invariant should be defensive: if a future code path ever calls resubscribe_all on a manager with prior failed-cycle state, stale stashed frames would replay alongside fresh ones. Added a clear() of both _stash and _stash_warned at the start of resubscribe_all (before _stashing = True). take_stash() still clears them on the happy path; this is purely the failure-mode guard. Regression test added (test_resubscribe_clears_stale_stash_at_start): inject stale _stash + _stash_warned state, run a clean resubscribe_all, verify both are empty regardless of what the per-sub loop did. 2126 unit tests pass. mypy strict + ruff clean. --- kalshi/ws/channels.py | 12 ++++++++++-- tests/ws/test_channels.py | 28 ++++++++++++++++++++++++++++ 2 files changed, 38 insertions(+), 2 deletions(-) diff --git a/kalshi/ws/channels.py b/kalshi/ws/channels.py index 9c130ef..72ee3e2 100644 --- a/kalshi/ws/channels.py +++ b/kalshi/ws/channels.py @@ -115,9 +115,9 @@ async def _wait_for_response( silently dropped — a real signal-loss bug under high-volume reconnect bursts. """ - deadline = asyncio.get_event_loop().time() + timeout + deadline = asyncio.get_running_loop().time() + timeout while True: - remaining = deadline - asyncio.get_event_loop().time() + remaining = deadline - asyncio.get_running_loop().time() if remaining <= 0: raise KalshiSubscriptionError( f"Timed out waiting for response to command {msg_id}" @@ -297,6 +297,14 @@ async def resubscribe_all(self) -> None: # F-P-03: clear sid->client mapping before sending any resubscribe # so stale in-flight frames with old sids can't mis-route. self._sid_to_client.clear() + # Defensive reset: if a prior resubscribe cycle raised before + # _drain_resubscribe_stash (which goes through take_stash()) had a + # chance to run, _stash and _stash_warned could carry stale state + # into this cycle. Within the current KalshiWebSocket instance the + # outer _handle_reconnect except sets _running=False so this case is + # rare; cheap to harden against anyway. (#187 review) + self._stash.clear() + self._stash_warned.clear() self._stashing = True try: diff --git a/tests/ws/test_channels.py b/tests/ws/test_channels.py index fc35076..a21a1bd 100644 --- a/tests/ws/test_channels.py +++ b/tests/ws/test_channels.py @@ -2,6 +2,7 @@ from __future__ import annotations +import collections import json import logging @@ -421,3 +422,30 @@ def test_stashing_flag_off_by_default( resubscribe_all (or future opt-in callers). Default is off so normal subscribe paths don't accumulate stale state.""" assert sub_mgr._stashing is False + + async def test_resubscribe_clears_stale_stash_at_start( + self, + sub_mgr: SubscriptionManager, + fake_ws, # type: ignore[no-untyped-def] + ) -> None: + """#187 review: if a prior resubscribe raised before take_stash() + ran, _stash + _stash_warned could survive into the next cycle and + replay stale frames or muddy overflow warnings. resubscribe_all + now defensively clears both at the start. + + Simulates the leak by injecting stale stash state, then runs a + clean resubscribe (no active subs → loop body skipped, but the + clear runs).""" + # Inject stale state from a hypothetical prior cycle that failed + # before draining. + sub_mgr._stash[42] = collections.deque(['{"sid": 42, "stale": true}']) + sub_mgr._stash_warned.add(42) + assert sub_mgr._stash != {} + assert sub_mgr._stash_warned == {42} + + await sub_mgr.resubscribe_all() + + # Stash + warned-set are both empty regardless of what the + # subscriptions loop did. + assert sub_mgr._stash == {} + assert sub_mgr._stash_warned == set() From 6b4db52b91c339ee24e1139886a30be405264d82 Mon Sep 17 00:00:00 2001 From: Jeff West Date: Wed, 20 May 2026 18:18:05 -0500 Subject: [PATCH 4/5] fixup: drop redundant import + PR-ref tag, document get_running_loop in CHANGELOG MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Round-3 bot review on PR #187: 1. Redundant `import logging` inside test_resubscribe_stash_skips_unmapped_sids removed. The module-level import at the top of the file already covers it. 2. Dropped the `(#187 review)` tag from the channels.py defensive-clear comment. The text explanation stays — it's the rationale that matters, not the PR number, which becomes a meaningless reference once merged. Issue numbers (#176, #172, #77) are durable and stay; PR numbers in production code rot. 3. Added a one-line CHANGELOG note for the asyncio.get_event_loop() -> get_running_loop() migration in _wait_for_response so the drive-by doesn't get lost in code review. The bot also suggested an explicit test that _stashing=False keeps the discard path intact. The existing test_stashing_flag_off_by_default already locks in the default state, and the structural invariant (_maybe_stash only called inside the `if self._stashing:` branch in _wait_for_response) is visible in the code. Adding a standalone test for the discard path would require mocking _wait_for_response and the underlying recv() to feed it a non-matching frame; the contract is already locked by the existing test plus code structure. Not adding. 2126 unit tests pass. ruff + mypy strict clean. --- CHANGELOG.md | 4 ++++ kalshi/ws/channels.py | 2 +- tests/ws/test_recv_loop_hardening.py | 2 -- 3 files changed, 5 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 7047eea..2329e1b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -36,6 +36,10 @@ Frames whose sid did not get re-mapped during ``resubscribe_all`` (a per-sub failure that #77's F-P-01 isolates) are dropped on drain with a debug log — there's no consumer to deliver them to. +Drive-by: ``SubscriptionManager._wait_for_response`` swapped two deprecated +``asyncio.get_event_loop().time()`` calls for ``asyncio.get_running_loop().time()`` +(the correct API inside an ``async def``). + ### WS `run_forever(stop_event=...)` cooperative shutdown (#177) `KalshiWebSocket.run_forever()` now accepts an optional diff --git a/kalshi/ws/channels.py b/kalshi/ws/channels.py index 72ee3e2..3b0e143 100644 --- a/kalshi/ws/channels.py +++ b/kalshi/ws/channels.py @@ -302,7 +302,7 @@ async def resubscribe_all(self) -> None: # chance to run, _stash and _stash_warned could carry stale state # into this cycle. Within the current KalshiWebSocket instance the # outer _handle_reconnect except sets _running=False so this case is - # rare; cheap to harden against anyway. (#187 review) + # rare; cheap to harden against anyway. self._stash.clear() self._stash_warned.clear() diff --git a/tests/ws/test_recv_loop_hardening.py b/tests/ws/test_recv_loop_hardening.py index 76b37b5..4067edb 100644 --- a/tests/ws/test_recv_loop_hardening.py +++ b/tests/ws/test_recv_loop_hardening.py @@ -564,8 +564,6 @@ async def test_resubscribe_stash_skips_unmapped_sids( frames stashed under a sid that the server did emit but that's no longer in `_sid_to_client` get dropped with a debug log rather than crashing the drain or routing to nowhere.""" - import logging - config = KalshiConfig(ws_base_url=fake_ws.url, timeout=5.0) ws = KalshiWebSocket(auth=test_auth, config=config) async with ws.connect() as session: From 2dfebc1019658eed15b26d46a5fd7c51649b50c6 Mon Sep 17 00:00:00 2001 From: Jeff West Date: Wed, 20 May 2026 18:26:48 -0500 Subject: [PATCH 5/5] fixup: isinstance(sid, int) guard + drop PR-refs from docstrings MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Round-4 bot review on PR #187: 1. Untyped sid from data.get("sid"). The stash is typed as `dict[int, ...]` but data is `dict[str, Any]`, so a malformed frame with a non-int sid would silently corrupt the stash. Replaced `if sid is None:` with `if not isinstance(sid, int):` to catch both None and string/float/etc. Added a comment noting the Kalshi protocol's monotonic-int-sid assumption and the within-cycle uniqueness assumption. Renamed the test `test_maybe_stash_drops_frame_without_sid` -> `test_maybe_stash_drops_frame_with_non_int_sid` and added two more cases: string sid and float sid. All drop with empty stash result. 2. Dropped `(#187 review fix)` from `test_maybe_stash_maxlen_evicts_oldest_with_one_warning_per_fill` and `#187 review:` from `test_resubscribe_clears_stale_stash_at_start`. PR refs rot once merged; the test docstrings already explain the what and why without the tag. Not actioned: - Bot's "broad except Exception in _drain_resubscribe_stash" — bot itself noted this is "a trade-off judgment call that's reasonably decided in the current direction." Leaving as-is; a malformed replayed frame shouldn't abort the entire reconnect orchestration. - Bot's "theoretical sid collision across resubscribes" — added a comment in _maybe_stash documenting the within-cycle uniqueness assumption. No code change needed. 2126 unit tests pass. ruff + mypy strict clean. --- kalshi/ws/channels.py | 13 ++++++++++--- tests/ws/test_channels.py | 24 ++++++++++++++---------- 2 files changed, 24 insertions(+), 13 deletions(-) diff --git a/kalshi/ws/channels.py b/kalshi/ws/channels.py index 3b0e143..0ae6db2 100644 --- a/kalshi/ws/channels.py +++ b/kalshi/ws/channels.py @@ -157,10 +157,17 @@ def _maybe_stash(self, raw: str, data: dict[str, Any]) -> None: notice congestion. """ sid = data.get("sid") - if sid is None: + # The Kalshi protocol uses monotonically-increasing integer sids; + # this guard catches malformed frames (non-int sid, missing sid) + # without corrupting the typed `dict[int, ...]` stash. Sids are + # also assumed unique within a single resubscribe cycle — keys + # collide only if the server reused a value, which would be a + # protocol bug, not a stash bug. + if not isinstance(sid, int): logger.debug( - "Stash mode: dropping non-matching frame with no sid: type=%s", - data.get("type"), + "Stash mode: dropping non-matching frame with non-int sid: " + "type=%s sid=%r", + data.get("type"), sid, ) return bucket = self._stash.get(sid) diff --git a/tests/ws/test_channels.py b/tests/ws/test_channels.py index a21a1bd..3857c72 100644 --- a/tests/ws/test_channels.py +++ b/tests/ws/test_channels.py @@ -358,21 +358,25 @@ def test_maybe_stash_collects_by_sid( ] assert list(sub_mgr._stash[2]) == ['{"sid": 2, "seq": 1}'] - def test_maybe_stash_drops_frame_without_sid( + def test_maybe_stash_drops_frame_with_non_int_sid( self, sub_mgr: SubscriptionManager ) -> None: - """Control envelopes without a sid have nowhere to replay to — drop.""" + """Control envelopes without a sid (or with a non-int sid from a + malformed/foreign frame) have nowhere to replay to and would + corrupt the typed `dict[int, ...]` stash. Drop with a debug log.""" sub_mgr._maybe_stash('{"type": "ok"}', {"type": "ok"}) + sub_mgr._maybe_stash('{"sid": "x"}', {"sid": "x"}) + sub_mgr._maybe_stash('{"sid": 1.5}', {"sid": 1.5}) assert sub_mgr._stash == {} def test_maybe_stash_maxlen_evicts_oldest_with_one_warning_per_fill( self, connected_mgr, caplog # type: ignore[no-untyped-def] ) -> None: """Per-sid deque is bounded; on overflow, oldest evicts and EXACTLY - ONE WARNING fires per (sid, resubscribe cycle) — not one per frame - (#187 review fix). Without the per-sid gate, every append after the - deque fills would log, producing per-frame spam on high-volume - channels during a prolonged stall.""" + ONE WARNING fires per (sid, resubscribe cycle) — not one per frame. + Without the per-sid gate, every append after the deque fills would + log, producing per-frame spam on high-volume channels during a + prolonged stall.""" mgr = SubscriptionManager(connected_mgr, stash_maxlen=3) with caplog.at_level(logging.WARNING, logger="kalshi.ws"): for i in range(10): @@ -428,10 +432,10 @@ async def test_resubscribe_clears_stale_stash_at_start( sub_mgr: SubscriptionManager, fake_ws, # type: ignore[no-untyped-def] ) -> None: - """#187 review: if a prior resubscribe raised before take_stash() - ran, _stash + _stash_warned could survive into the next cycle and - replay stale frames or muddy overflow warnings. resubscribe_all - now defensively clears both at the start. + """If a prior resubscribe raised before take_stash() ran, _stash + + _stash_warned could survive into the next cycle and replay stale + frames or muddy overflow warnings. resubscribe_all defensively + clears both at the start. Simulates the leak by injecting stale stash state, then runs a clean resubscribe (no active subs → loop body skipped, but the