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

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

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.
 

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:
@@ -526,7 +532,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 +543,26 @@ 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.
+
+            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).
@@ -550,7 +576,71 @@ 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:
+                # 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
+                # 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
+                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. 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()
+                with contextlib.suppress(asyncio.CancelledError):
+                    await stop_waiter
 
     # ------------------------------------------------------------------
     # Orderbook convenience

tests/ws/test_client.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import asyncio
+from typing import Any
 
 import pytest
 
@@ -384,6 +385,84 @@ 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. 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()
+            # 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
+
+    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)
+            # No frames ever sent through fake_ws, so the iterator should
+            # have terminated immediately on the sentinel.
+            assert collected == []
+
 
 # ---------------------------------------------------------------------------
 # Error callback