From fe41d62c2878969080d63b70ccc770d6d462e55c Mon Sep 17 00:00:00 2001 From: Jeff West Date: Fri, 22 May 2026 07:37:59 -0500 Subject: [PATCH 1/2] fix(ws): wrap TimeoutError; pop subscription on resubscribe failure MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ChannelsManager._wait_for_response only caught ConnectionClosed; the bare TimeoutError raised by asyncio.wait_for escaped to subscribe() / unsubscribe() / update_subscription() callers. The top-of-loop 'remaining <= 0' guard never reached because the next iteration never ran. Add a TimeoutError handler that raises the same structured KalshiSubscriptionError with channel / client_id / op populated, so callers branching on SDK exception types (#213) retain the recovery surface for the timeout case. SubscriptionManager.broadcast_error put an error sentinel on the queue but never popped the Subscription. After a failed resubscribe_one in _handle_seq_gap (or the KalshiBackpressureError rollback in _process_frame), the dead sub stayed in _subscriptions with server_sid cleared — user-driven unsubscribe short-circuited on the stuck entry, the next reconnect's resubscribe_all resurrected it on the server with a fresh sid, and the dispatcher routed frames into the closed queue (silent data loss + server-quota leak). Pop the Subscription from _subscriptions / _sid_to_client after delivering the sentinel so the zombie cannot resurrect. Closes #314, #315 --- kalshi/ws/channels.py | 25 ++++++ tests/ws/test_client.py | 169 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 194 insertions(+) diff --git a/kalshi/ws/channels.py b/kalshi/ws/channels.py index b4d753ef..6622e6b1 100644 --- a/kalshi/ws/channels.py +++ b/kalshi/ws/channels.py @@ -206,6 +206,19 @@ async def _wait_for_response( raise KalshiConnectionError( f"Connection closed while awaiting response to command {msg_id}" ) from e + except TimeoutError: + # #314: asyncio.wait_for surfaces a bare TimeoutError that the + # top-of-loop ``remaining <= 0`` guard never gets to reach + # (the next iteration doesn't run). Raise the same structured + # KalshiSubscriptionError so consumers branching on SDK + # exception types (#213) retain the channel/client_id/op + # surface for timeout recovery. + raise KalshiSubscriptionError( + f"Timed out waiting for response to command {msg_id}", + channel=channel, + client_id=client_id, + op=op, # type: ignore[arg-type] + ) from None data: dict[str, Any] = self._json_loads(raw) if data.get("id") == msg_id: return data @@ -562,6 +575,14 @@ async def broadcast_error( (it's assigned at subscribe-ack and may change on resubscribe) so enrich it here from the subscription's current ``server_sid`` when missing. + + #315: the error sentinel closes the queue permanently, so the + subscription is dead from the consumer's perspective. Pop it + from ``_subscriptions`` / ``_sid_to_client`` so the next + reconnect's ``resubscribe_all`` doesn't resurrect a zombie + sub on the server (silent data loss + server-quota leak) and + so user-driven ``unsubscribe`` doesn't short-circuit on a + permanently-stuck entry. """ sub = self._subscriptions.get(client_id) if sub is None: @@ -573,6 +594,10 @@ async def broadcast_error( ): exc.sid = sub.server_sid await sub.queue.put_error(exc) + # #315: drop the dead subscription so resubscribe_all skips it. + self._subscriptions.pop(client_id, None) + if sub.server_sid is not None: + self._sid_to_client.pop(sub.server_sid, None) def take_stash(self) -> dict[int, collections.deque[str]]: """Return and clear the resubscribe-window stash atomically (#176). diff --git a/tests/ws/test_client.py b/tests/ws/test_client.py index 98307d9a..a6fa198b 100644 --- a/tests/ws/test_client.py +++ b/tests/ws/test_client.py @@ -598,3 +598,172 @@ def my_loads(raw: bytes | str) -> Any: # At minimum the ticker frame was parsed by our loader. The recv loop # parses every frame off the socket — assert non-empty. assert calls, "ws_json_loads was never invoked on recv frames" + + +# --------------------------------------------------------------------------- +# Regression — #314, #315 +# --------------------------------------------------------------------------- + + +class TestIssue314WaitForResponseTimeout: + async def test_issue_314_subscribe_raises_kalshi_subscription_error_on_timeout( + self, + ) -> None: + """#314: ``asyncio.wait_for`` inside ``_wait_for_response`` raises a + bare ``TimeoutError`` when the server never answers a subscribe ack. + Before the fix it escaped to ``subscribe()`` callers as a plain + ``TimeoutError``, losing the structured ``channel`` / ``client_id`` / + ``op`` surface promised by #213. After the fix it surfaces as a + ``KalshiSubscriptionError`` with those fields populated.""" + from kalshi.ws.channels import SubscriptionManager + + class NeverRecv: + async def send(self, _cmd: Any) -> None: + return None + + async def recv(self) -> str: + # Hang forever; the wait_for inside _wait_for_response is + # what must time out and surface as KalshiSubscriptionError. + await asyncio.Event().wait() + raise AssertionError("unreachable") # pragma: no cover + + mgr = SubscriptionManager(NeverRecv()) # type: ignore[arg-type] + + with pytest.raises(KalshiSubscriptionError) as excinfo: + await mgr._wait_for_response( + msg_id=42, + timeout=0.05, + channel="orderbook_delta", + client_id=7, + op="subscribe", + ) + + err = excinfo.value + assert err.channel == "orderbook_delta" + assert err.client_id == 7 + assert err.op == "subscribe" + assert "42" in str(err) + + +class TestIssue315ZombieSubscriptionCleanup: + async def test_issue_315_failed_resubscribe_removes_zombie_subscription( + self, + ) -> None: + """#315: ``broadcast_error`` must pop the dead ``Subscription`` so the + next reconnect's ``resubscribe_all`` doesn't resurrect a zombie sub on + the server (silent data loss + server-quota leak) and ``unsubscribe`` + doesn't short-circuit on a permanently-stuck entry.""" + from kalshi.errors import KalshiSequenceGapError + from kalshi.ws.backpressure import MessageQueue, OverflowStrategy + from kalshi.ws.channels import Subscription, SubscriptionManager + + recorded: list[Any] = [] + + class StubConn: + async def send(self, cmd: Any) -> None: + recorded.append(cmd) + + async def recv(self) -> str: # pragma: no cover - not used + await asyncio.Event().wait() + raise AssertionError("unreachable") + + mgr = SubscriptionManager(StubConn()) # type: ignore[arg-type] + queue: MessageQueue[Any] = MessageQueue( + maxsize=10, + overflow=OverflowStrategy.ERROR, + channel="orderbook_delta", + client_id=7, + ) + sub = Subscription( + client_id=7, + channel="orderbook_delta", + params={"market_tickers": ["ABC-YES"]}, + queue=queue, + ) + sub.server_sid = 314 + mgr._subscriptions[7] = sub + mgr._sid_to_client[314] = 7 + + # Simulate the gap-recovery failure path: ``resubscribe_one`` clears + # ``server_sid`` and the ``_sid_to_client`` entry before raising, then + # ``_handle_seq_gap`` calls broadcast_error. After the fix + # ``broadcast_error`` must also pop the dead Subscription so the + # next reconnect's ``resubscribe_all`` doesn't resurrect it. + sub.server_sid = None + mgr._sid_to_client.pop(314, None) + await mgr.broadcast_error( + 7, + KalshiSequenceGapError( + "Resubscribe failed for orderbook_delta after gap", + channel="orderbook_delta", + sid=314, + client_id=7, + last_seq=10, + next_seq=12, + ), + ) + + # Subscription is gone — no zombie left behind. + assert 7 not in mgr._subscriptions + assert 314 not in mgr._sid_to_client + + # The iterator surfaces the error sentinel. + with pytest.raises(Exception) as excinfo: + async for _ in queue: + pass + assert "Resubscribe failed" in str(excinfo.value) + + # A subsequent reconnect's resubscribe_all must be a no-op — no + # subscribe command should hit the wire for the dead client_id. + recorded.clear() + await mgr.resubscribe_all() + assert recorded == [] + assert mgr._subscriptions == {} + + async def test_issue_315_backpressure_path_also_removes_subscription( + self, + ) -> None: + """#315: the ``KalshiBackpressureError`` rollback path in + ``_process_frame`` routes through the same ``broadcast_error``, so the + zombie cleanup must apply there too.""" + from kalshi.errors import KalshiBackpressureError + from kalshi.ws.backpressure import MessageQueue, OverflowStrategy + from kalshi.ws.channels import Subscription, SubscriptionManager + + class StubConn: + async def send(self, _cmd: Any) -> None: + return None + + async def recv(self) -> str: # pragma: no cover - not used + await asyncio.Event().wait() + raise AssertionError("unreachable") + + mgr = SubscriptionManager(StubConn()) # type: ignore[arg-type] + queue: MessageQueue[Any] = MessageQueue( + maxsize=1, + overflow=OverflowStrategy.ERROR, + channel="orderbook_delta", + client_id=11, + ) + sub = Subscription( + client_id=11, + channel="orderbook_delta", + params={}, + queue=queue, + ) + sub.server_sid = 999 + mgr._subscriptions[11] = sub + mgr._sid_to_client[999] = 11 + + await mgr.broadcast_error( + 11, + KalshiBackpressureError( + "queue full", + channel="orderbook_delta", + client_id=11, + maxsize=1, + ), + ) + + assert 11 not in mgr._subscriptions + assert 999 not in mgr._sid_to_client From 09ea50ae1f79285913a83f2c555b866e1a56a17e Mon Sep 17 00:00:00 2001 From: Jeff West Date: Fri, 22 May 2026 07:59:07 -0500 Subject: [PATCH 2/2] polish(ws): exercise _sid_to_client cleanup in gap-recovery test; trim comments --- kalshi/ws/channels.py | 15 ++------------- tests/ws/test_client.py | 13 ++++++------- 2 files changed, 8 insertions(+), 20 deletions(-) diff --git a/kalshi/ws/channels.py b/kalshi/ws/channels.py index 6622e6b1..a8ebbb51 100644 --- a/kalshi/ws/channels.py +++ b/kalshi/ws/channels.py @@ -207,12 +207,7 @@ async def _wait_for_response( f"Connection closed while awaiting response to command {msg_id}" ) from e except TimeoutError: - # #314: asyncio.wait_for surfaces a bare TimeoutError that the - # top-of-loop ``remaining <= 0`` guard never gets to reach - # (the next iteration doesn't run). Raise the same structured - # KalshiSubscriptionError so consumers branching on SDK - # exception types (#213) retain the channel/client_id/op - # surface for timeout recovery. + # #314: wrap bare TimeoutError so callers retain channel/client_id/op (#213). raise KalshiSubscriptionError( f"Timed out waiting for response to command {msg_id}", channel=channel, @@ -576,13 +571,7 @@ async def broadcast_error( so enrich it here from the subscription's current ``server_sid`` when missing. - #315: the error sentinel closes the queue permanently, so the - subscription is dead from the consumer's perspective. Pop it - from ``_subscriptions`` / ``_sid_to_client`` so the next - reconnect's ``resubscribe_all`` doesn't resurrect a zombie - sub on the server (silent data loss + server-quota leak) and - so user-driven ``unsubscribe`` doesn't short-circuit on a - permanently-stuck entry. + #315: pop the dead subscription so the next resubscribe_all can't resurrect a zombie sub. """ sub = self._subscriptions.get(client_id) if sub is None: diff --git a/tests/ws/test_client.py b/tests/ws/test_client.py index a6fa198b..775fec80 100644 --- a/tests/ws/test_client.py +++ b/tests/ws/test_client.py @@ -684,13 +684,12 @@ async def recv(self) -> str: # pragma: no cover - not used mgr._subscriptions[7] = sub mgr._sid_to_client[314] = 7 - # Simulate the gap-recovery failure path: ``resubscribe_one`` clears - # ``server_sid`` and the ``_sid_to_client`` entry before raising, then - # ``_handle_seq_gap`` calls broadcast_error. After the fix - # ``broadcast_error`` must also pop the dead Subscription so the - # next reconnect's ``resubscribe_all`` doesn't resurrect it. - sub.server_sid = None - mgr._sid_to_client.pop(314, None) + # Simulate the gap-recovery failure path: ``_handle_seq_gap`` calls + # broadcast_error on the live subscription (server_sid still set, sid + # still in ``_sid_to_client``) after ``resubscribe_one`` raises. The + # fix must clear *both* maps so the zombie can't resurrect. + assert sub.server_sid == 314 + assert mgr._sid_to_client[314] == 7 await mgr.broadcast_error( 7, KalshiSequenceGapError(