fix(ws): guard _start() against re-entry; clear managers on _stop (#297)

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.
 

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

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] = []