From e8039f32e856f46b2b008da8fe9584b80642fa23 Mon Sep 17 00:00:00 2001 From: Jeff West Date: Sun, 17 May 2026 13:34:14 -0500 Subject: [PATCH 1/3] fix(ws): wire SequenceTracker into MessageDispatcher MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit PR #137 added an optional seq_tracker kwarg on MessageDispatcher.__init__ so server-initiated unsubscribe (#81) can clear stale seq state alongside the sid mapping. The wiring at the construction site in KalshiWebSocket was never landed — the dispatcher always saw seq_tracker=None and could only clear half the state. Closes #27 (post-Wave-3 integration tracker). Co-Authored-By: Claude Opus 4.7 (1M context) --- kalshi/ws/client.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) 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 From d95314fbd68dc95257546c9460dedea590eac463 Mon Sep 17 00:00:00 2001 From: Jeff West Date: Sun, 17 May 2026 13:35:07 -0500 Subject: [PATCH 2/3] =?UTF-8?q?fix(ws):=20close=20#84=20F-O-05=20+=20#86?= =?UTF-8?q?=20=E2=80=94=20drop=20exc=5Finfo=20leak,=20single=20parse=20pat?= =?UTF-8?q?h?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit F-O-05 / Closes #84 (final half): The validation-failure log in MessageDispatcher.dispatch used `logger.warning("Failed to parse %s message", msg_type, exc_info=True)`. Pydantic's ValidationError stringification echoes the full input — for trade/fill/order_group channels that includes price, count, and sometimes user identifiers, dumped straight into stdout/Sentry/log infrastructure. Dropped exc_info; surface type + exception class name only. The underlying exception is still available via __cause__ for local debug. #86 dispatcher side: - Signature change: dispatch(raw: str) -> dispatch(data: dict, *, pre_validated: BaseModel | None = None). The recv loop now owns the json.loads (it already had one); the dispatcher's redundant second parse is gone. - pre_validated lets the recv loop hand off an orderbook message it already validated for the local OrderbookManager, so the dispatcher routes the same instance to the subscription queue without re-running Pydantic. Eliminates the double-validate the issue called out (~3 µs/msg). - json import dropped from dispatch.py (now unused). Recv loop client.py side (in the previous commit alongside the seq_tracker wiring): _process_frame captures the typed orderbook message and passes pre_validated to dispatch. Test fallout: - 21 tests in test_dispatch.py updated from `dispatcher.dispatch(raw)` to `dispatcher.dispatch(json.loads(raw))` via sed. - Dropped test_dispatch_invalid_json — that path now lives in the recv loop's exception ladder (covered by test_malformed_frame_logged_with_traceback_and_continues). - 2 tests in test_recv_loop_hardening.py monkey-patched dispatch with the old (raw: str) signature; updated to (data: dict[str, Any], **kw). - New regression: test_dispatch_pre_validated_skips_revalidation — asserts dispatch routes the same instance when pre_validated matches the expected model_cls. Closes #84 Closes #86 Co-Authored-By: Claude Opus 4.7 (1M context) --- kalshi/ws/dispatch.py | 43 +++++++++++----- tests/ws/test_dispatch.py | 73 +++++++++++++++++----------- tests/ws/test_recv_loop_hardening.py | 9 ++-- 3 files changed, 80 insertions(+), 45 deletions(-) diff --git a/kalshi/ws/dispatch.py b/kalshi/ws/dispatch.py index 37453ad..a2a9882 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,21 @@ 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; + # the underlying exception is preserved on __cause__ for debug. + 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] From aa98742d46dacb06c2c117fee5b45524cf0d3f92 Mon Sep 17 00:00:00 2001 From: Jeff West Date: Sun, 17 May 2026 13:39:12 -0500 Subject: [PATCH 3/3] review(#142): drop false __cause__ claim in dispatch parse-fail log MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per bot review on PR #142: the comment said "the underlying exception is preserved on __cause__ for debug" — but the except block catches, logs the type name, and returns. Nothing is re-raised, so no __cause__ chain is established. Dropped the misleading sentence; kept the security rationale (Pydantic input echo). Co-Authored-By: Claude Opus 4.7 (1M context) --- kalshi/ws/dispatch.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/kalshi/ws/dispatch.py b/kalshi/ws/dispatch.py index a2a9882..00aff54 100644 --- a/kalshi/ws/dispatch.py +++ b/kalshi/ws/dispatch.py @@ -215,8 +215,7 @@ async def dispatch( 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; - # the underlying exception is preserved on __cause__ for debug. + # straight into log sinks. Surface type + exception class only. logger.warning( "Failed to parse %s message: %s", msg_type, type(exc).__name__,