From 135660cee1e4c8e15cb58f262eeb95e86c84462b Mon Sep 17 00:00:00 2001 From: Jeff West Date: Thu, 21 May 2026 08:25:41 -0500 Subject: [PATCH] fix(ws): release orphan subscribe acks; require snapshot yes/no (#268) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two small WS polish items from the post-v2.4 audit: Item 1 — cancelled mid-subscribe leak. ``SubscriptionManager.subscribe`` sends a ``subscribe`` cmd then awaits ``_wait_for_response``; if the awaiting task is cancelled between send and ack, the lock releases but the server still completes the subscribe and emits a ``subscribed`` envelope. That ack lands in ``MessageDispatcher.dispatch`` with no ``sid -> client_id`` mapping installed, was silently dropped, and the server held an active sid until the WS closed (data frames thereafter logged "Message for unknown sid"; the leaked sub counted against any server-side quota). Detect the orphan in the dispatcher: when a ``subscribed`` control envelope arrives with a sid not in ``_sub_mgr._sid_to_client``, log a WARNING and send an ``unsubscribe`` back through the connection to release the server-side state. Best-effort: a send failure is logged and swallowed (socket is likely already gone, in which case the server reaps on close anyway). Sids that already have a mapping (normal-path race where the mapping landed first) are a no-op. Item 2 — ``OrderbookSnapshotPayload`` ``yes``/``no`` no longer have a ``default_factory=dict``. A wire frame missing both ``yes_dollars_fp`` and ``no_dollars_fp`` was previously parsing as a fully empty snapshot — silently resetting the local book to empty on schema drift. Now both sides are required so a malformed snapshot raises ``ValidationError``, which pairs with the #241 seq-rollback path so the watermark isn't silently advanced past the never-delivered frame. (The mutable-default anti-pattern was already fixed by #263 switching to ``default_factory``; this drops the default entirely.) Tests: - 4 orphan-subscribed dispatcher tests (sends unsubscribe with nested/top-level sid, no-op on known sid, send-failure swallowed). - Snapshot regression: missing ``yes`` / missing ``no`` / missing both each raise ``ValidationError``. Closes #268 --- kalshi/ws/dispatch.py | 55 +++++++++++++++ kalshi/ws/models/orderbook_delta.py | 9 ++- tests/ws/test_dispatch.py | 101 ++++++++++++++++++++++++++++ tests/ws/test_models.py | 21 ++++++ 4 files changed, 184 insertions(+), 2 deletions(-) diff --git a/kalshi/ws/dispatch.py b/kalshi/ws/dispatch.py index a85c0661..cc5b2fbd 100644 --- a/kalshi/ws/dispatch.py +++ b/kalshi/ws/dispatch.py @@ -131,6 +131,59 @@ async def _handle_server_unsubscribe(self, data: dict[str, Any]) -> None: sub.channel, sid, client_id, ) + async def _handle_orphan_subscribed(self, data: dict[str, Any]) -> None: + """Release a server-side sid whose subscribe ack has no client mapping. + + Normal subscribes consume the ``subscribed`` ack inside + ``SubscriptionManager._wait_for_response`` and install the + ``sid -> client_id`` mapping before the dispatcher ever sees it. + If the awaiting task is cancelled between send and ack (#268), + the lock is released but the server still completes the + subscribe; the ack lands here with a sid that has no mapping. + Without intervention the server holds an active sid the client + can never address again — subsequent data frames log "Message + for unknown sid" and the sub counts against the server quota + until the WS closes. + + Best-effort: log a WARNING and send an ``unsubscribe`` for the + orphan sid. Any send failure is swallowed (the connection is + likely already gone, in which case the server reaps the sid on + socket close anyway). + """ + sid = data.get("sid") + if sid is None: + msg_body = data.get("msg") + if isinstance(msg_body, dict): + sid = msg_body.get("sid") + if not isinstance(sid, int): + logger.debug( + "Orphan-subscribed handler: no int sid in envelope: %s", data, + ) + return + if sid in self._sub_mgr._sid_to_client: + # A normal subscribe completed and installed the mapping + # before this handler ran — nothing to do. + return + logger.warning( + "Orphan subscribed ack for sid=%d (no client mapping); " + "sending unsubscribe to release server-side state.", + sid, + ) + try: + msg_id = self._sub_mgr._get_msg_id() + cmd = { + "id": msg_id, + "cmd": "unsubscribe", + "params": {"sids": [sid]}, + } + await self._sub_mgr._connection.send(cmd) + except Exception: + logger.warning( + "Failed to send orphan-unsubscribe for sid=%d", + sid, + exc_info=True, + ) + async def _surface_channel_error( self, msg_type: str, data: dict[str, Any] ) -> None: @@ -200,6 +253,8 @@ async def dispatch( await self._on_error(error) elif msg_type == "unsubscribed": await self._handle_server_unsubscribe(data) + elif msg_type == "subscribed": + await self._handle_orphan_subscribed(data) return # Channel-level error semantics on a non-"error" envelope (#82): diff --git a/kalshi/ws/models/orderbook_delta.py b/kalshi/ws/models/orderbook_delta.py index ac5c10d1..b3d918fe 100644 --- a/kalshi/ws/models/orderbook_delta.py +++ b/kalshi/ws/models/orderbook_delta.py @@ -62,16 +62,21 @@ class OrderbookSnapshotPayload(BaseModel): These are validated directly into ``dict[Decimal, Decimal]`` (price -> count) in a single walk so ``OrderbookManager._apply_snapshot_inplace`` can adopt the map with no second iteration (#263). + + Both sides are required (#268): a snapshot envelope missing + ``yes_dollars_fp`` or ``no_dollars_fp`` is schema drift or a partial + server response, not "an empty book on that side". Surfacing a + ``ValidationError`` lets the recv loop's malformed-frame handler + log it and roll back the seq watermark for #241 rather than silently + resetting the local book to empty. """ market_ticker: str market_id: str yes: PriceCountMap = Field( - default_factory=dict, validation_alias=AliasChoices("yes_dollars_fp", "yes"), ) no: PriceCountMap = Field( - default_factory=dict, validation_alias=AliasChoices("no_dollars_fp", "no"), ) model_config = {"extra": "allow", "populate_by_name": True} diff --git a/tests/ws/test_dispatch.py b/tests/ws/test_dispatch.py index f368e7c7..b4d4abc0 100644 --- a/tests/ws/test_dispatch.py +++ b/tests/ws/test_dispatch.py @@ -478,6 +478,107 @@ async def test_dispatch_malformed_message_raises_for_seq_rollback(self) -> None: with pytest.raises(ValidationError): await dispatcher.dispatch(json.loads(raw)) + async def test_orphan_subscribed_sends_unsubscribe( + self, caplog: pytest.LogCaptureFixture + ) -> None: + """#268: A `subscribed` ack with no client_id mapping is the cancel- + mid-subscribe race. The dispatcher must release the server-side sid + by emitting an `unsubscribe` so the leak doesn't persist for the + lifetime of the connection. + """ + sent: list[dict[str, object]] = [] + + class _FakeConnection: + async def send(self, msg: dict[str, object]) -> None: + sent.append(msg) + + mgr = FakeSubManager() + mgr._connection = _FakeConnection() # type: ignore[attr-defined] + mgr._next_msg_id = 42 # type: ignore[attr-defined] + + def _get_msg_id() -> int: + mid = mgr._next_msg_id # type: ignore[attr-defined] + mgr._next_msg_id += 1 # type: ignore[attr-defined] + return mid + + mgr._get_msg_id = _get_msg_id # type: ignore[attr-defined] + dispatcher = MessageDispatcher(sub_mgr=mgr) # type: ignore[arg-type] + + raw = json.dumps({"type": "subscribed", "id": 1, "msg": {"sid": 777}}) + with caplog.at_level(logging.WARNING, logger="kalshi.ws"): + await dispatcher.dispatch(json.loads(raw)) + + assert sent == [ + {"id": 42, "cmd": "unsubscribe", "params": {"sids": [777]}} + ] + assert any( + "Orphan subscribed ack for sid=777" in rec.message + for rec in caplog.records + ) + + async def test_orphan_subscribed_with_known_sid_is_noop(self) -> None: + """When `subscribed` arrives for a sid already in _sid_to_client, + the orphan handler must NOT send a spurious unsubscribe (#268). + """ + sent: list[dict[str, object]] = [] + + class _FakeConnection: + async def send(self, msg: dict[str, object]) -> None: + sent.append(msg) + + mgr = FakeSubManager() + mgr.add(123, "ticker") + mgr._connection = _FakeConnection() # type: ignore[attr-defined] + mgr._get_msg_id = lambda: 1 # type: ignore[attr-defined] + dispatcher = MessageDispatcher(sub_mgr=mgr) # type: ignore[arg-type] + + raw = json.dumps({"type": "subscribed", "id": 1, "msg": {"sid": 123}}) + await dispatcher.dispatch(json.loads(raw)) + assert sent == [] + + async def test_orphan_subscribed_top_level_sid(self) -> None: + """`subscribed` envelope can also carry sid at the top level (#268).""" + sent: list[dict[str, object]] = [] + + class _FakeConnection: + async def send(self, msg: dict[str, object]) -> None: + sent.append(msg) + + mgr = FakeSubManager() + mgr._connection = _FakeConnection() # type: ignore[attr-defined] + mgr._get_msg_id = lambda: 99 # type: ignore[attr-defined] + dispatcher = MessageDispatcher(sub_mgr=mgr) # type: ignore[arg-type] + + raw = json.dumps({"type": "subscribed", "id": 1, "sid": 555}) + await dispatcher.dispatch(json.loads(raw)) + assert sent == [ + {"id": 99, "cmd": "unsubscribe", "params": {"sids": [555]}} + ] + + async def test_orphan_subscribed_send_failure_swallowed( + self, caplog: pytest.LogCaptureFixture + ) -> None: + """Send failure during orphan-unsubscribe is best-effort: log only, + do not propagate (the socket is likely already gone) (#268). + """ + + class _BrokenConnection: + async def send(self, msg: dict[str, object]) -> None: + raise ConnectionError("socket closed") + + mgr = FakeSubManager() + mgr._connection = _BrokenConnection() # type: ignore[attr-defined] + mgr._get_msg_id = lambda: 1 # type: ignore[attr-defined] + dispatcher = MessageDispatcher(sub_mgr=mgr) # type: ignore[arg-type] + + raw = json.dumps({"type": "subscribed", "id": 1, "msg": {"sid": 888}}) + with caplog.at_level(logging.WARNING, logger="kalshi.ws"): + await dispatcher.dispatch(json.loads(raw)) # must not raise + assert any( + "Failed to send orphan-unsubscribe for sid=888" in rec.message + for rec in caplog.records + ) + @pytest.mark.asyncio async def test_dispatch_routes_user_order_singular() -> None: diff --git a/tests/ws/test_models.py b/tests/ws/test_models.py index 5e73f4a4..6be0cccb 100644 --- a/tests/ws/test_models.py +++ b/tests/ws/test_models.py @@ -1085,6 +1085,27 @@ def test_orderbook_snapshot_yes_no_parse_as_decimal_tuples(self) -> None: assert isinstance(key, Decimal) assert isinstance(value, Decimal) + def test_orderbook_snapshot_missing_sides_raises(self) -> None: + """#268: dropping the mutable default means a snapshot envelope + missing both `yes_dollars_fp` and `no_dollars_fp` (schema drift or + a partial server response) must surface as a ValidationError + rather than silently materializing an empty book. + """ + from kalshi.ws.models.orderbook_delta import OrderbookSnapshotPayload + + with pytest.raises(ValidationError): + OrderbookSnapshotPayload.model_validate( + {"market_ticker": "T", "market_id": "x"} + ) + with pytest.raises(ValidationError): + OrderbookSnapshotPayload.model_validate( + {"market_ticker": "T", "market_id": "x", "yes": []} + ) + with pytest.raises(ValidationError): + OrderbookSnapshotPayload.model_validate( + {"market_ticker": "T", "market_id": "x", "no": []} + ) + def test_rfq_created_contracts_parses_as_decimal(self) -> None: payload = RfqCreatedPayload.model_validate( {