feat(ws): cooperative shutdown via run_forever(stop_event=...) (#177)

[TexasCoding]

Closes #177.

Sibling to the #175 fix: same method (run_forever()), different failure mode. #175 was "silent no-op on missing subscribe"; #177 is "no clean way to terminate."

The problem

async def run_forever(self) -> None:
    if self._recv_task is None:
        raise KalshiSubscriptionError(...)   # #175 guard
    await self._recv_task                    # blocks until recv loop exits

SIGINT (or any external cancellation) during await self._recv_task propagates CancelledError through both the run_forever caller AND the recv task itself, cancelling the loop mid-frame. The application sees CancelledError leak out and has no way to ask for a clean drain.

Fix (option B from the issue)

run_forever() gains an optional stop_event: asyncio.Event | None = None parameter. When set — typically from a SIGINT handler — run_forever() clears _running, closes the connection, and awaits the recv loop's natural exit via its existing if not self._running: break branch on ConnectionClosed. The recv task is NOT cancelled.

async def run_forever(self, stop_event: asyncio.Event | None = None) -> None:
    if self._recv_task is None:
        raise KalshiSubscriptionError(...)  # #175 guard unchanged

    if stop_event is None:
        await self._recv_task
        return

    stop_waiter = asyncio.create_task(stop_event.wait())
    try:
        done, _ = await asyncio.wait(
            {self._recv_task, stop_waiter},
            return_when=asyncio.FIRST_COMPLETED,
        )
        if stop_waiter in done and self._recv_task not in done:
            # Cooperative shutdown requested.
            self._running = False
            if self._connection is not None:
                await self._connection.close()
            await self._recv_task
        elif self._recv_task in done:
            self._recv_task.result()  # re-raise if errored
    finally:
        if not stop_waiter.done():
            stop_waiter.cancel()
            with contextlib.suppress(asyncio.CancelledError):
                await stop_waiter

Key invariant: the recv task is not cancelled. We close the connection, which makes the recv loop's existing ConnectionClosed handler see not self._running and break cleanly. Same teardown path as a deliberate _stop() call, just initiated externally instead of from the session's __aexit__.

Canonical SIGINT pattern (now in docs/websockets.md)

import asyncio, signal

stop = asyncio.Event()
asyncio.get_running_loop().add_signal_handler(signal.SIGINT, stop.set)

async with ws.connect() as session:
    await session.subscribe_ticker(tickers=["EXAMPLE-25-T"])
    await session.run_forever(stop_event=stop)

Ctrl+C → handler sets stop → run_forever() closes the WS gracefully → recv loop drains → method returns → async with runs _stop() for full teardown. No CancelledError, no half-dispatched frames, no orphaned tasks.

Tests

Two new in tests/ws/test_client.py::TestRunForever:

• test_run_forever_with_stop_event_returns_cleanly — subscribe, kick off run_forever(stop_event=stop) as task, sleep briefly to confirm it's blocking, stop.set(), assert clean return + ConnectionState.CLOSED + no exception on the task.
• test_run_forever_with_pre_set_stop_event_returns_immediately — stop event already set before the call; same cooperative-shutdown path, just no wait. Guards against future refactors of the asyncio.wait() arrangement breaking the race-free case.

The existing test_run_forever_blocks_until_close (no stop_event) and the #175 test_run_forever_without_subscription_raises are unchanged — confirms the new path is fully additive.

Verification

• uv run pytest tests/ --ignore=tests/integration → 2116 passed, 0 warnings.
• uv run ruff check . → clean.
• uv run mypy kalshi/ → Success: no issues found in 76 source files.

Acceptance criteria

• run_forever() can be terminated cooperatively without CancelledError leaking out.
• Test covering: stop-event set → run_forever() returns cleanly + connection is closed.
• docs/websockets.md callback section shows the canonical SIGINT + stop-event pattern.

(README WS section has no callback-mode example — only iterator-style — so no foot-gun to fix there. The cooperative-shutdown pattern lives in the full docs.)

Soft-breaking footprint

None. stop_event is an optional parameter; existing callers (await session.run_forever()) are unaffected. The recv-loop teardown path on stop_event reuses the existing _running + ConnectionClosed flow — no new code paths in the recv loop itself.

[claude]

Code Review — PR #186: cooperative shutdown via run_forever(stop_event=...)

Overview

Adds an optional stop_event: asyncio.Event | None = None parameter to run_forever(). When set, the method races _recv_task against the event via asyncio.wait, then — if the event fires first — sets _running = False, closes the connection, and drains the recv loop via its existing ConnectionClosed → not self._running → break path. The recv task is NOT cancelled, so no CancelledError leaks out.

The approach is sound. The asyncio.wait usage, the finally cleanup of stop_waiter, and the contextlib.suppress(CancelledError) are all correct. Tests and docs are thorough. Two issues below are worth addressing before merge.

---

Issues

1. _stop() divergence — missing queue sentinels (medium)

The cooperative shutdown path in run_forever manually replicates _stop() but skips its sentinel dispatch:

# _stop() does this — cooperative shutdown skips it:
if self._sub_mgr:
    for sub in self._sub_mgr.active_subscriptions.values():
        await sub.queue.put_sentinel()

Today this is harmless because __aexit__ calls _stop() after run_forever() returns, which sends the sentinels. But if _stop() gains new cleanup steps in the future, this path silently diverges. And the PR description says "same teardown path as _stop()" — which is not quite true.

Consider extracting the no-cancel teardown into a private helper so both paths stay in sync:

async def _teardown_recv(self) -> None:
    """Close the connection and drain the recv loop without cancelling it."""
    self._running = False
    if self._sub_mgr:
        for sub in self._sub_mgr.active_subscriptions.values():
            await sub.queue.put_sentinel()
    if self._connection is not None:
        await self._connection.close()
    if self._recv_task is not None:
        await self._recv_task

Or at minimum add a comment noting the intentional deviation from _stop() and why sentinels are deferred to __aexit__.

2. connection.close() exception leaves _recv_task un-awaited (low-medium)

In kalshi/ws/client.py, the cooperative shutdown block:

self._running = False
if self._connection is not None:
    await self._connection.close()   # if this raises…
await self._recv_task                # …this is skipped

If close() raises, _recv_task is left running with _running = False and a partially-closed connection. In practice a WebSocket close shouldn't throw, but a try/finally here costs nothing:

self._running = False
if self._connection is not None:
    try:
        await self._connection.close()
    finally:
        await self._recv_task

---

Minor notes

• asyncio.sleep(0.05) in test (test_run_forever_with_stop_event_returns_cleanly): This is a timing assertion. It's unlikely to flake in practice since asyncio is cooperative, but a asyncio.Event signalled from within the recv loop would be more robust if this ever becomes a CI pain point.

• External cancellation with stop_event set: If run_forever(stop_event=stop) is cancelled externally (CancelledError, not via the event), the finally block cleans up stop_waiter correctly, but _recv_task continues running. This is intentional, and __aexit__ → _stop() handles it. Worth adding one sentence to the docstring to set caller expectations, since it's a subtle invariant.

• Docstring length: The :param stop_event: block is thorough but quite long. Per project conventions, the "why" is non-obvious for async tear-down code, so the detail is defensible — just flagging for awareness.

---

Strengths

• The asyncio.wait race + stop_waiter cleanup in finally is correctly structured — no task leak on the happy path.
• Both done branches (stop_waiter wins vs. _recv_task wins) are handled, including the simultaneous-completion case (handled correctly by elif).
• Pre-set stop_event case is tested and works correctly.
• CHANGELOG and docs/websockets.md are clear and complete.
• Fully backward-compatible: existing await session.run_forever() callers are unaffected.
• 2116 tests pass, mypy clean, ruff clean.

The two issues above (sentinel divergence, exception safety around close()) are real but both have simple fixes. Happy to approve once those are addressed.

[TexasCoding]

Both real items addressed in ee0aeab:

1. Sentinel divergence — fixed. You're right that the PR description's "same teardown path" was imprecise: cooperative shutdown was skipping _stop()'s put_sentinel broadcast. Today that's masked by __aexit__ → _stop() running after run_forever() returns, but a caller using run_forever(stop_event=...) outside async with would leave iterator consumers hanging on their queues. The cooperative path now does the broadcast itself.

Sequencing is: _running = False → close() (recv loop sees ConnectionClosed, breaks via not self._running) → await self._recv_task (drain) → put_sentinel on each subscription queue. Matches _stop()'s post-recv-task-done ordering — sentinels arrive only after the recv loop has stopped emitting, so iterators don't see them prematurely.

I went with the inline duplication rather than extracting a shared helper because the cooperative path and _stop() diverge on the recv-task-cancel-vs-natural-drain decision, and the helper would have ended up with a cancel: bool flag that obscures more than it shares. A short comment cross-references the two paths.

2. close() exception safety — fixed. Nested try/finally so drain + sentinels happen regardless of close() outcome:

self._running = False
try:
    if self._connection is not None:
        await self._connection.close()
finally:
    try:
        await self._recv_task
    finally:
        if self._sub_mgr is not None:
            for sub in self._sub_mgr.active_subscriptions.values():
                await sub.queue.put_sentinel()

If close() raises, the exception still propagates (no suppress), but recv-task drain and sentinel broadcast happen on the way out. The await self._recv_task bound here is safe: websockets.close() either completes the handshake (recv loop's pending recv() returns ConnectionClosed) or hits its own close_timeout — either way _recv_task terminates.

Regression test — added test_run_forever_with_stop_event_broadcasts_sentinels: subscribes to a channel, kicks off run_forever(stop_event=stop), stops it, then asserts the iterator's async for terminates without further input. Pre-fix this hung; post-fix it drains cleanly.

Docstring note — added a paragraph explicitly covering the "external cancellation with stop_event provided" case: cancellation propagates, stop_waiter cleans up, but _recv_task keeps running until __aexit__ calls _stop(). "Use the event for graceful exit; rely on __aexit__ for hard cancellation."

Minor note on asyncio.sleep(0.05) in the existing test — kept as-is. The sleep is asserting "task hasn't completed yet"; an Event-based assertion would need the recv loop to signal back into the test, which complicates the test in service of robustness that hasn't actually been needed. Will revisit if it ever flakes on CI.

Verification: 2117 unit tests pass (+1 for the new sentinel regression test). mypy kalshi/ clean, ruff check . clean.

[claude]

Code Review — PR #186: run_forever(stop_event=...) cooperative shutdown

Overview

Solid approach. asyncio.wait() racing _recv_task against stop_waiter, paired with the existing _running + ConnectionClosed teardown path, is the right tool here. Backwards compatible, well-documented, and the third test (test_run_forever_with_stop_event_broadcasts_sentinels) catches the most dangerous failure mode (iterator consumers hanging after shutdown). A few things worth resolving or acknowledging before merge.

---

Issues

Medium — Simultaneous-completion case silently skips sentinel broadcast

If the server closes the connection at the exact moment stop_event fires, both tasks land in done in the same asyncio.wait() round. The condition:

if stop_waiter in done and self._recv_task not in done:
    # cooperative path — broadcasts sentinels
elif self._recv_task in done:
    self._recv_task.result()   # re-raises, no sentinels

…takes the elif branch — re-raises errors correctly, but does not broadcast sentinels. Whether this is a real problem depends on whether _recv_loop broadcasts sentinels itself in its normal ConnectionClosed exit paths. Scanning client.py, the sentinel broadcasts at lines 190–192, 206–208, 227–229 are inside exception handlers, but the clean not self._running: break exit at line 161 appears to skip all of them. If that's right, iterator consumers outside async with could still hang in this race.

Suggested fix — check the elif result first so you always land in the cooperative path when the stop event won the race, regardless of whether the recv task also happened to finish:

# Explicit priority: cooperative shutdown wins even if both fire at once.
if stop_waiter in done:
    if self._recv_task not in done:
        # Normal cooperative path — recv task still running.
        ...
    # Whether or not recv_task is also done, broadcast sentinels here.
    if self._sub_mgr is not None:
        for sub in self._sub_mgr.active_subscriptions.values():
            await sub.queue.put_sentinel()
elif self._recv_task in done:
    self._recv_task.result()

Alternatively, confirm that _recv_loop always broadcasts sentinels before exiting and document that explicitly.

---

Low — Sentinel broadcast duplicates _stop() logic

The loop:

if self._sub_mgr is not None:
    for sub in self._sub_mgr.active_subscriptions.values():
        await sub.queue.put_sentinel()

is a verbatim copy of what _stop() does (lines 122–124). _stop() can't be called here because it cancels _recv_task, which would reintroduce the CancelledError leak. That reasoning is sound, but the duplication means future additions to _stop()'s teardown won't automatically apply to the cooperative path.

A _broadcast_sentinels() private method would keep both callers in sync without introducing a dependency on the full _stop() logic. Not blocking, but worth a follow-up issue.

---

Low — Timing-dependent test assertion

await asyncio.sleep(0.05)
assert not run_task.done()

50 ms is fine on a loaded workstation but can be marginal in a resource-constrained CI runner. A more robust pattern is to inject a callback or event that fires when asyncio.wait() is reached — but that requires touching the production code, which may not be worth it for a test-only concern. At minimum, bumping to 100–150 ms costs nothing and reduces flake risk. Alternatively, document the 50 ms as a known flakiness risk so it's easy to triage if it ever surfaces.

---

Nit — drain() test doesn't assert final state

collected: list[Any] = []
async def drain() -> None:
    async for msg in stream:
        collected.append(msg)
await asyncio.wait_for(drain(), timeout=2.0)

The test correctly verifies the iterator terminates, but not asserting anything about collected (e.g., assert collected == []) makes the intent slightly opaque. Not a bug, but a one-liner assertion would document what "cleanly" means here.

---

What's done well

• asyncio.wait() usage is correct. Using a stop_waiter task (rather than wait_for or shield) gives precise control over which future wins and avoids inadvertently cancelling the recv task.
• Nested try/finally inside the cooperative branch ensures _recv_task is always drained and sentinels are always broadcast even if connection.close() raises.
• stop_waiter cleanup in the outer finally is correct — cancelling the waiter task prevents it from leaking if external cancellation fires while asyncio.wait() is blocked.
• Docstring is unusually thorough. The section on external cancellation while stop_event is provided is non-obvious and correctly described.
• docs/websockets.md and CHANGELOG are clean and cover the right context.
• Sentinel test (#177 review fix) is the most valuable of the three — it catches the concrete failure mode (hanging async for) that the feature exists to avoid.

---

Summary

The core logic is correct for the common case. The simultaneous-completion sentinel gap is the one thing worth resolving (or explicitly confirming is safe via _recv_loop always broadcasting sentinels before completing). Everything else is low-priority polish. Happy to approve once that edge case is addressed or ruled out.

[TexasCoding]

Round 2 addressed in 407a43a:

Simultaneous-completion (medium) — confirmed real and fixed. Verified by reading _recv_loop: only the error paths (KalshiBackpressureError, KalshiSubscriptionError, unexpected Exception) broadcast sentinels; the clean not self._running: break exits do not. So if both stop_waiter and _recv_task land in done together, my previous condition (stop_waiter in done and self._recv_task not in done) fell to the elif branch which re-raised but skipped sentinels.

Restructured so stop_waiter in done ALWAYS takes the cooperative-shutdown branch, regardless of _recv_task state:

if stop_waiter in done:
    self._running = False
    try:
        if self._connection is not None:
            await self._connection.close()
    finally:
        try:
            if not self._recv_task.done():
                await self._recv_task
            else:
                # Already done. .result() re-raises any exception,
                # surfaced as run_forever's exception.
                self._recv_task.result()
        finally:
            await self._broadcast_sentinels()
elif self._recv_task in done:
    # Server-side close / crash without stop event firing.
    self._recv_task.result()

Now the broadcast is unconditional once stop fires, and the .result() call on an already-done recv_task still surfaces its exception (if any) without re-awaiting.

_broadcast_sentinels() helper (low) — extracted. Reconsidered after the simultaneous-completion fix made the duplication grow to 6 call sites: _stop(), the cooperative shutdown branch, two paths in _recv_loop (cancel-cleanup KalshiBackpressureError/KalshiSubscriptionError, and the escape-hatch unexpected Exception), the recv-loop's KalshiBackpressureError/KalshiSubscriptionError non-cancel path, and _handle_reconnect()'s reconnect-failed path. The helper now centralises the loop with a docstring documenting the idempotent-put_sentinel invariant and which callers use it. Future cleanup additions land in one place.

Sleep bump (nit) — 0.05 → 0.1 with a comment about the trade-off (CI runner robustness vs. test speed).

Empty collected assertion (nit) — added. The drain test now explicitly asserts collected == [] to document what "clean termination" means here.

Not actioned:

External cancellation with stop_event set — already covered by the docstring paragraph added in the prior round. The bot acknowledged this above; no change needed.

Verification: 2117 unit tests pass. mypy kalshi/ clean, ruff check . clean.

[claude]

Code Review — PR #186: run_forever(stop_event=...) cooperative shutdown

Overall: solid implementation. The design is correct, the edge cases are handled, and the tests cover the important paths. A few style and coverage notes below.

---

What the PR does

• Adds stop_event: asyncio.Event | None = None to run_forever() — set the event to drain the recv loop cooperatively instead of leaking CancelledError.
• Extracts the repeated sentinel-broadcast idiom into _broadcast_sentinels() (5 call-sites → 1 helper). Clean DRY improvement.
• Adds 3 tests and updates CHANGELOG + docs.

---

Correctness

_stop() double-close is safe. After cooperative shutdown, __aexit__ calls _stop(), which calls connection.close() on an already-closed connection. ConnectionManager.close() sets _ws = None on first call and is effectively a no-op on subsequent calls — no issue here.

asyncio.wait() edge cases look right.

• Pre-set stop event fires on first scheduling tick — handled.
• Both tasks in done simultaneously — if stop_waiter in done: takes priority, as documented.
• _recv_task.result() in the elif branch correctly re-raises task exceptions.
• finally block always cancels stop_waiter — no orphaned tasks.

Sentinel ordering is correct. In _stop(): cancel task → sentinels → close connection. In the cooperative shutdown branch: set _running = False → close connection → drain recv → sentinels (in the innermost finally). The ordering difference is intentional and valid.

---

Issues

1. Docstring/comment verbosity violates CLAUDE.md

_broadcast_sentinels() has a multi-paragraph docstring. CLAUDE.md is explicit: "Never write multi-paragraph docstrings or multi-line comment blocks — one short line max." The idempotency note is non-obvious and worth keeping, but condense it to one line:

async def _broadcast_sentinels(self) -> None:
    # Idempotent: put_sentinel() twice is fine; iterators exit on the first.
    if self._sub_mgr is not None:
        for sub in self._sub_mgr.active_subscriptions.values():
            await sub.queue.put_sentinel()

The same issue applies to the long :param stop_event: docstring block in run_forever() (two full paragraphs explaining cancellation semantics). The PR description already explains the design; inline docs should stay minimal.

The lengthy inline # Cooperative shutdown takes priority... block comment inside the if stop_waiter in done: branch is also overly verbose for production code — again, design rationale belongs in the PR, not inline.

2. Missing test: external cancellation with stop_event provided

The docstring explicitly states:

External cancellation of run_forever() itself while stop_event is provided still propagates — the cancellation cleans up the internal stop_waiter task but does NOT trigger the cooperative shutdown branch. _recv_task keeps running until __aexit__ calls _stop().

This is a non-obvious invariant. There's currently no test asserting that task.cancel() on run_forever(stop_event=stop) still propagates CancelledError and leaves _recv_task running. Given the documented guarantee, this warrants a test case alongside the existing three.

3. Timing reliance in tests

asyncio.sleep(0.1) and asyncio.sleep(0.05) to "let the loop settle" are acknowledged in the comment, but these are the classic source of intermittent CI flakiness. A more robust pattern is to sync on a condition the session itself exposes — e.g. polling _recv_task is not None and not _recv_task.done() in a short loop, or adding a test-only hook. Low priority, but worth noting for a future pass.

---

Positive notes

• _broadcast_sentinels() extraction is exactly right — 5 identical 3-line blocks → 1 method, no behavior change.
• The try/finally nesting in the cooperative branch (close() exception still drains task AND broadcasts sentinels) is correct defensive async code.
• Test test_run_forever_with_stop_event_broadcasts_sentinels testing the sentinel-for-iterator-consumers path is a good catch.
• CHANGELOG entry and docs update are present and accurate.
• ConnectionState was already imported in the test file — no spurious new imports.

---

Summary

Required before merge: None — the implementation is correct as written.

Recommended:

1. Trim _broadcast_sentinels() docstring to one line per CLAUDE.md.
2. Trim the :param stop_event: docstring to the essential sentence.
3. Collapse the large block comment inside the cooperative-shutdown branch.
4. Add a test for external cancellation with stop_event provided.