diff --git a/kalshi/ws/client.py b/kalshi/ws/client.py index 1140d37..bc3aaa3 100644 --- a/kalshi/ws/client.py +++ b/kalshi/ws/client.py @@ -337,10 +337,38 @@ 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 + # #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 @@ -650,10 +678,40 @@ 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() + # #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, @@ -949,6 +1007,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 775fec8..ed09fe5 100644 --- a/tests/ws/test_client.py +++ b/tests/ws/test_client.py @@ -766,3 +766,161 @@ 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. + """ + 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" diff --git a/tests/ws/test_seq_correctness.py b/tests/ws/test_seq_correctness.py index 6bff320..16ed8ab 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." )