From 8bcbb750ac481560328f460615bd82a69474f411 Mon Sep 17 00:00:00 2001 From: Jeff West Date: Wed, 20 May 2026 05:51:31 -0500 Subject: [PATCH 1/3] feat(ws): cooperative shutdown via run_forever(stop_event=...) (#177) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `KalshiWebSocket.run_forever()` gains an optional `stop_event: asyncio.Event | None = None` parameter. When the event fires — typically from a SIGINT handler via `add_signal_handler(SIGINT, stop.set)` — run_forever() clears `_running`, closes the connection, and awaits the recv loop's natural exit. The recv task is NOT cancelled, so no CancelledError leaks out. The mechanism: asyncio.wait races recv_task vs stop_event.wait(); if stop wins, _running=False then connection.close() makes the recv loop see ConnectionClosed on its next read and exit via the existing `if not self._running: break` branch — same teardown path as a deliberate _stop() call, just initiated externally. No behavior change when stop_event is omitted; external cancellation still propagates as before. The #175 missing-subscription guard remains the first check. Sibling foot-gun to #175 (silent-no-op-on-no-subscribe) — both target the same method but different failure modes. Together they make run_forever() actually usable as a long-running callback-mode loop: loud when misused, cooperative when terminated. --- CHANGELOG.md | 24 +++++++++++++++++++++ docs/websockets.md | 24 +++++++++++++++++++++ kalshi/ws/client.py | 48 +++++++++++++++++++++++++++++++++++++++-- tests/ws/test_client.py | 45 ++++++++++++++++++++++++++++++++++++++ 4 files changed, 139 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 5d1a05d..cbb9f2d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,30 @@ All notable changes to kalshi-sdk will be documented in this file. ## Unreleased +### WS `run_forever(stop_event=...)` cooperative shutdown (#177) + +`KalshiWebSocket.run_forever()` now accepts an optional +`stop_event: asyncio.Event | None = None` parameter. When set — typically +from a SIGINT handler via `add_signal_handler(SIGINT, stop.set)` — +`run_forever()` clears `_running`, closes the connection, and drains the +recv loop via its existing `not self._running` branch. The recv task is +NOT cancelled, so no `CancelledError` leaks out. + +```python +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) +``` + +No behavior change when `stop_event` is omitted — external cancellation +still propagates as before, and the #175 "missing subscription" guard +remains in place. + ### WS `run_forever()` raises on missing subscription (#175) `KalshiWebSocket.run_forever()` previously returned immediately when no diff --git a/docs/websockets.md b/docs/websockets.md index c3304ae..35a4b5a 100644 --- a/docs/websockets.md +++ b/docs/websockets.md @@ -128,6 +128,30 @@ async with ws.connect() as session: has landed in the session — a callback alone doesn't tell the server to send frames, and the previous silent-no-op behavior was a foot-gun (#175). +### Cooperative shutdown + +Pass an ``asyncio.Event`` to ``run_forever(stop_event=...)`` to terminate +the recv loop without raising ``CancelledError``. The canonical pattern +wires the event to ``SIGINT`` so Ctrl+C drains in-flight dispatches, +closes the WebSocket cleanly, and returns: + +```python +import asyncio +import 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) +``` + +When the event fires, ``run_forever()`` clears ``_running``, closes the +connection, and awaits the recv loop's natural exit. No ``CancelledError`` +leaks out (#177). Without ``stop_event``, external cancellation still +propagates as before. + `on()` works both before and after `connect()`; callbacks registered before the socket opens are buffered and applied when the session starts. diff --git a/kalshi/ws/client.py b/kalshi/ws/client.py index f76266c..b16951f 100644 --- a/kalshi/ws/client.py +++ b/kalshi/ws/client.py @@ -526,7 +526,7 @@ def decorator( return func return decorator - async def run_forever(self) -> None: + async def run_forever(self, stop_event: asyncio.Event | None = None) -> None: """Block until the recv loop terminates. Use with the callback API. Requires at least one prior ``subscribe_*`` (or generic @@ -537,6 +537,17 @@ async def run_forever(self) -> None: have explicitly subscribed to, so a callback without a matching subscribe sees nothing. + :param stop_event: optional ``asyncio.Event`` used to terminate + ``run_forever()`` cooperatively (#177). When set — typically + from a signal handler such as ``add_signal_handler(SIGINT, + stop.set)`` — this method clears ``_running``, closes the + connection, and drains the recv loop. The recv loop sees + ``ConnectionClosed`` on its next read and exits via the + normal ``not self._running`` branch, NOT via cancellation, + so no ``CancelledError`` leaks out. When ``None`` (the + default) the method blocks on ``_recv_task`` directly and + external cancellation still propagates as before. + :raises KalshiSubscriptionError: ``run_forever()`` was called before any ``subscribe_*`` request landed (formerly a silent no-op return — fixed in #175). @@ -550,7 +561,40 @@ async def run_forever(self) -> None: "@ws.on(channel) callback does not subscribe — the server " "only sends frames for channels you explicitly subscribe to." ) - await self._recv_task + if stop_event is None: + await self._recv_task + return + + # Race the recv loop against the stop event. Whichever finishes first + # wins; if stop wins, signal the recv loop to exit cleanly (via + # _running + connection.close()) so the loop's existing + # ConnectionClosed handler breaks instead of attempting reconnect. + # NOTE: we do NOT cancel _recv_task — cancellation would re-introduce + # the very CancelledError leak this hook exists to avoid (#177). + 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. Clearing _running first + # means the recv loop's ConnectionClosed branch hits the + # `if not self._running: break` path instead of reconnecting. + self._running = False + if self._connection is not None: + await self._connection.close() + await self._recv_task + # If recv_task finished first (server-side close, etc.) just + # surface any exception it raised — same semantics as the + # no-stop-event branch above. + elif self._recv_task in done: + self._recv_task.result() # re-raises if the task errored + finally: + if not stop_waiter.done(): + stop_waiter.cancel() + with contextlib.suppress(asyncio.CancelledError): + await stop_waiter # ------------------------------------------------------------------ # Orderbook convenience diff --git a/tests/ws/test_client.py b/tests/ws/test_client.py index 0c1bfa7..a1dd9c2 100644 --- a/tests/ws/test_client.py +++ b/tests/ws/test_client.py @@ -384,6 +384,51 @@ async def test_run_forever_without_subscription_raises( with pytest.raises(KalshiSubscriptionError, match="at least one active subscription"): await session.run_forever() + async def test_run_forever_with_stop_event_returns_cleanly( + self, + fake_ws, # type: ignore[no-untyped-def] + test_auth, # type: ignore[no-untyped-def] + ) -> None: + """#177: stop_event signals cooperative shutdown. run_forever() + closes the connection and drains the recv loop without raising + CancelledError. The loop exits via its existing `not _running` + branch on the next ConnectionClosed, NOT via cancellation.""" + 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_ticker(tickers=["T1"]) + stop = asyncio.Event() + run_task = asyncio.create_task(session.run_forever(stop_event=stop)) + # Give the loop a tick to settle on the recv await. + await asyncio.sleep(0.05) + assert not run_task.done() + # Trigger cooperative shutdown. + stop.set() + # Returns cleanly, no exception leaked. + await asyncio.wait_for(run_task, timeout=2.0) + assert run_task.exception() is None + assert session._connection is not None + assert session._connection.state == ConnectionState.CLOSED + + async def test_run_forever_with_pre_set_stop_event_returns_immediately( + self, + fake_ws, # type: ignore[no-untyped-def] + test_auth, # type: ignore[no-untyped-def] + ) -> None: + """#177: a stop_event already set before run_forever() runs should + fire on the first scheduling tick — same cooperative-shutdown path, + just no wait. Guards against the race-free case being broken by + future refactors of the asyncio.wait() arrangement.""" + 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_ticker(tickers=["T1"]) + stop = asyncio.Event() + stop.set() + await asyncio.wait_for(session.run_forever(stop_event=stop), timeout=2.0) + assert session._connection is not None + assert session._connection.state == ConnectionState.CLOSED + # --------------------------------------------------------------------------- # Error callback From ee0aeabd8d7f4cdac44e1d7b6b9d1fdfded7065c Mon Sep 17 00:00:00 2001 From: Jeff West Date: Wed, 20 May 2026 06:01:18 -0500 Subject: [PATCH 2/3] =?UTF-8?q?fixup:=20address=20review=20on=20PR=20#186?= =?UTF-8?q?=20=E2=80=94=20sentinel=20broadcast=20+=20exception=20safety?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Bot raised two real items: 1. Cooperative shutdown skipped _stop()'s put_sentinel broadcast. Today this is masked by __aexit__ -> _stop() running after run_forever() returns, but a caller using run_forever() outside `async with` would leave iterator consumers hanging on their queues. PR description claimed "same teardown path" — not quite true. Fixed: cooperative path now broadcasts sentinels too. 2. `await connection.close()` raising would leave _recv_task un-awaited. In practice close() shouldn't throw, but a nested try/finally costs nothing and makes the contract honest. Now structured so drain + sentinels happen regardless of close() outcome. self._running = False try: await self._connection.close() finally: try: await self._recv_task finally: # broadcast queue sentinels Added regression test (test_run_forever_with_stop_event_broadcasts_sentinels): subscribes to a channel, stops via stop_event, then asserts the iterator's `async for` terminates without further input. Pre-fix this hung; post-fix it drains cleanly. Also added docstring note about external cancellation of run_forever() itself while stop_event is provided: the cancellation cleans up stop_waiter but does NOT trigger cooperative shutdown — _recv_task keeps running until __aexit__ runs _stop() for full teardown. Use the event for graceful exit; rely on __aexit__ for hard cancellation. --- kalshi/ws/client.py | 30 +++++++++++++++++++++++++++--- tests/ws/test_client.py | 29 +++++++++++++++++++++++++++++ 2 files changed, 56 insertions(+), 3 deletions(-) diff --git a/kalshi/ws/client.py b/kalshi/ws/client.py index b16951f..6007f17 100644 --- a/kalshi/ws/client.py +++ b/kalshi/ws/client.py @@ -548,6 +548,15 @@ async def run_forever(self, stop_event: asyncio.Event | None = None) -> None: default) the method blocks on ``_recv_task`` directly and external cancellation still propagates as before. + External cancellation of ``run_forever()`` itself (e.g., + ``task.cancel()`` on the awaiting task) 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 the session's ``__aexit__`` calls ``_stop()`` for the + full teardown. Use the event for graceful exit; rely on + ``__aexit__`` for hard cancellation. + :raises KalshiSubscriptionError: ``run_forever()`` was called before any ``subscribe_*`` request landed (formerly a silent no-op return — fixed in #175). @@ -581,10 +590,25 @@ async def run_forever(self, stop_event: asyncio.Event | None = None) -> None: # Cooperative shutdown requested. Clearing _running first # means the recv loop's ConnectionClosed branch hits the # `if not self._running: break` path instead of reconnecting. + # + # Nested try/finally so a close() exception still drains + # the recv task AND broadcasts queue sentinels — otherwise + # iterators hang on consumers' `async for` even though the + # connection is gone. _stop() does the same sentinel + # broadcast on __aexit__; duplicated here so cooperative + # shutdown is complete even for callers who use + # run_forever() outside an `async with` block. self._running = False - if self._connection is not None: - await self._connection.close() - await self._recv_task + 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 recv_task finished first (server-side close, etc.) just # surface any exception it raised — same semantics as the # no-stop-event branch above. diff --git a/tests/ws/test_client.py b/tests/ws/test_client.py index a1dd9c2..cd87540 100644 --- a/tests/ws/test_client.py +++ b/tests/ws/test_client.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio +from typing import Any import pytest @@ -429,6 +430,34 @@ async def test_run_forever_with_pre_set_stop_event_returns_immediately( assert session._connection is not None assert session._connection.state == ConnectionState.CLOSED + async def test_run_forever_with_stop_event_broadcasts_sentinels( + self, + fake_ws, # type: ignore[no-untyped-def] + test_auth, # type: ignore[no-untyped-def] + ) -> None: + """#177 review fix: cooperative shutdown must broadcast queue + sentinels so iterator consumers exit `async for` cleanly. Without + this an iterator outside an `async with` block would hang on the + empty queue after the connection 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: + stream = await session.subscribe_ticker(tickers=["T1"]) + stop = asyncio.Event() + run_task = asyncio.create_task(session.run_forever(stop_event=stop)) + await asyncio.sleep(0.05) + stop.set() + await asyncio.wait_for(run_task, timeout=2.0) + + # The iterator must terminate without further input. Pre-fix it + # would block on the empty queue waiting for a sentinel that + # only _stop() (on __aexit__) was sending. + collected: list[Any] = [] + async def drain() -> None: + async for msg in stream: + collected.append(msg) + await asyncio.wait_for(drain(), timeout=2.0) + # --------------------------------------------------------------------------- # Error callback From 407a43a43268d3df48c2405e5ab4ae8841a86597 Mon Sep 17 00:00:00 2001 From: Jeff West Date: Wed, 20 May 2026 06:11:04 -0500 Subject: [PATCH 3/3] fixup: simultaneous-completion broadcast + _broadcast_sentinels helper Round-2 bot review on PR #186: 1. Simultaneous-completion edge case (medium). If server-side close lands in the same asyncio.wait() round as stop_event.set(), my previous condition `if stop_waiter in done and self._recv_task not in done` fell through to the `elif self._recv_task in done:` branch which re-raises but does NOT broadcast sentinels. Combined with the recv loop's clean `not self._running: break` exits also not broadcasting, iterators outside `async with` would hang. Verified by reading the recv loop: only its error paths broadcast sentinels; clean ConnectionClosed exits do not. Fix: stop_event fires now ALWAYS takes the cooperative-shutdown branch, regardless of recv_task state. If recv_task happens to be done already, we skip `await self._recv_task` and call `.result()` instead (re-raises any exception). Sentinel broadcast happens in the inner finally either way. 2. Duplicate sentinel-broadcast logic (low). Bot suggested extracting to a private helper. Done: `_broadcast_sentinels()` is now the single point where the loop runs. Six callers consolidated: _stop(), the new cooperative shutdown path, and 4 fatal-error paths in _recv_loop and _handle_reconnect. The helper's docstring documents the idempotent-put_sentinel invariant and which paths use it. 3. Sleep bump in test (nit). asyncio.sleep(0.05) -> 0.1 with a comment explaining the trade-off. Reduces flake risk on resource-constrained CI runners. 4. Empty-collected assertion in sentinel test (nit). The drain test now explicitly asserts `collected == []` to document what "clean termination" means here. Verification: 2117 unit tests pass. ruff + mypy clean. --- kalshi/ws/client.py | 78 ++++++++++++++++++++++++++--------------- tests/ws/test_client.py | 9 +++-- 2 files changed, 57 insertions(+), 30 deletions(-) diff --git a/kalshi/ws/client.py b/kalshi/ws/client.py index 6007f17..0a7a822 100644 --- a/kalshi/ws/client.py +++ b/kalshi/ws/client.py @@ -118,14 +118,28 @@ async def _stop(self) -> None: with contextlib.suppress(asyncio.CancelledError): await self._recv_task - # Send sentinels to all active queues - if self._sub_mgr: - for sub in self._sub_mgr.active_subscriptions.values(): - await sub.queue.put_sentinel() + await self._broadcast_sentinels() if self._connection: await self._connection.close() + async def _broadcast_sentinels(self) -> None: + """Put a shutdown sentinel on every active subscription queue. + + Iterator consumers see the sentinel and exit their ``async for`` + loops. ``MessageQueue.put_sentinel`` is idempotent — calling + twice puts two sentinels in the queue, the iterator reads the + first and raises ``StopAsyncIteration`` without ever touching + the second. + + Used by ``_stop()`` (post-cancel cleanup), the cooperative + shutdown branch of ``run_forever(stop_event=...)`` (#177), and + ``_recv_loop``'s fatal-error broadcast paths. + """ + if self._sub_mgr is not None: + for sub in self._sub_mgr.active_subscriptions.values(): + await sub.queue.put_sentinel() + def _ensure_recv_loop(self) -> None: """Start the recv_loop background task if not already running.""" if self._recv_task is None or self._recv_task.done(): @@ -187,9 +201,7 @@ async def _recv_loop(self) -> None: try: await inner except (KalshiBackpressureError, KalshiSubscriptionError): - if self._sub_mgr: - for sub in self._sub_mgr.active_subscriptions.values(): - await sub.queue.put_sentinel() + await self._broadcast_sentinels() except Exception: logger.debug( "Shielded dispatch raised during cancel cleanup", @@ -203,9 +215,7 @@ async def _recv_loop(self) -> None: logger.error( "Fatal WS error in recv loop", exc_info=True, ) - if self._sub_mgr: - for sub in self._sub_mgr.active_subscriptions.values(): - await sub.queue.put_sentinel() + await self._broadcast_sentinels() break except (json.JSONDecodeError, ValidationError, KeyError): # #83: genuinely-non-fatal per-message errors. Log with @@ -224,9 +234,7 @@ async def _recv_loop(self) -> None: "Unexpected error in recv loop; broadcasting sentinels", exc_info=True, ) - if self._sub_mgr: - for sub in self._sub_mgr.active_subscriptions.values(): - await sub.queue.put_sentinel() + await self._broadcast_sentinels() raise async def _process_frame(self, raw: str) -> None: @@ -321,9 +329,7 @@ async def _handle_reconnect(self) -> None: logger.error( "Reconnect failed: %s", reconnect_err, exc_info=True, ) - if self._sub_mgr: - for sub in self._sub_mgr.active_subscriptions.values(): - await sub.queue.put_sentinel() + await self._broadcast_sentinels() self._running = False async def _handle_seq_gap(self, gap: SequenceGap) -> None: @@ -586,10 +592,21 @@ async def run_forever(self, stop_event: asyncio.Event | None = None) -> None: {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. Clearing _running first - # means the recv loop's ConnectionClosed branch hits the - # `if not self._running: break` path instead of reconnecting. + if stop_waiter in done: + # Cooperative shutdown takes priority — broadcast sentinels + # at the end even if _recv_task happens to also be in + # `done` (server-side close racing the stop signal). The + # recv loop's clean `not self._running: break` exits do + # NOT broadcast on their own; only its error paths do. So + # the simultaneous-completion case (both tasks done in the + # same wait round) MUST go through this branch to avoid + # leaving iterator consumers hanging. + # + # Clearing _running first means a still-running recv loop's + # ConnectionClosed branch hits `if not self._running: + # break` instead of reconnecting. We do NOT cancel + # _recv_task — cancellation would re-introduce the very + # CancelledError leak this hook exists to avoid (#177). # # Nested try/finally so a close() exception still drains # the recv task AND broadcasts queue sentinels — otherwise @@ -604,16 +621,21 @@ async def run_forever(self, stop_event: asyncio.Event | None = None) -> None: await self._connection.close() finally: try: - await self._recv_task + 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: - if self._sub_mgr is not None: - for sub in self._sub_mgr.active_subscriptions.values(): - await sub.queue.put_sentinel() - # If recv_task finished first (server-side close, etc.) just - # surface any exception it raised — same semantics as the - # no-stop-event branch above. + await self._broadcast_sentinels() elif self._recv_task in done: - self._recv_task.result() # re-raises if the task errored + # Server-side close / crash without stop event firing. The + # recv loop's error paths broadcast sentinels themselves; + # its clean ConnectionClosed break does not, but in that + # case the user didn't request shutdown and __aexit__ will + # handle the teardown. Just re-surface any exception. + self._recv_task.result() finally: if not stop_waiter.done(): stop_waiter.cancel() diff --git a/tests/ws/test_client.py b/tests/ws/test_client.py index cd87540..9bc2b38 100644 --- a/tests/ws/test_client.py +++ b/tests/ws/test_client.py @@ -400,8 +400,10 @@ async def test_run_forever_with_stop_event_returns_cleanly( await session.subscribe_ticker(tickers=["T1"]) stop = asyncio.Event() run_task = asyncio.create_task(session.run_forever(stop_event=stop)) - # Give the loop a tick to settle on the recv await. - await asyncio.sleep(0.05) + # Give the loop a tick to settle on the recv await. 100 ms is + # generous enough for resource-constrained CI runners while + # still keeping the test fast. + await asyncio.sleep(0.1) assert not run_task.done() # Trigger cooperative shutdown. stop.set() @@ -457,6 +459,9 @@ async def drain() -> None: async for msg in stream: collected.append(msg) await asyncio.wait_for(drain(), timeout=2.0) + # No frames ever sent through fake_ws, so the iterator should + # have terminated immediately on the sentinel. + assert collected == [] # ---------------------------------------------------------------------------