diff --git a/kalshi/ws/client.py b/kalshi/ws/client.py index 4fe3e2c..083a767 100644 --- a/kalshi/ws/client.py +++ b/kalshi/ws/client.py @@ -101,6 +101,7 @@ async def _start(self) -> None: self._dispatcher = MessageDispatcher( sub_mgr=self._sub_mgr, on_error=self._on_error, + seq_tracker=self._seq_tracker, ) self._running = True @@ -268,16 +269,22 @@ async def _process_frame(self, raw: str) -> None: # Only meaningful once we know we'll dispatch (and might roll back). tracked = True - # Check for orderbook messages + # Check for orderbook messages — validate ONCE for the local + # manager, then hand the typed message off to dispatch via + # pre_validated so the dispatcher routes the same instance to + # the queue without re-running Pydantic. + pre_validated: BaseModel | None = None if msg_type == "orderbook_snapshot" and self._orderbook_mgr: snapshot = OrderbookSnapshotMessage.model_validate(data) self._orderbook_mgr.apply_snapshot(snapshot) + pre_validated = snapshot elif msg_type == "orderbook_delta" and self._orderbook_mgr: delta = OrderbookDeltaMessage.model_validate(data) self._orderbook_mgr.apply_delta(delta) + pre_validated = delta try: - await self._dispatcher.dispatch(raw) + await self._dispatcher.dispatch(data, pre_validated=pre_validated) except KalshiBackpressureError: # Dispatch failed -> the consumer never saw this message. Roll # the seq watermark back so the dropped seq is treated as a diff --git a/kalshi/ws/dispatch.py b/kalshi/ws/dispatch.py index 37453ad..00aff54 100644 --- a/kalshi/ws/dispatch.py +++ b/kalshi/ws/dispatch.py @@ -1,7 +1,6 @@ """Message dispatcher: parse raw frames, route to queues/callbacks.""" from __future__ import annotations -import json import logging from collections.abc import Awaitable, Callable from typing import Any @@ -167,14 +166,22 @@ async def _surface_channel_error( msg_type, data, ) - async def dispatch(self, raw: str) -> None: - """Parse a raw JSON frame and route it.""" - try: - data = json.loads(raw) - except json.JSONDecodeError: - logger.warning("Received non-JSON frame: %s", raw[:100]) - return + async def dispatch( + self, + data: dict[str, Any], + *, + pre_validated: BaseModel | None = None, + ) -> None: + """Route a pre-parsed frame to its subscriber. + ``data`` is the already-decoded JSON envelope. The recv loop owns + ``json.loads`` (so the dispatcher doesn't parse twice). + + ``pre_validated`` lets the recv loop hand off a typed model it + already validated (orderbook snapshots/deltas applied to the + local book), so the dispatcher routes the same instance to the + queue without re-running Pydantic. + """ msg_type: str = data.get("type", "") # Skip control messages (handled by subscribe/unsubscribe flow) @@ -200,11 +207,20 @@ async def dispatch(self, raw: str) -> None: logger.warning("Unknown message type: %s", msg_type) return - try: - parsed = model_cls.model_validate(data) - except Exception: - logger.warning("Failed to parse %s message", msg_type, exc_info=True) - return + if pre_validated is not None and isinstance(pre_validated, model_cls): + parsed = pre_validated + else: + try: + parsed = model_cls.model_validate(data) + except Exception as exc: + # Drop exc_info: Pydantic's ValidationError __str__ echoes the + # full input including trade payload (price, count, user fields) + # straight into log sinks. Surface type + exception class only. + logger.warning( + "Failed to parse %s message: %s", + msg_type, type(exc).__name__, + ) + return # Route to subscription queue sid = data.get("sid") diff --git a/tests/ws/test_dispatch.py b/tests/ws/test_dispatch.py index ce2a86b..5687b98 100644 --- a/tests/ws/test_dispatch.py +++ b/tests/ws/test_dispatch.py @@ -60,10 +60,30 @@ async def test_dispatch_ticker(self) -> None: raw = json.dumps( {"type": "ticker", "sid": 1, "msg": {"market_ticker": "T", "market_id": "x"}} ) - await dispatcher.dispatch(raw) + await dispatcher.dispatch(json.loads(raw)) msg = await sub.queue.get() assert msg.msg.market_ticker == "T" + async def test_dispatch_pre_validated_skips_revalidation(self) -> None: + """Regression for #86: when recv loop hands us the typed message it + already validated, we must NOT re-run model_validate. The same + instance lands on the queue. + """ + from kalshi.ws.models.orderbook_delta import OrderbookSnapshotMessage + mgr = FakeSubManager() + sub = mgr.add(2, "orderbook_delta") + dispatcher = MessageDispatcher(sub_mgr=mgr) # type: ignore[arg-type] + data = { + "type": "orderbook_snapshot", + "sid": 2, + "seq": 1, + "msg": {"market_ticker": "M", "market_id": "x", "yes": [], "no": []}, + } + pre_validated = OrderbookSnapshotMessage.model_validate(data) + await dispatcher.dispatch(data, pre_validated=pre_validated) + delivered = await sub.queue.get() + assert delivered is pre_validated, "dispatcher re-validated instead of using pre_validated" + async def test_dispatch_orderbook_snapshot(self) -> None: mgr = FakeSubManager() sub = mgr.add(2, "orderbook_delta") @@ -76,7 +96,7 @@ async def test_dispatch_orderbook_snapshot(self) -> None: "msg": {"market_ticker": "M", "market_id": "x", "yes": [], "no": []}, } ) - await dispatcher.dispatch(raw) + await dispatcher.dispatch(json.loads(raw)) msg = await sub.queue.get() assert msg.type == "orderbook_snapshot" @@ -84,7 +104,7 @@ async def test_dispatch_unknown_type_no_crash(self) -> None: mgr = FakeSubManager() dispatcher = MessageDispatcher(sub_mgr=mgr) # type: ignore[arg-type] raw = json.dumps({"type": "future_type", "sid": 1, "msg": {}}) - await dispatcher.dispatch(raw) # should not crash + await dispatcher.dispatch(json.loads(raw)) # should not crash async def test_dispatch_control_message_skipped(self) -> None: mgr = FakeSubManager() @@ -93,21 +113,16 @@ async def test_dispatch_control_message_skipped(self) -> None: raw = json.dumps( {"type": "subscribed", "id": 1, "msg": {"channel": "ticker", "sid": 1}} ) - await dispatcher.dispatch(raw) + await dispatcher.dispatch(json.loads(raw)) assert sub.queue.qsize() == 0 # control messages don't go to queue - async def test_dispatch_invalid_json(self) -> None: - mgr = FakeSubManager() - dispatcher = MessageDispatcher(sub_mgr=mgr) # type: ignore[arg-type] - await dispatcher.dispatch("not json at all") # should not crash - async def test_dispatch_unknown_sid(self) -> None: mgr = FakeSubManager() dispatcher = MessageDispatcher(sub_mgr=mgr) # type: ignore[arg-type] raw = json.dumps( {"type": "ticker", "sid": 999, "msg": {"market_ticker": "T", "market_id": "x"}} ) - await dispatcher.dispatch(raw) # should not crash + await dispatcher.dispatch(json.loads(raw)) # should not crash async def test_callback_mode(self) -> None: """Callback AND queue both receive the message (fan-out, see #80).""" @@ -123,7 +138,7 @@ async def on_ticker(msg: object) -> None: raw = json.dumps( {"type": "ticker", "sid": 1, "msg": {"market_ticker": "T", "market_id": "x"}} ) - await dispatcher.dispatch(raw) + await dispatcher.dispatch(json.loads(raw)) assert len(received) == 1 assert sub.queue.qsize() == 1 # also fanned out to the iterator queue @@ -152,8 +167,8 @@ async def on_ticker(msg: object) -> None: raw = json.dumps( {"type": "ticker", "sid": 1, "msg": {"market_ticker": "T", "market_id": "x"}} ) - await dispatcher.dispatch(raw) - await dispatcher.dispatch(raw) + await dispatcher.dispatch(json.loads(raw)) + await dispatcher.dispatch(json.loads(raw)) # Both sinks observed both messages. assert len(received) == 2 @@ -185,7 +200,7 @@ async def on_error(err: object) -> None: dispatcher = MessageDispatcher(sub_mgr=mgr, on_error=on_error) # type: ignore[arg-type] raw = json.dumps({"type": "error", "id": 1, "msg": {"code": 5, "msg": "bad"}}) - await dispatcher.dispatch(raw) + await dispatcher.dispatch(json.loads(raw)) assert len(errors) == 1 async def test_channel_level_error_routed_to_on_error(self) -> None: @@ -209,7 +224,7 @@ async def on_error(err: object) -> None: "error": {"code": 7, "msg": "schema violation"}, } ) - await dispatcher.dispatch(raw) + await dispatcher.dispatch(json.loads(raw)) assert len(errors) == 1 async def test_channel_level_error_logged_when_no_handler( @@ -223,7 +238,7 @@ async def test_channel_level_error_logged_when_no_handler( {"type": "ticker", "sid": 1, "error": {"code": 7, "msg": "boom"}} ) with caplog.at_level(logging.WARNING, logger="kalshi.ws"): - await dispatcher.dispatch(raw) + await dispatcher.dispatch(json.loads(raw)) assert any( "Channel-level error envelope" in r.message for r in caplog.records ) @@ -244,7 +259,7 @@ async def on_error(err: object) -> None: "error": {"code": 9, "msg": "stale sid"}, } ) - await dispatcher.dispatch(raw) + await dispatcher.dispatch(json.loads(raw)) assert len(errors) == 1 async def test_null_error_field_not_misrouted(self) -> None: @@ -270,7 +285,7 @@ async def on_error(err: object) -> None: "msg": {"market_ticker": "T", "market_id": "x"}, } ) - await dispatcher.dispatch(raw) + await dispatcher.dispatch(json.loads(raw)) assert errors == [], "null error field misrouted as channel error" # Normal ticker flow happened: message landed on the queue. assert sub.queue.qsize() == 1 @@ -295,7 +310,7 @@ async def on_error(err: object) -> None: {"type": "ticker", "sid": 1, "error": 42} ) with caplog.at_level(logging.ERROR, logger="kalshi.ws"): - await dispatcher.dispatch(raw) + await dispatcher.dispatch(json.loads(raw)) assert len(errors) == 1, "on_error not called on validation failure" assert any( "failed strict ErrorMessage validation" in r.message @@ -339,7 +354,7 @@ async def on_ticker(msg: object) -> None: raw = json.dumps( {"type": "ticker", "sid": 1, "msg": {"market_ticker": "T", "market_id": "x"}} ) - await dispatcher.dispatch(raw) + await dispatcher.dispatch(json.loads(raw)) assert len(received) == 0 # callback was removed assert sub.queue.qsize() == 1 # routed to queue instead @@ -357,7 +372,7 @@ async def test_server_unsubscribe_reaps_state(self) -> None: assert 7 in mgr._sid_to_client and 7 in mgr._subscriptions raw = json.dumps({"type": "unsubscribed", "msg": {"sid": 7}}) - await dispatcher.dispatch(raw) + await dispatcher.dispatch(json.loads(raw)) assert 7 not in mgr._sid_to_client assert 7 not in mgr._subscriptions @@ -373,7 +388,7 @@ async def test_server_unsubscribe_top_level_sid(self) -> None: dispatcher = MessageDispatcher(sub_mgr=mgr) # type: ignore[arg-type] raw = json.dumps({"type": "unsubscribed", "sid": 8, "seq": 0}) - await dispatcher.dispatch(raw) + await dispatcher.dispatch(json.loads(raw)) assert 8 not in mgr._sid_to_client assert 8 not in mgr._subscriptions @@ -395,7 +410,7 @@ async def test_server_unsubscribe_resets_seq_tracker(self) -> None: seq_tracker=tracker, ) raw = json.dumps({"type": "unsubscribed", "msg": {"sid": 9}}) - await dispatcher.dispatch(raw) + await dispatcher.dispatch(json.loads(raw)) assert 9 not in tracker._last_seq @@ -415,7 +430,7 @@ async def test_server_unsubscribe_with_distinct_client_id(self) -> None: assert 500 not in mgr._subscriptions # not keyed by server sid raw = json.dumps({"type": "unsubscribed", "msg": {"sid": 500}}) - await dispatcher.dispatch(raw) + await dispatcher.dispatch(json.loads(raw)) assert 500 not in mgr._sid_to_client assert 1 not in mgr._subscriptions @@ -427,14 +442,14 @@ async def test_server_unsubscribe_unknown_sid_no_crash(self) -> None: mgr = FakeSubManager() dispatcher = MessageDispatcher(sub_mgr=mgr) # type: ignore[arg-type] raw = json.dumps({"type": "unsubscribed", "msg": {"sid": 999}}) - await dispatcher.dispatch(raw) # should not crash + await dispatcher.dispatch(json.loads(raw)) # should not crash async def test_dispatch_message_without_sid(self) -> None: """Messages without sid are logged but don't crash.""" mgr = FakeSubManager() dispatcher = MessageDispatcher(sub_mgr=mgr) # type: ignore[arg-type] raw = json.dumps({"type": "ticker", "msg": {"market_ticker": "T", "market_id": "x"}}) - await dispatcher.dispatch(raw) # should not crash + await dispatcher.dispatch(json.loads(raw)) # should not crash @pytest.mark.asyncio @@ -449,7 +464,7 @@ async def test_dispatch_routes_user_order_singular() -> None: sub = mgr.add(42, "user_orders") dispatcher = MessageDispatcher(sub_mgr=mgr) # type: ignore[arg-type] raw = '{"type":"user_order","sid":42,"msg":{"order_id":"ORD1"}}' - await dispatcher.dispatch(raw) + await dispatcher.dispatch(json.loads(raw)) msg = await asyncio.wait_for(sub.queue.get(), timeout=1.0) assert isinstance(msg, UserOrdersMessage) @@ -475,7 +490,7 @@ async def test_dispatch_routes_market_position_singular() -> None: sub = mgr.add(42, "market_positions") dispatcher = MessageDispatcher(sub_mgr=mgr) # type: ignore[arg-type] raw = '{"type":"market_position","sid":42,"msg":{"ticker":"X","market_ticker":"X"}}' - await dispatcher.dispatch(raw) + await dispatcher.dispatch(json.loads(raw)) msg = await asyncio.wait_for(sub.queue.get(), timeout=1.0) assert isinstance(msg, MarketPositionsMessage) @@ -499,7 +514,7 @@ async def test_dispatch_routes_multivariate_lookup() -> None: sub = mgr.add(17, "multivariate") dispatcher = MessageDispatcher(sub_mgr=mgr) # type: ignore[arg-type] raw = '{"type":"multivariate_lookup","sid":17,"msg":{"event_ticker":"E1"}}' - await dispatcher.dispatch(raw) + await dispatcher.dispatch(json.loads(raw)) msg = await asyncio.wait_for(sub.queue.get(), timeout=1.0) assert isinstance(msg, MultivariateMessage) diff --git a/tests/ws/test_recv_loop_hardening.py b/tests/ws/test_recv_loop_hardening.py index 9b9f5a8..4e2f4a2 100644 --- a/tests/ws/test_recv_loop_hardening.py +++ b/tests/ws/test_recv_loop_hardening.py @@ -10,6 +10,7 @@ import asyncio import json import logging +from typing import Any import pytest @@ -168,10 +169,12 @@ async def test_inflight_frame_dispatched_when_recv_task_cancelled( dispatch_started = asyncio.Event() allow_dispatch_finish = asyncio.Event() - async def slow_dispatch(raw: str) -> None: + async def slow_dispatch( + data: dict[str, Any], *, pre_validated: Any = None, + ) -> None: dispatch_started.set() await allow_dispatch_finish.wait() - await original_dispatch(raw) + await original_dispatch(data, pre_validated=pre_validated) session._dispatcher.dispatch = slow_dispatch # type: ignore[method-assign] @@ -388,7 +391,7 @@ async def test_unexpected_exception_broadcasts_sentinels_before_raising( # 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: + async def boom(_data: dict[str, Any], **_kw: Any) -> None: raise AttributeError("simulated user-callback bug") session._dispatcher.dispatch = boom # type: ignore[method-assign]