From 40a966c37d1675eae768e500b8e168210ddec09d Mon Sep 17 00:00:00 2001 From: Jeff West Date: Fri, 22 May 2026 08:21:24 -0500 Subject: [PATCH 1/3] fix(ws): close WS connection on KalshiBackpressureError to prevent orphan recv loop MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When the WS recv loop hit KalshiBackpressureError (or KalshiSubscriptionError) it broadcast queue sentinels and broke — but left the websocket connection open, _running=True, and the subscription manager populated. A subsequent subscribe_* on the same KalshiWebSocket instance (the #297 re-entry guard only fires when _connection is None or not _running) then resurrected a fresh recv task on top of the orphaned server-side subscriptions, with frames silently dropped onto the already-closed queues or mis-routed to stale sids without local snapshots. This commit mirrors the #297 reset pattern from _stop() inside the fatal-error branch of _recv_loop: broadcast sentinels, flip _running=False, close the underlying connection (suppressing close failures), and clear _connection / _sub_mgr / _seq_tracker / _orderbook_mgr / _dispatcher and pause state — without awaiting the recv task because we ARE the recv task. _do_subscribe gains an explicit guard that raises a clear RuntimeError when the session is no longer active (replacing a bare assert), so a follow-up subscribe on a torn-down session fails fast instead of detonating an AssertionError or restarting a dead session. The same KalshiWebSocket instance remains reusable for a fresh async with ws.connect() block — the #297 invariant is preserved end-to-end. tests/ws/test_seq_correctness.py: the existing #78 invariant test snapshots _seq_tracker before awaiting the recv task — the tracker is now cleared as part of teardown, so the assertion runs on the detached object. Substantive intent of the test is unchanged. Closes #332 --- kalshi/ws/client.py | 42 +++++++- tests/ws/test_client.py | 166 +++++++++++++++++++++++++++++++ tests/ws/test_seq_correctness.py | 8 +- 3 files changed, 211 insertions(+), 5 deletions(-) diff --git a/kalshi/ws/client.py b/kalshi/ws/client.py index 1140d373..86456079 100644 --- a/kalshi/ws/client.py +++ b/kalshi/ws/client.py @@ -337,10 +337,32 @@ async def _recv_loop(self) -> None: break except (KalshiBackpressureError, KalshiSubscriptionError): # #83: consumer-visible problem. Broadcast sentinels so - # iterators see the failure and the loop exits cleanly - # rather than spinning on the same error. - logger.error("Fatal WS error in recv loop", exc_info=True) + # iterators see the failure rather than spinning on the + # same error. #332: also tear the session down — close the + # WS connection and clear ``_running`` + manager refs so a + # subsequent ``subscribe_*`` on this instance does not + # resurrect a recv loop on top of orphaned server-side + # subscriptions (the queues here are now closed; new + # frames would silently drop or mis-route). Mirrors the + # #297 reset done by ``_stop()``; we skip the recv-task + # await because we ARE the recv task. + logger.error( + "Fatal WS error in recv loop; tearing down session", + exc_info=True, + ) await self._broadcast_sentinels() + self._running = False + if self._connection is not None: + with contextlib.suppress(Exception): + await self._connection.close() + self._connection = None + self._sub_mgr = None + self._seq_tracker = None + self._orderbook_mgr = None + self._dispatcher = None + self._pause_pending = False + self._pause_granted.clear() + self._resume_signal.clear() break except (json.JSONDecodeError, ValidationError, KeyError): # #83: genuinely-non-fatal per-message errors. #241: when @@ -650,8 +672,20 @@ async def _do_subscribe( """Pause recv_loop, subscribe, resume recv_loop, return queue. Serialized with asyncio.Lock to prevent concurrent subscribe races. + + #332: when the recv loop has torn itself down on a fatal error + (``KalshiBackpressureError`` / ``KalshiSubscriptionError``), the + managers are cleared. Fail fast with a clear error instead of + letting the assert below fire (or — worse — resurrecting a recv + loop on top of orphaned server-side subscriptions). """ - assert self._sub_mgr is not None + if not self._running or self._sub_mgr is None or self._connection is None: + raise RuntimeError( + "KalshiWebSocket session is not active. The recv loop tore " + "down after a fatal error (e.g. KalshiBackpressureError); " + "exit the current `async with ws.connect()` block and start " + "a fresh session before subscribing again." + ) async with self._subscribe_lock: await self._pause_recv_loop() try: diff --git a/tests/ws/test_client.py b/tests/ws/test_client.py index 775fec80..29fae755 100644 --- a/tests/ws/test_client.py +++ b/tests/ws/test_client.py @@ -766,3 +766,169 @@ async def recv(self) -> str: # pragma: no cover - not used assert 11 not in mgr._subscriptions assert 999 not in mgr._sid_to_client + +# --------------------------------------------------------------------------- +# #332 — KalshiBackpressureError tears down the session cleanly +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +class TestIssue332BackpressureTeardown: + async def test_issue_332_backpressure_closes_ws_cleanly( + self, + fake_ws, # type: ignore[no-untyped-def] + test_auth, # type: ignore[no-untyped-def] + ) -> None: + """#332: when the recv loop hits ``KalshiBackpressureError`` it must + tear the session down — close the WS, flip ``_running=False``, and + clear the manager refs. Without this, a subsequent ``subscribe_*`` + on the same instance resurrects a recv loop on top of orphaned + server-side subscriptions whose queues are already closed. + """ + 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"], 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 + + # Fill the queue with the snapshot. + 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) + + # Overflow: the recv loop's queue.put raises KalshiBackpressureError. + 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) + + # #332: session is fully torn down inside the recv loop. + assert session._running is False + assert session._connection is None + assert session._sub_mgr is None + assert session._seq_tracker is None + assert session._orderbook_mgr is None + assert session._dispatcher is None + + async def test_issue_332_next_subscribe_after_backpressure_starts_fresh( + self, + fake_ws, # type: ignore[no-untyped-def] + test_auth, # type: ignore[no-untyped-def] + ) -> None: + """#332: after the recv loop tears down on backpressure, a follow-up + ``subscribe_*`` on the SAME ``_WebSocketSession`` must fail fast + with a clear ``RuntimeError`` rather than resurrecting a recv loop + on top of orphaned server-side subscriptions. The user is expected + to exit the ``async with`` block and start a fresh session for + recovery — which is also exercised here. + """ + from kalshi.errors import KalshiBackpressureError + + 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"], 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 + + # Snapshot fills the maxsize=1 queue, delta overflows. + 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) + + # Subscribing again on the torn-down session must NOT silently + # restart the recv loop. It must raise a clear error. + with pytest.raises(RuntimeError, match="not active"): + await session.subscribe_ticker(tickers=["T2"]) + + # And after exiting the failed session, the same KalshiWebSocket + # instance can be reused for a clean fresh session (#297 invariant + # holds even though we tore down from the recv loop, not _stop). + assert ws._connection is None + assert ws._running is False + assert ws._sub_mgr is None + + async with ws.connect() as session2: + assert session2._connection is not None + stream = await session2.subscribe_ticker(tickers=["T2"]) + assert session2._sub_mgr is not None + sid2 = next( + iter(session2._sub_mgr.active_subscriptions.values()) + ).server_sid + await fake_ws.send_to_all( + { + "type": "ticker", + "sid": sid2, + "msg": ticker_payload_dict( + market_ticker="T2", market_id="x", yes_bid_dollars="55" + ), + } + ) + msg = await asyncio.wait_for(stream.__anext__(), timeout=2.0) + assert msg.msg.market_ticker == "T2" + + # Drain the prior iterator's KalshiBackpressureError tail (the + # first session's stream still holds the snapshot + error; reading + # the error here keeps it from surfacing as an unhandled exception + # in test teardown). + _ = KalshiBackpressureError diff --git a/tests/ws/test_seq_correctness.py b/tests/ws/test_seq_correctness.py index 6bff3204..16ed8ab9 100644 --- a/tests/ws/test_seq_correctness.py +++ b/tests/ws/test_seq_correctness.py @@ -112,13 +112,19 @@ async def snapshot_landed() -> None: recv_task = session._recv_task assert recv_task is not None + # #332: the recv loop now clears manager refs on + # ``KalshiBackpressureError`` teardown (close the WS and reset + # state so the next subscribe doesn't resurrect a dead session). + # Snapshot the tracker BEFORE awaiting the task so we can still + # validate the #78 invariant on the now-detached object. + seq_tracker = session._seq_tracker 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, ( + assert seq_tracker.peek(sid) == 1, ( "#78 regression: ERROR-overflow advanced last_seq past a " "dropped message; orderbook is now silently desynced." ) From 1aebf3cf5b35b75ebe7c4fc0e3165293e4718063 Mon Sep 17 00:00:00 2001 From: Jeff West Date: Fri, 22 May 2026 08:34:37 -0500 Subject: [PATCH 2/3] polish(ws): drop dead KalshiBackpressureError reference; document _recv_task teardown invariant --- kalshi/ws/client.py | 3 +++ tests/ws/test_client.py | 9 +-------- 2 files changed, 4 insertions(+), 8 deletions(-) diff --git a/kalshi/ws/client.py b/kalshi/ws/client.py index 86456079..591890f7 100644 --- a/kalshi/ws/client.py +++ b/kalshi/ws/client.py @@ -983,6 +983,9 @@ async def run_forever(self, stop_event: asyncio.Event | None = None) -> None: self._recv_task.result() finally: await self._broadcast_sentinels() + # NOTE: _recv_task is intentionally NOT reset to None here — + # session lifecycle state is owned by ``_stop()`` (called from + # ``__aexit__``), which performs the full reset atomically. elif self._recv_task in done: # Server-side close / crash without stop event firing. The # recv loop's error paths broadcast sentinels themselves; diff --git a/tests/ws/test_client.py b/tests/ws/test_client.py index 29fae755..23750abd 100644 --- a/tests/ws/test_client.py +++ b/tests/ws/test_client.py @@ -851,8 +851,7 @@ async def test_issue_332_next_subscribe_after_backpressure_starts_fresh( to exit the ``async with`` block and start a fresh session for recovery — which is also exercised here. """ - from kalshi.errors import KalshiBackpressureError - + config = KalshiConfig(ws_base_url=fake_ws.url, timeout=5.0) config = KalshiConfig(ws_base_url=fake_ws.url, timeout=5.0) ws = KalshiWebSocket(auth=test_auth, config=config) async with ws.connect() as session: @@ -926,9 +925,3 @@ async def test_issue_332_next_subscribe_after_backpressure_starts_fresh( ) msg = await asyncio.wait_for(stream.__anext__(), timeout=2.0) assert msg.msg.market_ticker == "T2" - - # Drain the prior iterator's KalshiBackpressureError tail (the - # first session's stream still holds the snapshot + error; reading - # the error here keeps it from surfacing as an unhandled exception - # in test teardown). - _ = KalshiBackpressureError From 7ea2a23f6cd294d5c6f77b46a85ecfd11c52f492 Mon Sep 17 00:00:00 2001 From: Jeff West Date: Fri, 22 May 2026 08:43:32 -0500 Subject: [PATCH 3/3] fix(ws): wake _pause_recv_loop waiter on inline teardown; remove dup test line --- kalshi/ws/client.py | 28 ++++++++++++++++++++++++++-- tests/ws/test_client.py | 1 - 2 files changed, 26 insertions(+), 3 deletions(-) diff --git a/kalshi/ws/client.py b/kalshi/ws/client.py index 591890f7..bc3aaa37 100644 --- a/kalshi/ws/client.py +++ b/kalshi/ws/client.py @@ -361,8 +361,14 @@ async def _recv_loop(self) -> None: self._orderbook_mgr = None self._dispatcher = None self._pause_pending = False - self._pause_granted.clear() - self._resume_signal.clear() + # #332: wake any subscriber blocked on ``_pause_granted.wait()``; + # ``Event.clear()`` does NOT unblock existing waiters, so a + # pending ``_pause_recv_loop`` would deadlock here. ``set()`` + # releases the waiter; the post-pause guard in + # ``_do_subscribe`` re-checks ``_running`` so the woken + # subscriber sees the torn-down session and raises. + self._pause_granted.set() + self._resume_signal.set() break except (json.JSONDecodeError, ValidationError, KeyError): # #83: genuinely-non-fatal per-message errors. #241: when @@ -688,6 +694,24 @@ async def _do_subscribe( ) async with self._subscribe_lock: await self._pause_recv_loop() + # #332: ``_pause_recv_loop`` can return because the recv loop + # tore the session down on backpressure (inline teardown sets + # ``_pause_granted`` to release any waiter — ``Event.clear()`` + # alone does NOT unblock a coroutine already in ``.wait()``). + # Re-check before proceeding so a woken-up subscriber doesn't + # poke a dead session. + if ( + not self._running + or self._sub_mgr is None + or self._connection is None + ): + raise RuntimeError( + "KalshiWebSocket session is not active. The recv loop " + "tore down after a fatal error (e.g. " + "KalshiBackpressureError); exit the current `async with " + "ws.connect()` block and start a fresh session before " + "subscribing again." + ) try: sub = await self._sub_mgr.subscribe( channel, params=params, overflow=overflow, maxsize=maxsize, diff --git a/tests/ws/test_client.py b/tests/ws/test_client.py index 23750abd..ed09fe58 100644 --- a/tests/ws/test_client.py +++ b/tests/ws/test_client.py @@ -852,7 +852,6 @@ async def test_issue_332_next_subscribe_after_backpressure_starts_fresh( recovery — which is also exercised here. """ config = KalshiConfig(ws_base_url=fake_ws.url, timeout=5.0) - 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"], maxsize=1)