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

[TexasCoding]

Summary

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 (a documented recovery path; 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 PR 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/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 self._sub_mgr is not None), 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.

Issues closed

• Closes KalshiBackpressureError in recv loop breaks the loop but leaves the WS open and _running=True #332 — KalshiBackpressureError in recv loop now closes the WS connection and clears _running + manager state so the next subscribe either starts a fresh session or fails fast

Behavioral changes

• After a KalshiBackpressureError / KalshiSubscriptionError in the recv loop, the session is torn down: _connection/_running/manager refs are cleared and the WS is closed. Previously the same KalshiWebSocket session would silently accept further subscribe_* calls that orphaned server-side state.
• _do_subscribe now raises RuntimeError("KalshiWebSocket session is not active. …") instead of AssertionError when called on a session whose recv loop has torn down. Callers must exit the async with block and re-enter for a fresh session.

Tests

• tests/ws/test_client.py::TestIssue332BackpressureTeardown::test_issue_332_backpressure_closes_ws_cleanly — asserts the recv loop fully clears connection + manager refs and flips _running=False after overflow.
• tests/ws/test_client.py::TestIssue332BackpressureTeardown::test_issue_332_next_subscribe_after_backpressure_starts_fresh — asserts a follow-up subscribe_* on the torn-down session raises RuntimeError, and that the same instance can be reused for a fresh session that successfully streams a ticker frame.
• tests/ws/test_seq_correctness.py::TestBackpressureDoesNotAdvanceSeq::test_dropped_delta_keeps_seq_watermark_unchanged — consequential update only: snapshot the _seq_tracker reference before awaiting the recv task, since the tracker is now cleared as part of the teardown. The #78 invariant assertion is preserved on the detached object.

Full tests/ws/ (374 tests) pass, plus ruff and mypy clean.

Source

Round-3 independent audit closure plan, wave W2 (medium-severity correctness).

[claude]

Code Review — PR #367: fix(ws): backpressure error in recv loop tears down connection cleanly

Overview

Solid, targeted fix. When the recv loop hits KalshiBackpressureError / KalshiSubscriptionError it now fully mirrors the _stop() teardown: broadcasts sentinels, flips _running=False, closes the connection (suppressing close failures), and clears all manager refs. _do_subscribe gets a proper RuntimeError guard in place of the bare assert. The #297 reuse invariant is preserved — the same KalshiWebSocket instance works again after exiting and re-entering async with ws.connect().

---

Strengths

• Correct teardown sequence: sentinels first, _running=False, then connection close, then ref nulling — matches _stop() exactly and is safe.
• contextlib already imported: no new imports needed; contextlib.suppress(Exception) is the right tool for cleanup-path close failures.
• RuntimeError replaces bare assert: asserts are for invariants that "can't happen"; this is a real boundary that callers can hit. The error message is clear and actionable.
• _stop() is safe after recv-loop teardown: since _connection is None and _recv_task.done() is True, all the _stop() branches short-circuit cleanly — no double-close, no second sentinel broadcast.
• Test coverage: teardown assertions, fail-fast guard, and the full reuse cycle (async with → fail → exit → new async with → subscribe → stream) are all exercised.

---

Issues

1. Dead code / misleading comment in test (minor correctness concern)

In test_issue_332_next_subscribe_after_backpressure_starts_fresh, the final two lines are:

# 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

_ = KalshiBackpressureError binds the class to _; it does not consume anything from an iterator or drain any queue. The comment describes intent the code doesn't fulfil. If the goal is to prevent an unhandled-exception warning at test teardown, the actual fix would be to explicitly drain stream (the iterator from the first session) until it raises or exhausts — something like:

with contextlib.suppress(KalshiBackpressureError, StopAsyncIteration):
    await stream.__anext__()  # consume the buffered snapshot or the error sentinel

As written, the line is dead code and the comment could mislead future readers into thinking it has an effect.

2. asyncio.sleep(0.1) timing dependence in both new tests

Both new tests use await asyncio.sleep(0.1) to let the recv loop process the snapshot before sending the delta. This works in practice but is a timing assumption that can flip on a heavily loaded CI runner. A more robust pattern is to poll for the expected queue state (e.g., while queue.qsize() < 1: await asyncio.sleep(0)) or use a shared asyncio.Event that the test fixture sets when a frame arrives.

3. _recv_task not cleared during recv-loop teardown (harmless, but worth a note)

