From f64b49912a4a188100901e6812e9d471181f1372 Mon Sep 17 00:00:00 2001 From: Jeff West Date: Thu, 21 May 2026 18:24:35 -0500 Subject: [PATCH 1/3] fix(ws): guard KalshiWebSocket._start() against re-entry; clear managers on _stop Re-entering an active session (nested connect() or restart-on-error with a cached instance) used to silently rebuild _connection, _sub_mgr, _seq_tracker, _orderbook_mgr, and _dispatcher, orphaning the outer session's subscriptions and recv task with no error. _start() now refuses re-entry while _connection/_running are live, with a message that points at the two legitimate paths (await the existing async-with, or build a fresh KalshiWebSocket()). _stop() now nulls _connection / _sub_mgr / _seq_tracker / _orderbook_mgr / _dispatcher / _recv_task and clears the pause/resume events AFTER awaiting connection close, so the same instance can be reused for a fresh connect() once the previous session has fully exited. Tests: test_double_start_raises pins the overlap rejection; test_reconnect_after_stop_works pins that clean restart still works (subscribe + receive a frame on the reused instance). Updated test_close_sets_state and the integration shutdown test to match the new clear-on-stop contract. Closes #297. --- kalshi/ws/client.py | 21 ++++++++++++++++ tests/ws/test_client.py | 49 ++++++++++++++++++++++++++++++++++-- tests/ws/test_integration.py | 6 ++--- 3 files changed, 71 insertions(+), 5 deletions(-) diff --git a/kalshi/ws/client.py b/kalshi/ws/client.py index 5ac03623..5e246074 100644 --- a/kalshi/ws/client.py +++ b/kalshi/ws/client.py @@ -123,6 +123,13 @@ def connect(self) -> _WebSocketSession: async def _start(self) -> None: """Connect and initialize managers. Does NOT start recv_loop yet.""" + if self._connection is not None or self._running: + raise RuntimeError( + "KalshiWebSocket session is already started. Each instance " + "supports one active session at a time — await the existing " + "`async with ws.connect()` to exit, or create a fresh " + "KalshiWebSocket() for a new session." + ) self._connection = ConnectionManager( auth=self._auth, config=self._config, @@ -172,6 +179,20 @@ async def _stop(self) -> None: if self._connection: await self._connection.close() + # #297: clear instance state so the same KalshiWebSocket() can be + # reused for a fresh `connect()` after a clean `__aexit__`. Done + # after close so any final teardown above still has access to the + # managers via self.* attributes. + self._connection = None + self._sub_mgr = None + self._seq_tracker = None + self._orderbook_mgr = None + self._dispatcher = None + self._recv_task = None + self._pause_pending = False + self._pause_granted.clear() + self._resume_signal.clear() + async def _broadcast_sentinels(self) -> None: """Put a shutdown sentinel on every active subscription queue. diff --git a/tests/ws/test_client.py b/tests/ws/test_client.py index 081955d3..5b51e382 100644 --- a/tests/ws/test_client.py +++ b/tests/ws/test_client.py @@ -31,12 +31,14 @@ async def test_connect_and_close(self, fake_ws, test_auth) -> None: # type: ign assert session._connection.state == ConnectionState.CONNECTED async def test_close_sets_state(self, fake_ws, test_auth) -> None: # type: ignore[no-untyped-def] + """#297: after a clean __aexit__, _connection (and the other managers) + are cleared so the same instance can be reused for a fresh connect().""" config = KalshiConfig(ws_base_url=fake_ws.url, timeout=5.0) ws = KalshiWebSocket(auth=test_auth, config=config) async with ws.connect(): pass - assert ws._connection is not None - assert ws._connection.state == ConnectionState.CLOSED + assert ws._connection is None + assert ws._running is False async def test_connect_returns_session(self, fake_ws, test_auth) -> None: # type: ignore[no-untyped-def] config = KalshiConfig(ws_base_url=fake_ws.url, timeout=5.0) @@ -57,6 +59,49 @@ async def on_state(old: ConnectionState, new: ConnectionState) -> None: # Should have at least DISCONNECTED->CONNECTING and CONNECTING->CONNECTED assert any(new == ConnectionState.CONNECTED for _, new in states) + async def test_double_start_raises(self, fake_ws, test_auth) -> None: # type: ignore[no-untyped-def] + """#297: re-entering connect() on an active session must fail loudly + rather than silently clobbering the managers.""" + config = KalshiConfig(ws_base_url=fake_ws.url, timeout=5.0) + ws = KalshiWebSocket(auth=test_auth, config=config) + async with ws.connect(): + with pytest.raises(RuntimeError, match="already started"): + async with ws.connect(): + pass + + async def test_reconnect_after_stop_works(self, fake_ws, test_auth) -> None: # type: ignore[no-untyped-def] + """#297: after a clean __aexit__, the same instance can be reused + for a fresh connect() — the guard rejects overlap, not legitimate + restart.""" + config = KalshiConfig(ws_base_url=fake_ws.url, timeout=5.0) + ws = KalshiWebSocket(auth=test_auth, config=config) + + async with ws.connect() as session: + assert session._connection is not None + # State cleared after exit + assert ws._connection is None + assert ws._sub_mgr is None + assert ws._dispatcher is None + assert ws._recv_task is None + assert ws._running is False + + # Reuse same instance: subscribe and receive a frame. + async with ws.connect() as session: + assert session._connection is not None + assert session._connection.state == ConnectionState.CONNECTED + stream = await session.subscribe_ticker(tickers=["T1"]) + await fake_ws.send_to_all( + { + "type": "ticker", + "sid": 1, + "msg": ticker_payload_dict( + market_ticker="T1", market_id="x", yes_bid_dollars="55" + ), + } + ) + msg = await asyncio.wait_for(stream.__anext__(), timeout=2.0) + assert msg.msg.market_ticker == "T1" + # --------------------------------------------------------------------------- # Typed subscribe methods diff --git a/tests/ws/test_integration.py b/tests/ws/test_integration.py index c540e737..75d6f4f2 100644 --- a/tests/ws/test_integration.py +++ b/tests/ws/test_integration.py @@ -591,9 +591,9 @@ async def test_graceful_shutdown_stops_iterators( assert msg.msg.yes_bid == 42 # Context manager exit happens here - # After exit, the connection should be closed - assert ws._connection is not None - assert ws._connection.state == ConnectionState.CLOSED + # After exit, #297: managers are cleared so the instance can be reused. + assert ws._connection is None + assert ws._running is False # Sentinel was pushed; iteration should stop immediately items: list[object] = [] From 288d78b07ceee260c29b9cecb2392b9b5afe9b1a Mon Sep 17 00:00:00 2001 From: Jeff West Date: Thu, 21 May 2026 18:41:59 -0500 Subject: [PATCH 2/3] fix-up(ws): reset state on _start failure; pin regression test; trim comment - Wrap _start() body in try/except BaseException that clears _connection, _sub_mgr, _seq_tracker, _orderbook_mgr, _dispatcher, and _running on any failure. Without this, a raise in ConnectionManager.connect() (auth/network) left _connection set while _running stayed False; __aexit__ doesn't run when __aenter__ raises, so the next connect() permanently tripped the 'already started' guard. - Add test_failed_connect_allows_retry pointing at an unreachable port to lock in the state-clear invariant. - Collapse the 4-line #297 cleanup comment in _stop() to a single line preserving the ordering WHY. --- kalshi/ws/client.py | 60 +++++++++++++++++++++++------------------ tests/ws/test_client.py | 20 ++++++++++++++ 2 files changed, 54 insertions(+), 26 deletions(-) diff --git a/kalshi/ws/client.py b/kalshi/ws/client.py index 5e246074..8e17617c 100644 --- a/kalshi/ws/client.py +++ b/kalshi/ws/client.py @@ -130,29 +130,40 @@ async def _start(self) -> None: "`async with ws.connect()` to exit, or create a fresh " "KalshiWebSocket() for a new session." ) - self._connection = ConnectionManager( - auth=self._auth, - config=self._config, - heartbeat_timeout=self._heartbeat_timeout, - on_state_change=self._on_state_change, - ) - await self._connection.connect() - - self._sub_mgr = SubscriptionManager(self._connection, json_loads=self._json_loads) - self._seq_tracker = SequenceTracker(on_gap=self._handle_seq_gap) - self._orderbook_mgr = OrderbookManager() - self._dispatcher = MessageDispatcher( - sub_mgr=self._sub_mgr, - on_error=self._on_error, - seq_tracker=self._seq_tracker, - orderbook_mgr=self._orderbook_mgr, - ) - self._running = True + try: + self._connection = ConnectionManager( + auth=self._auth, + config=self._config, + heartbeat_timeout=self._heartbeat_timeout, + on_state_change=self._on_state_change, + ) + await self._connection.connect() + + self._sub_mgr = SubscriptionManager(self._connection, json_loads=self._json_loads) + self._seq_tracker = SequenceTracker(on_gap=self._handle_seq_gap) + self._orderbook_mgr = OrderbookManager() + self._dispatcher = MessageDispatcher( + sub_mgr=self._sub_mgr, + on_error=self._on_error, + seq_tracker=self._seq_tracker, + orderbook_mgr=self._orderbook_mgr, + ) + self._running = True - # Register any callbacks that were buffered before connect() - for channel, func in self._pending_callbacks: - self._dispatcher.register_callback(channel, func) - self._pending_callbacks.clear() + # Register any callbacks that were buffered before connect() + for channel, func in self._pending_callbacks: + self._dispatcher.register_callback(channel, func) + self._pending_callbacks.clear() + except BaseException: + # #297: partial-failure cleanup. __aexit__ won't run if __aenter__ + # raises, so reset state here to keep the instance reusable. + self._connection = None + self._sub_mgr = None + self._seq_tracker = None + self._orderbook_mgr = None + self._dispatcher = None + self._running = False + raise async def _stop(self) -> None: """Stop the receive loop and close the connection.""" @@ -179,10 +190,7 @@ async def _stop(self) -> None: if self._connection: await self._connection.close() - # #297: clear instance state so the same KalshiWebSocket() can be - # reused for a fresh `connect()` after a clean `__aexit__`. Done - # after close so any final teardown above still has access to the - # managers via self.* attributes. + # Reset state AFTER awaiting close so teardown still has manager refs (#297). self._connection = None self._sub_mgr = None self._seq_tracker = None diff --git a/tests/ws/test_client.py b/tests/ws/test_client.py index 5b51e382..98307d9a 100644 --- a/tests/ws/test_client.py +++ b/tests/ws/test_client.py @@ -102,6 +102,26 @@ async def test_reconnect_after_stop_works(self, fake_ws, test_auth) -> None: # msg = await asyncio.wait_for(stream.__anext__(), timeout=2.0) assert msg.msg.market_ticker == "T1" + async def test_failed_connect_allows_retry(self, test_auth) -> None: # type: ignore[no-untyped-def] + """#297: a failed connect() must not brick the instance. + + ``__aexit__`` does not run if ``__aenter__`` raises, so ``_start()`` + is responsible for resetting any partially-assigned managers when + ``ConnectionManager.connect()`` blows up. Otherwise the guard at the + top of ``_start()`` permanently rejects every subsequent attempt. + """ + # Unreachable port -> ConnectionManager.connect() will raise. + config = KalshiConfig(ws_base_url="ws://127.0.0.1:1", timeout=1.0) + ws = KalshiWebSocket(auth=test_auth, config=config) + with pytest.raises(Exception): # noqa: B017 — any connect failure must clear state + async with ws.connect(): + pass + # Core invariant: state cleared so the guard doesn't trip on retry. + assert ws._connection is None + assert ws._sub_mgr is None + assert ws._dispatcher is None + assert ws._running is False + # --------------------------------------------------------------------------- # Typed subscribe methods From 97955020b18aba0b0966685a746c8399296c50d4 Mon Sep 17 00:00:00 2001 From: Jeff West Date: Thu, 21 May 2026 19:01:14 -0500 Subject: [PATCH 3/3] fix-up(ws): defensive close on partial _start failure; document callback retry path - Close underlying ConnectionManager in the except block when connect() succeeded but a later step raised, instead of dropping the reference to None and relying on GC. - Document _pending_callbacks retry semantics: on mid-loop failure the unregistered tail replays on the next _start; already-registered entries are discarded with the failed dispatcher. Addresses round-2 review on #307. --- kalshi/ws/client.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/kalshi/ws/client.py b/kalshi/ws/client.py index 8e17617c..a06657f9 100644 --- a/kalshi/ws/client.py +++ b/kalshi/ws/client.py @@ -151,12 +151,21 @@ async def _start(self) -> None: self._running = True # Register any callbacks that were buffered before connect() + # On failure mid-loop, _pending_callbacks is intentionally NOT cleared: + # any unregistered tail replays on the next _start (already-registered + # entries are discarded with the failed dispatcher). for channel, func in self._pending_callbacks: self._dispatcher.register_callback(channel, func) self._pending_callbacks.clear() except BaseException: # #297: partial-failure cleanup. __aexit__ won't run if __aenter__ # raises, so reset state here to keep the instance reusable. + # Belt-and-suspenders close: if connect() succeeded but a subsequent + # step raised, drop the underlying socket explicitly rather than + # relying on GC of the orphaned ConnectionManager. + 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