diff --git a/kalshi/_base_client.py b/kalshi/_base_client.py index 0a33c5f..e2965e2 100644 --- a/kalshi/_base_client.py +++ b/kalshi/_base_client.py @@ -144,7 +144,11 @@ def request( headers=auth_headers, ) except httpx.TimeoutException as e: - last_error = KalshiError(f"Request timed out: {e}") + # F-O-09: don't interpolate httpx exception strings — they + # include the full URL with query string and can leak + # token-like values into uncaught-exception sinks. The + # underlying exception is preserved via `__cause__`. + last_error = KalshiError(f"Request timed out: {method.upper()} {path}") if method.upper() in RETRYABLE_METHODS and attempt < self._config.max_retries: delay = _compute_backoff(attempt, self._config) logger.warning("Timeout on %s %s, retrying in %.1fs", method, path, delay) @@ -152,7 +156,7 @@ def request( continue raise last_error from e except httpx.HTTPError as e: - raise KalshiError(f"HTTP error: {e}") from e + raise KalshiError(f"HTTP error: {method.upper()} {path}") from e logger.debug("Response: %s %s → %d", method.upper(), path, response.status_code) @@ -261,7 +265,8 @@ async def request( headers=auth_headers, ) except httpx.TimeoutException as e: - last_error = KalshiError(f"Request timed out: {e}") + # F-O-09: see sync transport above. + last_error = KalshiError(f"Request timed out: {method.upper()} {path}") if method.upper() in RETRYABLE_METHODS and attempt < self._config.max_retries: delay = _compute_backoff(attempt, self._config) logger.warning("Timeout on %s %s, retrying in %.1fs", method, path, delay) @@ -269,7 +274,7 @@ async def request( continue raise last_error from e except httpx.HTTPError as e: - raise KalshiError(f"HTTP error: {e}") from e + raise KalshiError(f"HTTP error: {method.upper()} {path}") from e logger.debug( "Async response: %s %s → %d", method.upper(), path, response.status_code diff --git a/kalshi/ws/channels.py b/kalshi/ws/channels.py index 27333e9..b43c1ef 100644 --- a/kalshi/ws/channels.py +++ b/kalshi/ws/channels.py @@ -7,6 +7,9 @@ import logging from typing import Any +from websockets.exceptions import ConnectionClosed + +from kalshi.errors import KalshiConnectionError, KalshiSubscriptionError from kalshi.ws.backpressure import MessageQueue, OverflowStrategy from kalshi.ws.connection import ConnectionManager @@ -83,14 +86,20 @@ async def _wait_for_response( while True: remaining = deadline - asyncio.get_event_loop().time() if remaining <= 0: - from kalshi.errors import KalshiSubscriptionError - raise KalshiSubscriptionError( f"Timed out waiting for response to command {msg_id}" ) - raw = await asyncio.wait_for( - self._connection.recv(), timeout=remaining - ) + try: + raw = await asyncio.wait_for( + self._connection.recv(), timeout=remaining + ) + except ConnectionClosed as e: + # F-P-05: surface as KalshiConnectionError instead of raw + # websockets exception. The recv loop's reconnect path will + # restore the subscription via resubscribe_all. + raise KalshiConnectionError( + f"Connection closed while awaiting response to command {msg_id}" + ) from e data: dict[str, Any] = json.loads(raw) if data.get("id") == msg_id: return data @@ -127,8 +136,6 @@ async def subscribe( # Read frames until we get our subscribe ack (by matching id) data = await self._wait_for_response(msg_id) if data.get("type") == "error": - from kalshi.errors import KalshiSubscriptionError - error_msg = data.get("msg", {}) raise KalshiSubscriptionError( str(error_msg.get("msg", "Subscribe failed")), @@ -160,6 +167,9 @@ async def unsubscribe(self, client_id: int) -> None: await self._connection.send(cmd) await self._wait_for_response(msg_id) + # F-P-08: push sentinel before deleting so any held iterator exits + # cleanly via StopAsyncIteration instead of hanging on queue.get(). + await sub.queue.put_sentinel() # Clean up mappings self._sid_to_client.pop(sub.server_sid, None) del self._subscriptions[client_id] @@ -178,8 +188,6 @@ async def update_subscription( """Add or remove markets from an existing subscription.""" sub = self._subscriptions.get(client_id) if not sub or sub.server_sid is None: - from kalshi.errors import KalshiSubscriptionError - raise KalshiSubscriptionError("Subscription not found or not active") msg_id = self._get_msg_id() @@ -201,31 +209,58 @@ async def resubscribe_all(self) -> None: Gets new server sids and updates the mapping. Iterators/callbacks continue working because they reference client_ids. + + F-P-01: Each resubscribe is independent. If one fails, the failed + subscription is removed and its iterator gets a sentinel, but other + subscriptions continue working. The whole reconnect is not aborted. """ old_subs = dict(self._subscriptions) + # F-P-03: clear sid->client mapping before sending any resubscribe + # so stale in-flight frames with old sids can't mis-route. self._sid_to_client.clear() for client_id, sub in old_subs.items(): sub.server_sid = None # Clear old sid - msg_id = self._get_msg_id() - # Re-subscribe with send_initial_snapshot for orderbook channels - params = sub.to_subscribe_params() - if sub.channel == "orderbook_delta": - params["send_initial_snapshot"] = True - cmd = {"id": msg_id, "cmd": "subscribe", "params": params} - await self._connection.send(cmd) - - data = await self._wait_for_response(msg_id) - new_sid = data.get("msg", {}).get("sid") - if new_sid is not None: - sub.server_sid = new_sid - self._sid_to_client[new_sid] = client_id - logger.debug( - "Resubscribed %s: client_id=%d, new_sid=%s", - sub.channel, - client_id, - new_sid, - ) + try: + msg_id = self._get_msg_id() + # Re-subscribe with send_initial_snapshot for orderbook channels + params = sub.to_subscribe_params() + if sub.channel == "orderbook_delta": + params["send_initial_snapshot"] = True + cmd = {"id": msg_id, "cmd": "subscribe", "params": params} + await self._connection.send(cmd) + + data = await self._wait_for_response(msg_id) + if data.get("type") == "error": + error_msg = data.get("msg", {}) + raise KalshiSubscriptionError( + str(error_msg.get("msg", "Resubscribe failed")), + error_code=error_msg.get("code"), + ) + new_sid = data.get("msg", {}).get("sid") + if new_sid is not None: + sub.server_sid = new_sid + self._sid_to_client[new_sid] = client_id + logger.debug( + "Resubscribed %s: client_id=%d, new_sid=%s", + sub.channel, + client_id, + new_sid, + ) + except Exception: + # F-P-01: per-sub failure is isolated. Push sentinel so the + # iterator exits cleanly, drop the subscription, continue with + # the rest. `exc_info=True` so an unexpected AttributeError / + # programming bug doesn't lose its traceback the way the + # earlier `: %s, e` formulation did. + logger.warning( + "Resubscribe failed for client_id=%d channel=%s", + client_id, + sub.channel, + exc_info=True, + ) + await sub.queue.put_sentinel() + self._subscriptions.pop(client_id, None) def get_subscription_by_sid(self, server_sid: int) -> Subscription | None: """Look up a subscription by current server sid.""" diff --git a/kalshi/ws/client.py b/kalshi/ws/client.py index 498e80c..b499888 100644 --- a/kalshi/ws/client.py +++ b/kalshi/ws/client.py @@ -9,11 +9,15 @@ from types import TracebackType from typing import Any -from pydantic import BaseModel +from pydantic import BaseModel, ValidationError from websockets.exceptions import ConnectionClosed from kalshi.auth import KalshiAuth from kalshi.config import KalshiConfig +from kalshi.errors import ( + KalshiBackpressureError, + KalshiSubscriptionError, +) from kalshi.models.markets import Orderbook from kalshi.ws.backpressure import OverflowStrategy from kalshi.ws.channels import SubscriptionManager @@ -127,7 +131,13 @@ def _ensure_recv_loop(self) -> None: self._recv_task = asyncio.create_task(self._recv_loop()) async def _pause_recv_loop(self) -> None: - """Cancel the recv_loop so subscribe can safely call recv.""" + """Cancel the recv_loop so subscribe can safely call recv. + + F-P-04: cancellation only fires while the loop is awaiting + ``connection.recv()``. The recv-then-process critical section is + shielded from cancellation so any frame already off the socket is + fully dispatched before the loop exits. + """ if self._recv_task and not self._recv_task.done(): self._recv_task.cancel() with contextlib.suppress(asyncio.CancelledError): @@ -142,69 +152,147 @@ async def _recv_loop(self) -> None: while self._running: try: raw = await self._connection.recv() - # Check for sequenced messages - try: - data = json.loads(raw) - sid = data.get("sid") - seq = data.get("seq") - msg_type = data.get("type", "") - - if sid is not None and seq is not None and self._seq_tracker: - channel = "" - sub = ( - self._sub_mgr.get_subscription_by_sid(sid) - if self._sub_mgr - else None - ) - if sub: - channel = sub.channel - ok = await self._seq_tracker.track( - sid, seq, msg_type if msg_type else channel - ) - if not ok: - # Gap detected — skip dispatching this message - # The gap handler will trigger resync - continue - - # Check for orderbook messages - if msg_type == "orderbook_snapshot" and self._orderbook_mgr: - snapshot = OrderbookSnapshotMessage.model_validate(data) - self._orderbook_mgr.apply_snapshot(snapshot) - elif msg_type == "orderbook_delta" and self._orderbook_mgr: - delta = OrderbookDeltaMessage.model_validate(data) - self._orderbook_mgr.apply_delta(delta) - except json.JSONDecodeError: - pass # dispatch will handle parse errors - - await self._dispatcher.dispatch(raw) - except asyncio.CancelledError: + # F-P-04: cancellation while awaiting recv = no frame read, + # no data lost. Safe to exit. break except ConnectionClosed: if not self._running: break - logger.warning("Connection lost, attempting reconnect...") - # Attempt reconnect + await self._handle_reconnect() + if not self._running: + break + continue + + # F-P-04: shield the recv->dispatch critical section. Once a + # frame has been read off the socket, we MUST dispatch it before + # honoring cancellation, otherwise the frame is silently dropped + # (seq trackers miss the gap, ticker frames vanish, etc.). + # + # IMPORTANT: hold the inner task in a local. `asyncio.shield(coro)` + # creates a detached background task and returns control on outer + # cancel — without the explicit `await inner`, dispatch keeps + # running after `break` and races a subsequent subscribe. + inner = asyncio.create_task(self._process_frame(raw)) + try: + await asyncio.shield(inner) + except asyncio.CancelledError: + # Block until the shielded dispatch actually finishes, then + # honor cancellation. Loop exit now strictly post-dispatch. + # KalshiBackpressureError / KalshiSubscriptionError must + # still broadcast sentinels even when we're being cancelled — + # the consumer-visible invariant (no hanging iterators) holds + # unconditionally. Parse errors are best-effort logged. try: - await self._connection.reconnect() - if self._sub_mgr: - if self._seq_tracker: - self._seq_tracker.reset_all() - if self._orderbook_mgr: - self._orderbook_mgr.clear() - await self._sub_mgr.resubscribe_all() - await self._connection._set_state(ConnectionState.STREAMING) - except Exception as reconnect_err: - logger.error("Reconnect failed: %s", reconnect_err) - # Send sentinels so consumers don't hang forever + await inner + except (KalshiBackpressureError, KalshiSubscriptionError): if self._sub_mgr: for sub in self._sub_mgr.active_subscriptions.values(): await sub.queue.put_sentinel() - break - except Exception as e: - # Application error (backpressure, callback, parse) — log, don't reconnect - logger.warning("Error processing message: %s", e) + except Exception: + logger.debug( + "Shielded dispatch raised during cancel cleanup", + exc_info=True, + ) + break + except (KalshiBackpressureError, KalshiSubscriptionError): + # #83: these signal real consumer-visible problems. Propagate + # via sentinel-broadcast so iterators see the failure and the + # loop exits cleanly rather than spinning on the same error. + 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() + break + except (json.JSONDecodeError, ValidationError, KeyError): + # #83: genuinely-non-fatal per-message errors. Log with + # traceback so users can debug, then continue. + logger.warning( + "Skipping malformed WS frame", exc_info=True, + ) continue + except Exception: + # #83 escape hatch: anything else (AttributeError from a + # user-supplied callback, TypeError, programming bug, ...) + # MUST broadcast sentinels before propagating — otherwise + # consumers hang on their queues with no signal. Re-raise + # after broadcast so the task failure is still visible. + logger.error( + "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() + raise + + async def _process_frame(self, raw: str) -> None: + """Parse, track seq, update orderbook, and dispatch a single frame.""" + assert self._dispatcher is not None + data = json.loads(raw) + sid = data.get("sid") + seq = data.get("seq") + msg_type = data.get("type", "") + + if sid is not None and seq is not None and self._seq_tracker: + channel = "" + sub = ( + self._sub_mgr.get_subscription_by_sid(sid) + if self._sub_mgr + else None + ) + if sub: + channel = sub.channel + ok = await self._seq_tracker.track( + sid, seq, msg_type if msg_type else channel + ) + if not ok: + # Gap detected — skip dispatching this message. + # The gap handler will trigger resync. + return + + # Check for orderbook messages + if msg_type == "orderbook_snapshot" and self._orderbook_mgr: + snapshot = OrderbookSnapshotMessage.model_validate(data) + self._orderbook_mgr.apply_snapshot(snapshot) + elif msg_type == "orderbook_delta" and self._orderbook_mgr: + delta = OrderbookDeltaMessage.model_validate(data) + self._orderbook_mgr.apply_delta(delta) + + await self._dispatcher.dispatch(raw) + + async def _handle_reconnect(self) -> None: + """Reconnect after ConnectionClosed and re-establish subscriptions. + + F-P-03: acquire `_subscribe_lock` for the entire reconnect+resubscribe + sequence so any in-flight user-level `subscribe_*` blocks until we've + rebuilt `_sid_to_client`. Without the lock, messages arriving on a + freshly-assigned sid during the rebuild can mis-route. + """ + assert self._connection is not None + logger.warning("Connection lost, attempting reconnect...") + async with self._subscribe_lock: + try: + await self._connection.reconnect() + if self._sub_mgr: + if self._seq_tracker: + self._seq_tracker.reset_all() + if self._orderbook_mgr: + self._orderbook_mgr.clear() + await self._sub_mgr.resubscribe_all() + # #88: use the public transition; no reach-through to + # ConnectionManager's name-mangled _set_state. + await self._connection.mark_streaming() + except Exception as reconnect_err: + 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() + self._running = False async def _handle_seq_gap(self, gap: SequenceGap) -> None: """Handle a sequence gap by logging and triggering resync.""" @@ -245,7 +333,7 @@ async def _do_subscribe( channel, params=params, overflow=overflow, maxsize=maxsize, ) finally: - # Always restart recv loop, even if subscribe fails + # Always restart/resume recv loop, even if subscribe fails. self._ensure_recv_loop() return sub.queue diff --git a/kalshi/ws/connection.py b/kalshi/ws/connection.py index 0722759..d5aace3 100644 --- a/kalshi/ws/connection.py +++ b/kalshi/ws/connection.py @@ -89,6 +89,15 @@ async def _set_state(self, new_state: ConnectionState) -> None: if self._on_state_change is not None: await self._on_state_change(old, new_state) + async def mark_streaming(self) -> None: + """Public transition from CONNECTED to STREAMING. + + The recv loop calls this after a successful reconnect+resubscribe + so the manager can fire `on_state_change` with the right caller. + Replaces the previous reach-through to private `_set_state` (#88). + """ + await self._set_state(ConnectionState.STREAMING) + def _build_auth_headers(self) -> dict[str, str]: """Build RSA-PSS auth headers for the WebSocket upgrade request.""" ws_path = urlparse(self._config.ws_base_url).path @@ -115,8 +124,14 @@ async def connect(self) -> None: await self._set_state(ConnectionState.CONNECTED) except Exception as e: await self._set_state(ConnectionState.CLOSED) + # F-O-09: don't interpolate the underlying exception string — + # it may include the full ws URL with query params and leak into + # uncaught-exception sinks. The cause is preserved via `__cause__`. + # Include the path (no query string) for debug context; per + # F-O-09 the original exception's URL stays out of __str__. + ws_path = urlparse(self._config.ws_base_url).path raise KalshiConnectionError( - f"WebSocket connection failed: {e}" + f"WebSocket connection failed: {ws_path}" ) from e async def reconnect(self) -> None: diff --git a/tests/test_client.py b/tests/test_client.py index a7a3eb9..5bff88e 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -361,6 +361,59 @@ def test_close(self, test_auth: KalshiAuth, config: KalshiConfig) -> None: transport.close() # should not raise +class TestErrorMessageDoesNotLeakUrl: + """Issue #84 F-O-09: KalshiError str() must not include the URL. + + httpx exception strings often contain the full request URL with query + string. Interpolating them into a KalshiError leaks any token-like + query parameter into Sentry/stderr/uncaught-exception sinks. + """ + + @respx.mock + def test_timeout_message_does_not_contain_url( + self, + transport: SyncTransport, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + monkeypatch.setattr("kalshi._base_client.time.sleep", lambda d: None) + # Use a sensitive-looking query param: if it shows up in the + # KalshiError __str__, the leak is real. + respx.post( + "https://test.kalshi.com/trade-api/v2/portfolio/orders" + ).mock(side_effect=httpx.TimeoutException( + "Request timed out for https://test.kalshi.com/trade-api/v2/" + "portfolio/orders?secret=SUPER_SECRET_TOKEN" + )) + with pytest.raises(KalshiError) as exc_info: + transport.request("POST", "/portfolio/orders", json={"ticker": "T"}) + msg = str(exc_info.value) + assert "SUPER_SECRET_TOKEN" not in msg + assert "kalshi.com" not in msg + # Method + relative path are safe debug context and SHOULD appear. + assert "POST" in msg + assert "/portfolio/orders" in msg + + @respx.mock + def test_httperror_message_does_not_contain_url( + self, + transport: SyncTransport, + ) -> None: + respx.get( + "https://test.kalshi.com/trade-api/v2/markets" + ).mock(side_effect=httpx.HTTPError( + "boom for https://test.kalshi.com/trade-api/v2/markets?" + "auth_token=LEAKY" + )) + with pytest.raises(KalshiError) as exc_info: + transport.request("GET", "/markets") + msg = str(exc_info.value) + assert "LEAKY" not in msg + assert "kalshi.com" not in msg + # Method + relative path are safe debug context and SHOULD appear. + assert "GET" in msg + assert "/markets" in msg + + class TestKalshiClientConstructor: """Tests for KalshiClient constructor branches and from_env().""" diff --git a/tests/ws/test_connection.py b/tests/ws/test_connection.py index 537c9a9..b5aaaf9 100644 --- a/tests/ws/test_connection.py +++ b/tests/ws/test_connection.py @@ -74,6 +74,48 @@ async def test_connect_invalid_url(self, test_auth: object) -> None: await mgr.connect() assert mgr.state == ConnectionState.CLOSED + async def test_mark_streaming_transitions_state( + self, fake_ws: FakeKalshiWS, test_auth: object + ) -> None: + """Issue #88: public mark_streaming() transitions CONNECTED -> + STREAMING and fires the state-change callback.""" + states: list[tuple[ConnectionState, ConnectionState]] = [] + + async def on_state(old: ConnectionState, new: ConnectionState) -> None: + states.append((old, new)) + + config = KalshiConfig(ws_base_url=fake_ws.url, timeout=5.0) + mgr = ConnectionManager( + auth=test_auth, config=config, on_state_change=on_state, # type: ignore[arg-type] + ) + await mgr.connect() + states.clear() + await mgr.mark_streaming() + assert mgr.state == ConnectionState.STREAMING + assert states == [(ConnectionState.CONNECTED, ConnectionState.STREAMING)] + await mgr.close() + + async def test_connect_error_does_not_leak_url( + self, test_auth: object + ) -> None: + """Issue #84 F-O-09: connection-failure str() must not include the + ws URL (which may contain token-like query params).""" + # URL with a sensitive-looking query param + config = KalshiConfig( + ws_base_url="ws://127.0.0.1:1?secret=SUPER_SECRET_TOKEN", + timeout=1.0, + ) + mgr = ConnectionManager(auth=test_auth, config=config) # type: ignore[arg-type] + with pytest.raises(KalshiConnectionError) as exc_info: + await mgr.connect() + msg = str(exc_info.value) + assert "SUPER_SECRET_TOKEN" not in msg + assert "127.0.0.1" not in msg + # The ws path (no query) is safe context and SHOULD appear. + # Our test config sets ws_base_url to a URL without an explicit path, + # so urlparse returns "" — assert by not crashing rather than substring. + assert "WebSocket connection failed" in msg + async def test_close_when_already_disconnected( self, test_auth: object ) -> None: diff --git a/tests/ws/test_recv_loop_hardening.py b/tests/ws/test_recv_loop_hardening.py new file mode 100644 index 0000000..9b9f5a8 --- /dev/null +++ b/tests/ws/test_recv_loop_hardening.py @@ -0,0 +1,414 @@ +"""Regression tests for WS recv-loop hardening (issues #77, #83). + +Covers the 5 reconnect/resubscribe race fixes from #77 plus the narrowed +exception catch from #83. Race scenarios use asyncio.Event barriers, not +sleeps, so they're deterministic. +""" + +from __future__ import annotations + +import asyncio +import json +import logging + +import pytest + +from kalshi.config import KalshiConfig +from kalshi.errors import KalshiConnectionError +from kalshi.ws.channels import SubscriptionManager +from kalshi.ws.client import KalshiWebSocket +from kalshi.ws.connection import ConnectionManager + +# --------------------------------------------------------------------------- +# F-P-01 — per-sub resubscribe failure no longer aborts the whole reconnect +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +class TestResubscribeIsolation: + async def test_one_failed_resubscribe_does_not_kill_others( + self, fake_ws, test_auth # type: ignore[no-untyped-def] + ) -> None: + """F-P-01: if one sub's resubscribe errors, the others still succeed.""" + config = KalshiConfig( + ws_base_url=fake_ws.url, + timeout=5.0, + retry_base_delay=0.01, + retry_max_delay=0.05, + ) + conn = ConnectionManager(auth=test_auth, config=config) + await conn.connect() + sub_mgr = SubscriptionManager(conn) + + sub1 = await sub_mgr.subscribe("ticker", params={"market_tickers": ["A"]}) + sub2 = await sub_mgr.subscribe("ticker", params={"market_tickers": ["B"]}) + sub3 = await sub_mgr.subscribe("ticker", params={"market_tickers": ["C"]}) + + # Simulate disconnect + await conn.close() + fake_ws.received_commands.clear() + fake_ws._next_sid = 50 + await conn.connect() + + # Make the SECOND resubscribe fail. We do this by monkey-patching the + # fake_ws subscribe handler to error on the second command id only. + original_handle = fake_ws._handle_command + call_count = {"n": 0} + + async def selective_handler(ws, msg): # type: ignore[no-untyped-def] + if msg.get("cmd") == "subscribe": + call_count["n"] += 1 + if call_count["n"] == 2: + msg_id = msg.get("id", 0) + await ws.send(json.dumps({ + "id": msg_id, + "type": "error", + "msg": {"code": 400, "msg": "simulated failure"}, + })) + return + await original_handle(ws, msg) + + fake_ws._handle_command = selective_handler # type: ignore[assignment] + + await sub_mgr.resubscribe_all() + + # sub2 should be removed; sub1 and sub3 should still be active + active = sub_mgr.active_subscriptions + assert sub1.client_id in active + assert sub2.client_id not in active + assert sub3.client_id in active + + # sub2's iterator must receive a sentinel (StopAsyncIteration) + with pytest.raises(StopAsyncIteration): + await asyncio.wait_for(sub2.queue.__anext__(), timeout=1.0) + + await conn.close() + + +# --------------------------------------------------------------------------- +# F-P-05 — ConnectionClosed during subscribe surfaces as KalshiConnectionError +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +class TestSubscribeConnectionClosed: + async def test_connection_closed_mid_subscribe_raises_kalshi_error( + self, fake_ws, test_auth # type: ignore[no-untyped-def] + ) -> None: + """F-P-05: a connection drop during _wait_for_response surfaces as + KalshiConnectionError, not raw websockets ConnectionClosed.""" + config = KalshiConfig(ws_base_url=fake_ws.url, timeout=5.0) + conn = ConnectionManager(auth=test_auth, config=config) + await conn.connect() + sub_mgr = SubscriptionManager(conn) + + # Override the handler to close the socket instead of replying. + async def closing_handler(ws, msg): # type: ignore[no-untyped-def] + await ws.close() + + fake_ws._handle_command = closing_handler # type: ignore[assignment] + + with pytest.raises(KalshiConnectionError): + await sub_mgr.subscribe("ticker") + + await conn.close() + + +# --------------------------------------------------------------------------- +# F-P-08 — unsubscribe pushes sentinel so held iterator exits +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +class TestUnsubscribeSentinel: + async def test_unsubscribe_pushes_sentinel( + self, fake_ws, test_auth # type: ignore[no-untyped-def] + ) -> None: + """F-P-08: unsubscribe must push a sentinel so iterators stop hanging.""" + config = KalshiConfig(ws_base_url=fake_ws.url, timeout=5.0) + conn = ConnectionManager(auth=test_auth, config=config) + await conn.connect() + sub_mgr = SubscriptionManager(conn) + + sub = await sub_mgr.subscribe("ticker") + assert sub.server_sid is not None + + await sub_mgr.unsubscribe(sub.client_id) + + # Iterator on the queue should terminate, not hang + with pytest.raises(StopAsyncIteration): + await asyncio.wait_for(sub.queue.__anext__(), timeout=1.0) + + await conn.close() + + +# --------------------------------------------------------------------------- +# F-P-04 — _pause_recv_loop does not drop in-flight frames +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +class TestPauseDoesNotDropFrames: + async def test_inflight_frame_dispatched_when_recv_task_cancelled( + self, fake_ws, test_auth # type: ignore[no-untyped-def] + ) -> None: + """F-P-04: a frame mid-dispatch when the recv task is cancelled must + still be delivered. We monkey-patch the dispatcher's dispatch() with + a barrier so we can force cancellation precisely while it's running. + Without the asyncio.shield in _process_frame, the frame is lost. + """ + 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"]) + + assert session._dispatcher is not None + original_dispatch = session._dispatcher.dispatch + + dispatch_started = asyncio.Event() + allow_dispatch_finish = asyncio.Event() + + async def slow_dispatch(raw: str) -> None: + dispatch_started.set() + await allow_dispatch_finish.wait() + await original_dispatch(raw) + + session._dispatcher.dispatch = slow_dispatch # type: ignore[method-assign] + + sid = next(iter(session._sub_mgr.active_subscriptions.values())).server_sid # type: ignore[union-attr] + # Send one frame. recv loop reads it, enters _process_frame, + # calls slow_dispatch which blocks on the event. + await fake_ws.send_to_all({ + "type": "ticker", "sid": sid, + "msg": { + "market_ticker": "T1", "market_id": "x", "yes_bid": 22, + }, + }) + await asyncio.wait_for(dispatch_started.wait(), timeout=2.0) + + # Now request a pause (cancellation). The shield must keep the + # dispatch alive. + pause_task = asyncio.create_task(session._pause_recv_loop()) + # Give the cancel a chance to propagate + await asyncio.sleep(0.05) + + # Release the dispatch barrier — shielded coroutine runs to + # completion and puts the message on the queue. + allow_dispatch_finish.set() + await asyncio.wait_for(pause_task, timeout=2.0) + + # The message must have arrived despite the cancellation. + msg = await asyncio.wait_for(stream.__anext__(), timeout=2.0) + assert msg.msg.yes_bid == 22 + + # Restore original dispatch and resume + session._dispatcher.dispatch = original_dispatch # type: ignore[method-assign] + session._ensure_recv_loop() + + +# --------------------------------------------------------------------------- +# F-P-03 — reconnect path takes subscribe lock +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +class TestReconnectHoldsSubscribeLock: + async def test_handle_reconnect_acquires_subscribe_lock( + self, fake_ws, test_auth # type: ignore[no-untyped-def] + ) -> None: + """F-P-03: reconnect must hold _subscribe_lock so concurrent user + subscribe_* cannot race the sid-remap. + + Verified structurally: while _handle_reconnect runs, a user subscribe + coroutine cannot complete because it competes for the same lock. + """ + config = KalshiConfig( + ws_base_url=fake_ws.url, + timeout=5.0, + retry_base_delay=0.01, + retry_max_delay=0.05, + ws_max_retries=3, + ) + ws = KalshiWebSocket(auth=test_auth, config=config) + async with ws.connect() as session: + # Acquire the lock externally to simulate user subscribe in flight + assert session._subscribe_lock is not None + await session._subscribe_lock.acquire() + + # _handle_reconnect should wait for the lock + reconnect_task = asyncio.create_task(session._handle_reconnect()) + + # Deterministic: yield until reconnect_task is parked on the lock + # (lock._waiters is non-empty). Bounded by wait_for so a regression + # where _handle_reconnect doesn't take the lock fails loudly within + # 2s rather than passing vacuously after a fixed sleep. + async def wait_for_lock_contention() -> None: + while not session._subscribe_lock._waiters: # type: ignore[attr-defined] + await asyncio.sleep(0) + + await asyncio.wait_for(wait_for_lock_contention(), timeout=2.0) + assert not reconnect_task.done(), ( + "F-P-03 regression: _handle_reconnect did not take " + "_subscribe_lock" + ) + + # Release; reconnect proceeds + session._subscribe_lock.release() + await asyncio.wait_for(reconnect_task, timeout=3.0) + + +# --------------------------------------------------------------------------- +# #83 — recv loop narrows broad except +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +class TestRecvLoopExceptionPolicy: + async def test_backpressure_error_breaks_loop_and_sentinels( + self, fake_ws, test_auth # type: ignore[no-untyped-def] + ) -> None: + """#83: KalshiBackpressureError no longer swallowed. The recv loop + exits and consumers see sentinel. + + Deterministic setup: subscribe but do NOT iterate, pre-fill the queue + to maxsize, then send one more frame so the recv-loop's queue.put + raises BackpressureError. Then start iterating — must see sentinel. + """ + config = KalshiConfig(ws_base_url=fake_ws.url, timeout=5.0) + ws = KalshiWebSocket(auth=test_auth, config=config) + async with ws.connect() as session: + # ERROR strategy with maxsize=1 -> second message overflows + stream = await session.subscribe_orderbook_delta( + tickers=["T1"], maxsize=1, + ) + + assert session._sub_mgr is not None + sid = next(iter(session._sub_mgr.active_subscriptions.values())).server_sid + assert sid is not None + + # Fill the queue: snapshot lands as the single allowed item. + await fake_ws.send_to_all({ + "type": "orderbook_snapshot", "sid": sid, "seq": 1, + "msg": { + "market_ticker": "T1", "market_id": "x", + "yes": [["0.50", "100"]], "no": [], + }, + }) + # Yield to let recv loop process the snapshot + await asyncio.sleep(0.1) + + # Second frame overflows -> KalshiBackpressureError + await fake_ws.send_to_all({ + "type": "orderbook_delta", "sid": sid, "seq": 2, + "msg": { + "market_ticker": "T1", "market_id": "x", + "price": "0.51", "delta": "10", "side": "yes", + }, + }) + + # Wait for the recv loop to process the delta and exit on + # BackpressureError. The recv task should complete (loop broke). + recv_task = session._recv_task + assert recv_task is not None + await asyncio.wait_for(recv_task, timeout=2.0) + + # Iterator must terminate. Buffer holds [snapshot, SENTINEL] + # (sentinel is appended unconditionally by put_sentinel). + collected: list[object] = [] + try: + async with asyncio.timeout(3.0): + async for msg in stream: + collected.append(msg) + except TimeoutError: + pytest.fail( + "#83 regression: BackpressureError was swallowed; " + "iterator did not receive sentinel." + ) + assert len(collected) == 1 # only the snapshot was queued + + async def test_malformed_frame_logged_with_traceback_and_continues( + self, + fake_ws, # type: ignore[no-untyped-def] + test_auth, # type: ignore[no-untyped-def] + caplog: pytest.LogCaptureFixture, + ) -> None: + """#83: ValidationError on a single frame logs with exc_info but + does NOT take down the loop.""" + 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"]) + + assert session._sub_mgr is not None + sid = next(iter(session._sub_mgr.active_subscriptions.values())).server_sid + + # Send a frame that will fail validation (missing required field + # for ticker), then send a valid frame. The loop must keep going. + with caplog.at_level(logging.WARNING, logger="kalshi.ws"): + # Inject a parse-failure: the recv loop's _process_frame + # validates orderbook messages and json.loads anything. Send + # invalid JSON directly to trigger json.JSONDecodeError. + for conn in fake_ws.connections: + await conn.send("{not valid json") + + await fake_ws.send_to_all({ + "type": "ticker", "sid": sid, + "msg": { + "market_ticker": "T1", "market_id": "x", + "yes_bid": 77, + }, + }) + + msg = await asyncio.wait_for(stream.__anext__(), timeout=2.0) + assert msg.msg.yes_bid == 77 + + # The warning should have been logged + assert any( + "malformed" in r.message.lower() or "non-json" in r.message.lower() + for r in caplog.records + ) + + async def test_unexpected_exception_broadcasts_sentinels_before_raising( + self, + fake_ws, # type: ignore[no-untyped-def] + test_auth, # type: ignore[no-untyped-def] + caplog: pytest.LogCaptureFixture, + ) -> None: + """#83 escape hatch: a TypeError/AttributeError/etc. from inside + dispatch (e.g. a buggy user callback) must broadcast sentinels to + every consumer before propagating so iterators don't hang. + """ + 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"]) + assert session._sub_mgr is not None + sid = next(iter(session._sub_mgr.active_subscriptions.values())).server_sid + + # Monkey-patch dispatch to raise an unexpected exception class + # (AttributeError is neither in the JSON/Validation/Key bucket + # nor in the KalshiBackpressure/Subscription bucket). + async def boom(_raw: str) -> None: + raise AttributeError("simulated user-callback bug") + + session._dispatcher.dispatch = boom # type: ignore[method-assign] + + with caplog.at_level(logging.ERROR, logger="kalshi.ws"): + await fake_ws.send_to_all({ + "type": "ticker", "sid": sid, + "msg": {"market_ticker": "T1", "market_id": "x", "yes_bid": 1}, + }) + + # Iterator must see sentinel (StopAsyncIteration) within timeout. + with pytest.raises(StopAsyncIteration): + await asyncio.wait_for(stream.__anext__(), timeout=2.0) + + assert any( + "Unexpected error in recv loop" in r.message + for r in caplog.records + ) + # Recv task is in failed state (we re-raised after sentinel). + recv_task = session._recv_task + assert recv_task is not None + assert recv_task.done() + assert isinstance(recv_task.exception(), AttributeError)