Skip to content

_stop() cancels recv task and broadcasts sentinels before closing the connection — late frames drop silently #357

Closed

Description

@TexasCoding

Problem

KalshiWebSocket._stop() runs in this order: (1) clear _running, (2) wait/cancel the recv task, (3) broadcast queue sentinels, (4) await self._connection.close(). The canonical close-then-drain pattern is reversed:

  • After step 2 the recv loop is gone, so any frame still buffered inside the websockets library at step 3 will never be drained.
  • After step 3 the queues are closed; any frame that arrives between steps 2 and 4 is silently dropped with no signal to the user.
  • connection.close() should run before the recv-task cancel so the recv loop can drain pending frames that arrive between the client-side close and the server's close-ack. Today the recv loop is forcibly cancelled first, then close() runs with no reader behind it.

For orderbook_delta consumers that need to observe the last few frames before a clean shutdown (e.g. final fills against a position) the data is gone with no log indication.

Evidence

  • kalshi/ws/client.py:177-211
async def _stop(self) -> None:
    self._running = False
    self._resume_signal.set()
    if self._recv_task and not self._recv_task.done():
        with contextlib.suppress(asyncio.CancelledError, Exception):
            await asyncio.wait_for(self._recv_task, timeout=2.0)
        if not self._recv_task.done():
            self._recv_task.cancel()
            with contextlib.suppress(asyncio.CancelledError, Exception):
                await self._recv_task

    await self._broadcast_sentinels()

    if self._connection:
        await self._connection.close()

Suggested fix

Reorder _stop() to mirror the close-then-drain pattern already used by run_forever(stop_event=...):

  1. self._running = False; set _resume_signal.
  2. await self._connection.close() first — server stops sending, in-flight frames drain through the still-running recv loop, which sees ConnectionClosed and exits at the not self._running check.
  3. Wait/cancel the recv task as today.
  4. Broadcast sentinels last so iterator consumers see the final drained frames before the sentinel.

Add a test that publishes a final frame from the server side immediately before close and asserts the iterator yields it before the sentinel.

Source

Round-3 independent audit (reviewer: websocket).

Metadata

Metadata

Assignees

No one assigned

    Labels

    bugSomething isn't workingwsWebSocket-related

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions