diff --git a/CHANGELOG.md b/CHANGELOG.md index 30efe4a..33ee3c5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -21,6 +21,15 @@ All notable changes to kalshi-sdk will be documented in this file. guard (`KalshiError` on a repeating cursor) provides the runaway protection it always did. Callers who want a cap pass `max_pages=N` explicitly. No user-visible change for anyone iterating <1000 pages (#98). +- **WS callbacks no longer suppress queue delivery (#80).** Previously, + registering a callback for a channel silently disabled the iterator queue + for that channel — a user who held both an `@on()` callback AND an + `async for msg in subscription` consumer would never see the iterator + fire. Now messages fan out to both. A WARNING is logged at + `register_callback` time if an active subscription already exists, so + users relying on the suppress-side-effect are alerted. Callback-only + users now accumulate up to the queue's `maxsize=1000`; existing + `DROP_OLDEST` backpressure prevents unbounded growth. ### Fixed diff --git a/kalshi/ws/dispatch.py b/kalshi/ws/dispatch.py index ca49c13..37453ad 100644 --- a/kalshi/ws/dispatch.py +++ b/kalshi/ws/dispatch.py @@ -9,7 +9,7 @@ from pydantic import BaseModel from kalshi.ws.channels import SubscriptionManager -from kalshi.ws.models.base import ErrorMessage +from kalshi.ws.models.base import ErrorMessage, ErrorPayload from kalshi.ws.models.communications import CommunicationsMessage from kalshi.ws.models.fill import FillMessage from kalshi.ws.models.market_lifecycle import MarketLifecycleMessage @@ -20,6 +20,7 @@ from kalshi.ws.models.ticker import TickerMessage from kalshi.ws.models.trade import TradeMessage from kalshi.ws.models.user_orders import UserOrdersMessage +from kalshi.ws.sequence import SequenceTracker logger = logging.getLogger("kalshi.ws") @@ -50,21 +51,122 @@ def __init__( self, sub_mgr: SubscriptionManager, on_error: Callable[[ErrorMessage], Awaitable[None]] | None = None, + seq_tracker: SequenceTracker | None = None, ) -> None: self._sub_mgr = sub_mgr self._on_error = on_error + self._seq_tracker = seq_tracker self._callbacks: dict[str, Callable[[Any], Awaitable[None]]] = {} def register_callback( self, channel: str, callback: Callable[[Any], Awaitable[None]] ) -> None: - """Register a callback for a specific channel type.""" + """Register a callback for a specific channel type. + + Messages on this channel are fanned out to BOTH the callback and any + active subscription queue for the same channel. If an active + subscription queue exists, a warning is logged so users know the + callback was registered after the fact. + """ + if any( + sub.channel == channel + for sub in self._sub_mgr.active_subscriptions.values() + ): + logger.warning( + "register_callback(%r): an active subscription exists on this " + "channel; messages will be delivered to BOTH the callback and " + "the existing iterator queue.", + channel, + ) self._callbacks[channel] = callback def unregister_callback(self, channel: str) -> None: """Remove a callback for a channel type.""" self._callbacks.pop(channel, None) + async def _handle_server_unsubscribe(self, data: dict[str, Any]) -> None: + """Reap state for a server-initiated unsubscribe envelope (#81). + + Client-initiated unsubscribes consume the ack inside + SubscriptionManager._wait_for_response, so any unsubscribed frame + that reaches the dispatcher is server-initiated (admin action, + session expiry) or a late ack after teardown. In either case the + sid mappings, seq tracker state, and any held iterator must be + cleaned up so they don't leak across reconnects. + """ + # The envelope can carry sid at the top level or nested under msg. + sid = data.get("sid") + if sid is None: + msg_body = data.get("msg") + if isinstance(msg_body, dict): + sid = msg_body.get("sid") + if sid is None: + logger.debug("server unsubscribed envelope without sid: %s", data) + return + + client_id = self._sub_mgr._sid_to_client.pop(sid, None) + if self._seq_tracker is not None: + self._seq_tracker.reset(sid) + + if client_id is None: + logger.debug( + "server unsubscribed for unknown sid %d (already reaped)", sid + ) + return + + sub = self._sub_mgr._subscriptions.pop(client_id, None) + if sub is not None: + # Wake any held iterator so async-for exits cleanly. + await sub.queue.put_sentinel() + logger.info( + "Server-initiated unsubscribe: channel=%s sid=%d client_id=%d", + sub.channel, sid, client_id, + ) + + async def _surface_channel_error( + self, msg_type: str, data: dict[str, Any] + ) -> None: + """Route a channel-level error envelope to on_error (#82). + + AsyncAPI allows errors to ride on a typed message (e.g. a payload + with an `error` field, or an unrecognized type that still carries + error semantics) rather than the top-level type="error" envelope. + Previously these fell through to a debug log and were silently + dropped; surface them via on_error if registered, or log at + WARNING with the envelope contents so the user has a signal. + """ + if self._on_error is not None: + # Best-effort: synthesize an ErrorMessage. If the payload + # doesn't match the strict schema, fall back to model_construct + # so the handler still fires with the raw payload visible. + try: + synthesized = { + "id": data.get("id", 0), + "type": "error", + "msg": data.get("error") + if isinstance(data.get("error"), dict) + else data.get("msg"), + } + error = ErrorMessage.model_validate(synthesized) + except Exception: + logger.error( + "Channel-level error envelope (type=%s) failed strict " + "ErrorMessage validation; firing on_error with raw payload: %s", + msg_type, data, + ) + error = ErrorMessage.model_construct( + id=data.get("id", 0), + type="error", + msg=ErrorPayload.model_construct(code="unknown", msg=str(data)), + ) + await self._on_error(error) + else: + logger.warning( + "Channel-level error envelope (type=%s) with no on_error " + "handler registered: %s", + msg_type, data, + ) + async def dispatch(self, raw: str) -> None: """Parse a raw JSON frame and route it.""" try: @@ -80,6 +182,16 @@ async def dispatch(self, raw: str) -> None: if msg_type == "error" and self._on_error is not None: error = ErrorMessage.model_validate(data) await self._on_error(error) + elif msg_type == "unsubscribed": + await self._handle_server_unsubscribe(data) + return + + # Channel-level error semantics on a non-"error" envelope (#82): + # `data.get("error") is not None` (NOT `"error" in data`) so a + # legitimate `"error": null` field on a typed envelope isn't + # falsely routed as an error. + if data.get("error") is not None: + await self._surface_channel_error(msg_type, data) return # Parse into typed model @@ -105,9 +217,9 @@ async def dispatch(self, raw: str) -> None: logger.debug("Message for unknown sid %d (may be stale)", sid) return - # Check for callback first + # Fan out: deliver to the callback (if any) AND the subscription queue. + # Previously the queue was skipped when a callback was registered, + # which silently stranded any iterator on the same channel (#80). if sub.channel in self._callbacks: await self._callbacks[sub.channel](parsed) - else: - # Route to queue - await sub.queue.put(parsed) + await sub.queue.put(parsed) diff --git a/tests/ws/test_dispatch.py b/tests/ws/test_dispatch.py index 73c5205..ce2a86b 100644 --- a/tests/ws/test_dispatch.py +++ b/tests/ws/test_dispatch.py @@ -3,6 +3,7 @@ import asyncio import json +import logging import pytest @@ -12,6 +13,7 @@ from kalshi.ws.models.market_positions import MarketPositionsMessage from kalshi.ws.models.multivariate import MultivariateMessage from kalshi.ws.models.user_orders import UserOrdersMessage +from kalshi.ws.sequence import SequenceTracker class FakeSubManager: @@ -19,16 +21,34 @@ class FakeSubManager: def __init__(self) -> None: self._subs: dict[int, Subscription] = {} - - def add(self, sid: int, channel: str) -> Subscription: + # Mirror the real manager so dispatcher's cleanup paths work. + self._subscriptions: dict[int, Subscription] = self._subs + self._sid_to_client: dict[int, int] = {} + + def add( + self, sid: int, channel: str, *, client_id: int | None = None, + ) -> Subscription: + """Add a sub. ``client_id`` defaults to ``sid`` for simple tests, but + can be passed distinct to exercise the production mapping path + (server-assigned sid != client-side id). + """ + cid = sid if client_id is None else client_id queue: MessageQueue[object] = MessageQueue(maxsize=100) - sub = Subscription(client_id=sid, channel=channel, params={}, queue=queue) + sub = Subscription(client_id=cid, channel=channel, params={}, queue=queue) sub.server_sid = sid - self._subs[sid] = sub + self._subs[cid] = sub + self._sid_to_client[sid] = cid return sub def get_subscription_by_sid(self, sid: int) -> Subscription | None: - return self._subs.get(sid) + client_id = self._sid_to_client.get(sid) + if client_id is None: + return None + return self._subs.get(client_id) + + @property + def active_subscriptions(self) -> dict[int, Subscription]: + return dict(self._subs) @pytest.mark.asyncio @@ -90,6 +110,7 @@ async def test_dispatch_unknown_sid(self) -> None: await dispatcher.dispatch(raw) # should not crash async def test_callback_mode(self) -> None: + """Callback AND queue both receive the message (fan-out, see #80).""" mgr = FakeSubManager() sub = mgr.add(1, "ticker") received: list[object] = [] @@ -104,7 +125,56 @@ async def on_ticker(msg: object) -> None: ) await dispatcher.dispatch(raw) assert len(received) == 1 - assert sub.queue.qsize() == 0 # callback consumed it, not queue + assert sub.queue.qsize() == 1 # also fanned out to the iterator queue + + async def test_callback_and_iterator_same_channel( + self, caplog: pytest.LogCaptureFixture + ) -> None: + """Regression for #80: callback + iterator on same channel. + + Both must receive every message; the iterator must not silently + stop. A WARNING is logged at register time so the fan-out is not + invisible. + """ + + mgr = FakeSubManager() + sub = mgr.add(1, "ticker") + received: list[object] = [] + + async def on_ticker(msg: object) -> None: + received.append(msg) + + dispatcher = MessageDispatcher(sub_mgr=mgr) # type: ignore[arg-type] + with caplog.at_level(logging.WARNING, logger="kalshi.ws"): + dispatcher.register_callback("ticker", on_ticker) + assert any("active subscription exists" in r.message for r in caplog.records) + + raw = json.dumps( + {"type": "ticker", "sid": 1, "msg": {"market_ticker": "T", "market_id": "x"}} + ) + await dispatcher.dispatch(raw) + await dispatcher.dispatch(raw) + + # Both sinks observed both messages. + assert len(received) == 2 + assert sub.queue.qsize() == 2 + + async def test_register_callback_without_active_sub_no_warning( + self, caplog: pytest.LogCaptureFixture + ) -> None: + """No warning when callback is registered before any subscription.""" + + mgr = FakeSubManager() # no subs + + async def cb(msg: object) -> None: + return None + + dispatcher = MessageDispatcher(sub_mgr=mgr) # type: ignore[arg-type] + with caplog.at_level(logging.WARNING, logger="kalshi.ws"): + dispatcher.register_callback("ticker", cb) + assert not any( + "active subscription exists" in r.message for r in caplog.records + ) async def test_error_callback(self) -> None: mgr = FakeSubManager() @@ -118,6 +188,120 @@ async def on_error(err: object) -> None: await dispatcher.dispatch(raw) assert len(errors) == 1 + async def test_channel_level_error_routed_to_on_error(self) -> None: + """Regression for #82: error semantics on a typed envelope reach on_error. + + Previously a message with type other than "error" but carrying an + `error` payload fell through to the unknown-type log path and was + silently dropped. + """ + mgr = FakeSubManager() + errors: list[object] = [] + + async def on_error(err: object) -> None: + errors.append(err) + + dispatcher = MessageDispatcher(sub_mgr=mgr, on_error=on_error) # type: ignore[arg-type] + raw = json.dumps( + { + "type": "ticker", + "sid": 1, + "error": {"code": 7, "msg": "schema violation"}, + } + ) + await dispatcher.dispatch(raw) + assert len(errors) == 1 + + async def test_channel_level_error_logged_when_no_handler( + self, caplog: pytest.LogCaptureFixture + ) -> None: + """Without on_error, a channel-level error is logged at WARNING (#82).""" + + mgr = FakeSubManager() + dispatcher = MessageDispatcher(sub_mgr=mgr) # type: ignore[arg-type] + raw = json.dumps( + {"type": "ticker", "sid": 1, "error": {"code": 7, "msg": "boom"}} + ) + with caplog.at_level(logging.WARNING, logger="kalshi.ws"): + await dispatcher.dispatch(raw) + assert any( + "Channel-level error envelope" in r.message for r in caplog.records + ) + + async def test_channel_level_error_unrecognized_sid_still_surfaced(self) -> None: + """Error with a sid the dispatcher no longer knows still reaches on_error.""" + mgr = FakeSubManager() # no subs + errors: list[object] = [] + + async def on_error(err: object) -> None: + errors.append(err) + + dispatcher = MessageDispatcher(sub_mgr=mgr, on_error=on_error) # type: ignore[arg-type] + raw = json.dumps( + { + "type": "ticker", + "sid": 12345, # post-unsubscribe race + "error": {"code": 9, "msg": "stale sid"}, + } + ) + await dispatcher.dispatch(raw) + assert len(errors) == 1 + + async def test_null_error_field_not_misrouted(self) -> None: + """Regression: `"error": null` on a typed envelope must parse normally. + + The dispatcher uses `data.get("error") is not None`, not + `"error" in data`, so a legitimate optional error field set to null + doesn't trip the channel-level-error path. + """ + mgr = FakeSubManager() + sub = mgr.add(1, "ticker") + errors: list[object] = [] + + async def on_error(err: object) -> None: + errors.append(err) + + dispatcher = MessageDispatcher(sub_mgr=mgr, on_error=on_error) # type: ignore[arg-type] + raw = json.dumps( + { + "type": "ticker", + "sid": 1, + "error": None, + "msg": {"market_ticker": "T", "market_id": "x"}, + } + ) + await dispatcher.dispatch(raw) + assert errors == [], "null error field misrouted as channel error" + # Normal ticker flow happened: message landed on the queue. + assert sub.queue.qsize() == 1 + + async def test_channel_level_error_validation_failure_still_fires_handler( + self, caplog: pytest.LogCaptureFixture + ) -> None: + """If the error payload fails strict ErrorMessage validation, on_error + still fires with a model_construct'd fallback so handlers always see + a signal. The dispatcher also logs at ERROR (not WARNING) because a + registered handler was almost-not-called. + """ + mgr = FakeSubManager() + errors: list[object] = [] + + async def on_error(err: object) -> None: + errors.append(err) + + dispatcher = MessageDispatcher(sub_mgr=mgr, on_error=on_error) # type: ignore[arg-type] + # Error field is a non-dict, non-string sentinel that fails ErrorPayload validation. + raw = json.dumps( + {"type": "ticker", "sid": 1, "error": 42} + ) + with caplog.at_level(logging.ERROR, logger="kalshi.ws"): + await dispatcher.dispatch(raw) + assert len(errors) == 1, "on_error not called on validation failure" + assert any( + "failed strict ErrorMessage validation" in r.message + for r in caplog.records + ) + async def test_all_channel_types_have_models(self) -> None: """Verify every expected channel type is in the dispatch map.""" expected = { @@ -159,6 +343,92 @@ async def on_ticker(msg: object) -> None: assert len(received) == 0 # callback was removed assert sub.queue.qsize() == 1 # routed to queue instead + async def test_server_unsubscribe_reaps_state(self) -> None: + """Regression for #81: server-initiated unsubscribe reaps sid maps. + + Previously the dispatcher logged and dropped the unsubscribed + envelope, leaking _sid_to_client entries. Now it must pop both + the sid-mapping and the subscription, and push a sentinel so any + held iterator exits cleanly. + """ + mgr = FakeSubManager() + sub = mgr.add(7, "ticker") + dispatcher = MessageDispatcher(sub_mgr=mgr) # type: ignore[arg-type] + assert 7 in mgr._sid_to_client and 7 in mgr._subscriptions + + raw = json.dumps({"type": "unsubscribed", "msg": {"sid": 7}}) + await dispatcher.dispatch(raw) + + assert 7 not in mgr._sid_to_client + assert 7 not in mgr._subscriptions + + # Held iterator exits cleanly via the pushed sentinel. + with pytest.raises(StopAsyncIteration): + await sub.queue.get() + + async def test_server_unsubscribe_top_level_sid(self) -> None: + """Server can also send sid at the top level (UnsubscribedMessage shape).""" + mgr = FakeSubManager() + sub = mgr.add(8, "ticker") + dispatcher = MessageDispatcher(sub_mgr=mgr) # type: ignore[arg-type] + + raw = json.dumps({"type": "unsubscribed", "sid": 8, "seq": 0}) + await dispatcher.dispatch(raw) + + assert 8 not in mgr._sid_to_client + assert 8 not in mgr._subscriptions + with pytest.raises(StopAsyncIteration): + await sub.queue.get() + + async def test_server_unsubscribe_resets_seq_tracker(self) -> None: + """If a SequenceTracker is wired in, reset(sid) is called.""" + + mgr = FakeSubManager() + mgr.add(9, "orderbook_delta") + tracker = SequenceTracker() + # Seed the tracker as if a delta had been processed. + await tracker.track(9, 5, "orderbook_delta") + assert 9 in tracker._last_seq + + dispatcher = MessageDispatcher( + sub_mgr=mgr, # type: ignore[arg-type] + seq_tracker=tracker, + ) + raw = json.dumps({"type": "unsubscribed", "msg": {"sid": 9}}) + await dispatcher.dispatch(raw) + + assert 9 not in tracker._last_seq + + async def test_server_unsubscribe_with_distinct_client_id(self) -> None: + """Server-assigned sid != client-side id is the production case. + + Exercises the two-step lookup: server's `sid` → `_sid_to_client[sid]` + gives `client_id`, then `_subscriptions.pop(client_id)`. Collapsing + sid == client_id (default) lets the test pass by accident even if + the lookup is wrong. + """ + mgr = FakeSubManager() + sub = mgr.add(sid=500, channel="ticker", client_id=1) + dispatcher = MessageDispatcher(sub_mgr=mgr) # type: ignore[arg-type] + assert 500 in mgr._sid_to_client + assert 1 in mgr._subscriptions + assert 500 not in mgr._subscriptions # not keyed by server sid + + raw = json.dumps({"type": "unsubscribed", "msg": {"sid": 500}}) + await dispatcher.dispatch(raw) + + assert 500 not in mgr._sid_to_client + assert 1 not in mgr._subscriptions + with pytest.raises(StopAsyncIteration): + await sub.queue.get() + + async def test_server_unsubscribe_unknown_sid_no_crash(self) -> None: + """A late/duplicate unsubscribed for an already-reaped sid is a no-op.""" + 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 + async def test_dispatch_message_without_sid(self) -> None: """Messages without sid are logged but don't crash.""" mgr = FakeSubManager()