fix(ws): backpressure error in recv loop tears down connection cleanly

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;

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"

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."
             )