_stop() sets self._recv_task = None as part of its cleanup (line 208). The recv-loop teardown path intentionally skips this — correctly, because you can't cancel or await yourself. The omission is safe: _do_subscribe raises RuntimeError before _ensure_recv_loop could create a duplicate task, and _stop() handles a done task gracefully. But the existing comment only explains skipping the await; adding a one-liner like # _recv_task intentionally not cleared here — _stop() does it on __aexit__ would save the next reader from having to trace through the same logic.

4. TOCTOU window in _do_subscribe guard (pre-existing, not introduced here)

The new guard:

if not self._running or self._sub_mgr is None or self._connection is None:
    raise RuntimeError(...)
async with self._subscribe_lock:
    ...

The check happens outside the lock. In asyncio cooperative scheduling, a task switch at the await self._broadcast_sentinels() yield point in the teardown could allow another task to pass the guard before _running is flipped. The window is extremely narrow and the overall risk is low. This is a pre-existing design limitation (the old assert had the same gap), so no blocker — but noting it here for awareness.

---

Summary

The core fix is correct and well-executed. Two things worth addressing before merge:

1. Fix or remove the _ = KalshiBackpressureError line — it does nothing and the misleading comment is a maintenance hazard.
2. Consider annotating _recv_task not being cleared in the teardown block for future readers.

The timing-dependent sleep(0.1) in tests is low priority but worth a follow-up issue if the WS test suite sees intermittent failures in CI.

[TexasCoding]

Round-2 review addressed in 1aebf3c:

Addressed

• Dropped the dead _ = KalshiBackpressureError line (and its misleading drain-the-tail comment) in tests/ws/test_client.py; the only remaining import of the symbol in that test went with it, so ruff/pyflakes now police usage naturally.
• Added a comment in run_forever()'s cooperative-shutdown branch documenting that _recv_task is intentionally NOT reset to None there — session lifecycle reset is owned by _stop() via __aexit__.

Deferred (low-priority follow-up)

• Replacing the await asyncio.sleep(0.1) test timing in the KalshiBackpressureError in recv loop breaks the loop but leaves the WS open and _running=True #332 backpressure tests with deterministic synchronization. Not blocking; tracked for a later polish pass.

Verified: ruff check + mypy clean on kalshi/ws/client.py; pytest tests/ws/test_client.py -k TestIssue332 → 2 passed.

[claude]

Code Review — PR #367: fix(ws): backpressure error in recv loop tears down connection cleanly

Overview

