diff --git a/kalshi/ws/client.py b/kalshi/ws/client.py index 5ac03623..a06657f9 100644 --- a/kalshi/ws/client.py +++ b/kalshi/ws/client.py @@ -123,29 +123,56 @@ def connect(self) -> _WebSocketSession: async def _start(self) -> None: """Connect and initialize managers. Does NOT start recv_loop yet.""" - 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 + 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." + ) + 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() + # 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 + self._orderbook_mgr = None + self._dispatcher = None + self._running = False + raise async def _stop(self) -> None: """Stop the receive loop and close the connection.""" @@ -172,6 +199,17 @@ async def _stop(self) -> None: if self._connection: await self._connection.close() + # Reset state AFTER awaiting close so teardown still has manager refs (#297). + 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..98307d9a 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,69 @@ 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" + + 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 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] = []