This PR fixes a real correctness bug (#332): after KalshiBackpressureError in _recv_loop, the connection was left open with _running=True and managers intact, allowing a subsequent subscribe_* to resurrect a recv loop on top of orphaned server-side subscriptions. The fix mirrors the _stop() teardown pattern inline, and replaces assert self._sub_mgr is not None with a fast-failing RuntimeError. The approach is correct and the scope is surgical — exactly what the problem requires.

---

Issues

Medium — Potential deadlock in _pause_recv_loop during concurrent subscribe + backpressure

File: kalshi/ws/client.py, teardown block (~line 353)

The teardown does:

self._pause_pending = False
self._pause_granted.clear()   # ← does NOT wake waiters
self._resume_signal.clear()

asyncio.Event.clear() does not unblock a coroutine already blocked on .wait(). If _do_subscribe has already passed the recv_task.done() check in _pause_recv_loop, set _pause_pending = True, and is now awaiting self._pause_granted.wait() — and the recv loop crashes on backpressure before reaching the _pause_pending check at the top of its next iteration — _pause_recv_loop will wait forever.

The existing _stop() avoids this because it cancels the recv task first, which propagates CancelledError into any pending await in the recv loop; the inline teardown has no equivalent unblock mechanism.

Suggested fix — change clear() to set() in the teardown's pause-state reset:

self._pause_pending = False
self._pause_granted.set()    # wake any _pause_recv_loop waiter so it can observe _running=False
self._resume_signal.clear()

And add a post-pause guard inside _do_subscribe so a woken-up subscriber doesn't proceed on a dead session:

async with self._subscribe_lock:
    await self._pause_recv_loop()
    if not self._running or self._sub_mgr is None or self._connection is None:
        raise RuntimeError("KalshiWebSocket session is not active …")
    try:
        ...

This is an inherently narrow race (concurrent subscribe + overflow), but it's reachable and would produce an infinite hang rather than a clear error.

---

Minor Issues

Minor — Duplicate assignment in test

File: tests/ws/test_client.py, test_issue_332_next_subscribe_after_backpressure_starts_fresh (~line 853)

config = KalshiConfig(ws_base_url=fake_ws.url, timeout=5.0)
config = KalshiConfig(ws_base_url=fake_ws.url, timeout=5.0)   # ← duplicate; dead code

Remove the second line.

---

What's Good

• Teardown order is correct: sentinels are broadcast before state is cleared, so consumers see the failure before managers disappear. _stop() follows the same ordering.
• _stop() double-clear is safe: after the recv loop clears managers, _stop() (called from __aexit__) re-clears the same fields — all are already None/False, so it's a clean no-op. No interaction risk.
• contextlib already imported: no missing import.
• RuntimeError over AssertionError: the new guard message ("exit the current async with block…") gives callers actionable recovery guidance. Correct and helpful.
• _recv_task intentionally not cleared here: the comment in run_forever explains why lifetime is owned by _stop(). This is consistent with the existing design.
• test_seq_correctness.py adaptation: snapshotting _seq_tracker before awaiting the recv task is exactly right — validates the WS backpressure: ERROR-overflow advances last_seq, silently desyncing orderbook #78 invariant on the detached object rather than racing with the teardown.
• Test coverage: the two new test cases cover the core paths well — teardown state assertions, fail-fast on re-subscribe, and the reuse-after-exit invariant from bug(ws): KalshiWebSocket._start() silently clobbers managers on repeat/nested connect() #297.

---

Summary

The core fix is correct and the tests validate the important behaviors. One real (if narrow) deadlock risk exists in the interaction between _pause_recv_loop and the new inline teardown — worth addressing before merge. The duplicate config line in the test is trivial.

[TexasCoding]

Round-2 review addressed in 7ea2a23: closed the _pause_recv_loop deadlock race (set() not clear() in inline teardown + post-pause subscriber guard re-checks _running/_sub_mgr/_connection so a woken waiter raises RuntimeError instead of poking a dead session); removed duplicate config = KalshiConfig(...) line in test_issue_332_next_subscribe_after_backpressure_starts_fresh. Ruff/mypy clean; TestIssue332 tests pass. Regression test for the deadlock race omitted — recv loop's normal pause-grant path makes it non-deterministic to schedule a waiter that survives until inline teardown; invariant is documented inline.

[claude]

Code Review — PR #367 (Round 3)

Status after Round 2 fixes

All round-2 feedback has been addressed correctly. Quick verification:

• Deadlock fix — inline teardown now calls _pause_granted.set() (not .clear()), correctly waking any subscriber blocked in _pause_recv_loop. ✓
• Post-pause guard — _do_subscribe re-checks _running / _sub_mgr / _connection after _pause_recv_loop returns and raises RuntimeError if torn down. ✓
• Duplicate config line removed from test_issue_332_next_subscribe_after_backpressure_starts_fresh. ✓
• Dead _ = KalshiBackpressureError line dropped. ✓
• run_forever comment explaining _recv_task lifecycle ownership added. ✓

---

Final correctness pass

The critical race is handled correctly. When _do_subscribe is blocked in _pause_recv_loop and the recv loop simultaneously hits KalshiBackpressureError in _process_frame:

1. Teardown calls _pause_granted.set() → unblocks the subscriber's _pause_granted.wait()
2. _pause_recv_loop returns; _pause_pending is already False (teardown cleared it)
3. Post-pause guard catches _running=False → raises RuntimeError
4. _ensure_recv_loop() is not reached (raise is before the try: block) — no spurious task creation on a dead session

_stop() called from __aexit__ after recv-loop teardown is clean:

• _recv_task.done() is True → await/cancel block skipped
• _broadcast_sentinels() is a no-op (_sub_mgr is None)
• _connection is None → close skipped
• _pause_granted.clear() and _resume_signal.clear() restore events to initial state for the next async with ws.connect()

No double-close, no double-sentinel, no leaked task. ✓

_ensure_recv_loop guard is adequate. It checks self._recv_task is None or self._recv_task.done() at line 242, so if somehow called on a done-but-not-None task it would create a new task that exits immediately because _running=False. The paths that could reach this are already blocked by the RuntimeError guards.

---

Remaining low-priority item

The asyncio.sleep(0.1) timing in both new tests is the only deferred item. Correctly tracked and not a blocker.

---

Verdict

LGTM. The implementation is correct, the teardown sequence is safe, the _stop()-after-teardown path is clean, and the tests cover the failure path, the fail-fast guard, and the reuse invariant. Ready to merge.