perf(ws): inplace orderbook updates; config knobs; pluggable JSON; small polish

kalshi/config.py

@@ -3,7 +3,9 @@
 from __future__ import annotations
 
 import logging
+from collections.abc import Callable
 from dataclasses import dataclass, field
+from typing import Any
 from urllib.parse import urlparse
 
 import httpx
@@ -44,6 +46,14 @@ class KalshiConfig:
             Requires the ``h2`` package (install ``httpx[http2]`` or ``h2``).
         limits: Custom ``httpx.Limits`` for connection pool tuning. ``None``
             uses httpx defaults.
+        ws_ping_interval: Interval (s) between WS keepalive pings. Default 20.
+        ws_close_timeout: Time (s) to wait for graceful WS close handshake.
+            Default 5.
+        ws_json_loads: Optional callable used to parse WS frames. ``None``
+            falls back to :func:`json.loads`. Set to e.g. ``orjson.loads``
+            for ~2-3x faster recv-loop parsing on high-volume streams.
+        ws_json_dumps: Optional callable used to serialize outbound WS
+            commands. ``None`` falls back to :func:`json.dumps`.
     """
 
     base_url: str = PRODUCTION_BASE_URL  # trailing slash is stripped automatically
@@ -57,6 +67,10 @@ class KalshiConfig:
     ws_max_retries: int = DEFAULT_WS_MAX_RETRIES
     http2: bool = False
     limits: httpx.Limits | None = None
+    ws_ping_interval: float = 20.0
+    ws_close_timeout: float = 5.0
+    ws_json_loads: Callable[[bytes | str], Any] | None = None
+    ws_json_dumps: Callable[[Any], bytes | str] | None = None
 
     def __post_init__(self) -> None:
         # Strip trailing slash to prevent double-slash in auth signing paths

kalshi/ws/channels.py

@@ -6,6 +6,7 @@
 import collections
 import json
 import logging
+from collections.abc import Callable
 from typing import Any
 
 from websockets.exceptions import ConnectionClosed
@@ -116,8 +117,14 @@ def __init__(
         connection: ConnectionManager,
         *,
         stash_maxlen: int = 1000,
+        json_loads: Callable[[bytes | str], Any] | None = None,
     ) -> None:
         self._connection = connection
+        # #209: pluggable JSON loader for subscribe-ack frames. ``None``
+        # falls back to stdlib ``json.loads``.
+        self._json_loads: Callable[[bytes | str], Any] = (
+            json_loads if json_loads is not None else json.loads
+        )
         self._subscriptions: dict[int, Subscription] = {}  # client_id -> Subscription
         self._sid_to_client: dict[int, int] = {}  # server_sid -> client_id
         self._next_client_id = 1
@@ -195,7 +202,7 @@ async def _wait_for_response(
                 raise KalshiConnectionError(
                     f"Connection closed while awaiting response to command {msg_id}"
                 ) from e
-            data: dict[str, Any] = json.loads(raw)
+            data: dict[str, Any] = self._json_loads(raw)
             if data.get("id") == msg_id:
                 return data
             # Non-matching frame. During resubscribe, stash it by sid for

kalshi/ws/client.py

@@ -90,6 +90,11 @@ def __init__(
         self._running = False
         self._subscribe_lock = asyncio.Lock()
         self._pending_callbacks: list[tuple[str, Callable[..., Awaitable[None]]]] = []
+        # #209: pluggable JSON loader (None -> stdlib json.loads). Resolved
+        # once here so the recv hot path doesn't keep dereferencing config.
+        self._json_loads: Callable[[bytes | str], Any] = (
+            config.ws_json_loads if config.ws_json_loads is not None else json.loads
+        )
 
     def connect(self) -> _WebSocketSession:
         """Return an async context manager for the WebSocket session."""
@@ -105,7 +110,7 @@ async def _start(self) -> None:
         )
         await self._connection.connect()
 
-        self._sub_mgr = SubscriptionManager(self._connection)
+        self._sub_mgr = SubscriptionManager(self._connection, json_loads=self._json_loads)
         self._seq_tracker = SequenceTracker(on_gap=self._handle_seq_gap)
         self._orderbook_mgr = OrderbookManager()
         self._dispatcher = MessageDispatcher(
@@ -138,10 +143,9 @@ async def _broadcast_sentinels(self) -> None:
         """Put a shutdown sentinel on every active subscription queue.
 
         Iterator consumers see the sentinel and exit their ``async for``
-        loops. ``MessageQueue.put_sentinel`` is idempotent — calling
-        twice puts two sentinels in the queue, the iterator reads the
-        first and raises ``StopAsyncIteration`` without ever touching
-        the second.
+        loops. ``MessageQueue.put_sentinel`` is idempotent — once the
+        queue is closed, subsequent calls are no-ops (no extra sentinel
+        appended, no eviction of an in-flight item).
 
         Used by ``_stop()`` (post-cancel cleanup), the cooperative
         shutdown branch of ``run_forever(stop_event=...)`` (#177), and
@@ -283,7 +287,7 @@ async def _process_frame(self, raw: str) -> None:
         desynced-and-still-consumed.
         """
         assert self._dispatcher is not None
-        data = json.loads(raw)
+        data = self._json_loads(raw)
         sid = data.get("sid")
         seq = data.get("seq")
         msg_type = data.get("type", "")
@@ -317,13 +321,18 @@ async def _process_frame(self, raw: str) -> None:
         pre_validated: BaseModel | None = None
         if msg_type == "orderbook_snapshot" and self._orderbook_mgr:
             snapshot = OrderbookSnapshotMessage.model_validate(data)
+            # #199: mutate in place — skip the O(n log n) sort + 2N
+            # OrderbookLevel allocation that the public ``apply_snapshot``
+            # wrapper performs. Consumers materialize via ``get(ticker)``
+            # on demand.
             # Record per-sid ticker ownership so all-markets subscriptions
             # can be torn down on gap recovery / unsubscribe (#189, #206).
-            self._orderbook_mgr.apply_snapshot(snapshot, sid=snapshot.sid)
+            self._orderbook_mgr._apply_snapshot_inplace(snapshot, sid=snapshot.sid)
             pre_validated = snapshot
         elif msg_type == "orderbook_delta" and self._orderbook_mgr:
             delta = OrderbookDeltaMessage.model_validate(data)
-            self._orderbook_mgr.apply_delta(delta)
+            # #199: in-place mutation only — see snapshot branch above.
+            self._orderbook_mgr._apply_delta_inplace(delta)
             pre_validated = delta
 
         try:

kalshi/ws/connection.py

@@ -64,6 +64,11 @@ def __init__(
         self._on_state_change = on_state_change
         self._ws: ClientConnection | None = None
         self._state = ConnectionState.DISCONNECTED
+        # #209: pluggable JSON dumper for outbound frames. Resolved once at
+        # init so the send hot path doesn't keep dereferencing config.
+        self._json_dumps: Callable[[Any], bytes | str] = (
+            config.ws_json_dumps if config.ws_json_dumps is not None else json.dumps
+        )
 
     @property
     def state(self) -> ConnectionState:
@@ -107,6 +112,22 @@ async def _build_auth_headers(self) -> dict[str, str]:
         ws_path = urlparse(self._config.ws_base_url).path
         return await self._auth.sign_request_async("GET", ws_path)
 
+    async def _open_socket(self) -> ClientConnection:
+        """Open the websocket. Single source of truth for connect args.
+
+        Both ``connect`` and ``reconnect`` route through this so the
+        ping/close timeouts and auth-header build stay in lockstep
+        (#208).
+        """
+        headers = await self._build_auth_headers()
+        return await connect(
+            self._config.ws_base_url,
+            additional_headers=headers,
+            ping_interval=self._config.ws_ping_interval,
+            ping_timeout=self._heartbeat_timeout,
+            close_timeout=self._config.ws_close_timeout,
+        )
+
     async def connect(self) -> None:
         """Establish a WebSocket connection with RSA-PSS auth headers.
 
@@ -117,14 +138,7 @@ async def connect(self) -> None:
         """
         await self._set_state(ConnectionState.CONNECTING)
         try:
-            headers = await self._build_auth_headers()
-            self._ws = await connect(
-                self._config.ws_base_url,
-                additional_headers=headers,
-                ping_interval=20,
-                ping_timeout=self._heartbeat_timeout,
-                close_timeout=5.0,
-            )
+            self._ws = await self._open_socket()
             await self._set_state(ConnectionState.CONNECTED)
         except Exception as e:
             await self._set_state(ConnectionState.CLOSED)
@@ -139,11 +153,9 @@ async def connect(self) -> None:
             ) from e
 
     async def reconnect(self) -> None:
-        """Reconnect with exponential backoff and jitter.
+        """Reconnect with AWS Full Jitter backoff (matches REST transport).
 
-        Uses the same backoff pattern as the REST transport:
-            delay = retry_base_delay * (2 ** attempt) + random.uniform(0, 0.5)
-            delay = min(delay, retry_max_delay)
+            delay = random.uniform(0, min(retry_max_delay, retry_base_delay * 2 ** attempt))
 
         Transitions: -> RECONNECTING -> CONNECTING -> CONNECTED
                      or -> CLOSED (if max retries exceeded)
@@ -159,10 +171,18 @@ async def reconnect(self) -> None:
 
         await self._set_state(ConnectionState.RECONNECTING)
         for attempt in range(self._config.ws_max_retries):
-            delay = self._config.retry_base_delay * (
-                2**attempt
-            ) + random.uniform(0, 0.5)
-            delay = min(delay, self._config.retry_max_delay)
+            # #221 P2.1: match REST transport's AWS Full Jitter formula.
+            # The previous ``base * 2**attempt + uniform(0, 0.5)`` reduced
+            # to a flat 0.5s jitter window once the exponential pinned the
+            # delay at ``retry_max_delay`` — thundering-herd risk on a
+            # global Kalshi WS outage.
+            delay = random.uniform(
+                0,
+                min(
+                    self._config.retry_max_delay,
+                    self._config.retry_base_delay * (2**attempt),
+                ),
+            )
             logger.warning(
                 "Reconnecting in %.1fs (attempt %d/%d)",
                 delay,
@@ -172,14 +192,7 @@ async def reconnect(self) -> None:
             await asyncio.sleep(delay)
             try:
                 await self._set_state(ConnectionState.CONNECTING)
-                headers = await self._build_auth_headers()
-                self._ws = await connect(
-                    self._config.ws_base_url,
-                    additional_headers=headers,
-                    ping_interval=20,
-                    ping_timeout=self._heartbeat_timeout,
-                    close_timeout=5.0,
-                )
+                self._ws = await self._open_socket()
                 await self._set_state(ConnectionState.CONNECTED)
                 return
             except Exception:
@@ -212,7 +225,7 @@ async def send(self, msg: dict[str, Any]) -> None:
         """
         if self._ws is None:
             raise KalshiConnectionError("Not connected")
-        await self._ws.send(json.dumps(msg))
+        await self._ws.send(self._json_dumps(msg))
 
     async def recv(self) -> str:
         """Receive a raw string message from the WebSocket.

kalshi/ws/models/orderbook_delta.py

@@ -1,6 +1,8 @@
 """Orderbook delta and snapshot message models."""
 from __future__ import annotations
 
+from typing import Literal
+
 from pydantic import AliasChoices, BaseModel, Field
 
 from kalshi.types import DollarDecimal, FixedPointCount
@@ -45,7 +47,7 @@ class OrderbookDeltaPayload(BaseModel):
     delta: FixedPointCount = Field(
         validation_alias=AliasChoices("delta_fp", "delta"),
     )
-    side: str
+    side: Literal["yes", "no"]
     client_order_id: str | None = None
     subaccount: int | None = None
     ts: str | None = None

kalshi/ws/orderbook.py

@@ -90,10 +90,14 @@ def __init__(self) -> None:
         # number of tickers owned by the sid, not O(n) over every book.
         self._sid_tickers: dict[int, set[str]] = {}
 
-    def apply_snapshot(
+    def _apply_snapshot_inplace(
         self, msg: OrderbookSnapshotMessage, sid: int | None = None
-    ) -> Orderbook:
-        """Initialize (or reset) a book from a full snapshot.
+    ) -> None:
+        """In-place state mutation for a snapshot. No materialization.
+
+        Recv-loop hot path entry point (#199). The public
+        :meth:`apply_snapshot` wraps this with a ``to_orderbook()`` call
+        for direct callers that still want a fully built ``Orderbook``.
 
         ``sid`` records which subscription produced this snapshot so the
         ticker can be torn down later via :meth:`remove_by_sid` — required
@@ -131,23 +135,19 @@ def apply_snapshot(
             len(no_levels),
             sid_val,
         )
-        return state.to_orderbook()
-
-    def apply_delta(self, msg: OrderbookDeltaMessage) -> Orderbook | None:
-        """Apply an incremental delta to an existing book.
 
-        O(1) on the update itself (price-indexed dict lookup); the returned
-        Orderbook is materialized at O(n log n) only because consumers expect
-        a sorted level list.
+    def _apply_delta_inplace(self, msg: OrderbookDeltaMessage) -> bool:
+        """In-place state mutation for a delta. No materialization.
 
-        Returns the updated Orderbook, or None if no book exists for this ticker
-        (delta arrived before snapshot -- should not happen in normal flow).
+        Returns True if the mutation was applied, False if no book exists
+        for the ticker yet (delta arrived before snapshot — caller may
+        choose to skip materializing in that case).
         """
         ticker = msg.msg.market_ticker
         state = self._books.get(ticker)
         if state is None:
             logger.warning("Delta for unknown ticker %s (no snapshot yet)", ticker)
-            return None
+            return False
 
         price = msg.msg.price  # Decimal via DollarDecimal
         delta = msg.msg.delta  # Decimal via FixedPointCount
@@ -165,8 +165,35 @@ def apply_delta(self, msg: OrderbookDeltaMessage) -> Orderbook | None:
         elif delta > 0:
             levels[price] = delta
 
+        return True
+
+    def apply_snapshot(
+        self, msg: OrderbookSnapshotMessage, sid: int | None = None
+    ) -> Orderbook:
+        """Initialize (or reset) a book from a full snapshot.
+
+        Public wrapper that mutates in place then materializes the
+        resulting :class:`Orderbook`. The recv loop bypasses this and
+        calls :meth:`_apply_snapshot_inplace` directly to avoid the
+        unused O(n log n) materialization (#199).
+        """
+        self._apply_snapshot_inplace(msg, sid=sid)
+        ticker = msg.msg.market_ticker
+        state = self._books[ticker]
         return state.to_orderbook()
 
+    def apply_delta(self, msg: OrderbookDeltaMessage) -> Orderbook | None:
+        """Apply an incremental delta to an existing book.
+
+        Public wrapper. Returns the materialized Orderbook on success or
+        None if no book exists for the ticker. The recv loop bypasses this
+        and calls :meth:`_apply_delta_inplace` directly (#199).
+        """
+        if not self._apply_delta_inplace(msg):
+            return None
+        state = self._books.get(msg.msg.market_ticker)
+        return state.to_orderbook() if state else None
+
     def get(self, ticker: str) -> Orderbook | None:
         """Get current book state (non-blocking).
 

tests/test_config.py

@@ -228,3 +228,36 @@ def test_total_timeout_field_exists_and_defaults_None(self) -> None:  # noqa: N8
     def test_total_timeout_can_be_set(self) -> None:
         config = KalshiConfig(total_timeout=30.0)
         assert config.total_timeout == 30.0
+
+
+class TestWsExtraConfigFields:
+    """#208 + #209: ws_ping_interval / ws_close_timeout / ws_json_loads /
+    ws_json_dumps fields exist on KalshiConfig with sane defaults.
+    """
+
+    def test_ws_ping_interval_default_20(self) -> None:
+        assert KalshiConfig().ws_ping_interval == 20.0
+
+    def test_ws_close_timeout_default_5(self) -> None:
+        assert KalshiConfig().ws_close_timeout == 5.0
+
+    def test_ws_json_loads_default_None(self) -> None:  # noqa: N802
+        assert KalshiConfig().ws_json_loads is None
+
+    def test_ws_json_dumps_default_None(self) -> None:  # noqa: N802
+        assert KalshiConfig().ws_json_dumps is None
+
+    def test_ws_ping_interval_settable(self) -> None:
+        cfg = KalshiConfig(ws_ping_interval=45.0)
+        assert cfg.ws_ping_interval == 45.0
+
+    def test_ws_close_timeout_settable(self) -> None:
+        cfg = KalshiConfig(ws_close_timeout=1.5)
+        assert cfg.ws_close_timeout == 1.5
+
+    def test_ws_json_loads_callable_stored(self) -> None:
+        import json
+
+        cfg = KalshiConfig(ws_json_loads=json.loads, ws_json_dumps=json.dumps)
+        assert cfg.ws_json_loads is json.loads
+        assert cfg.ws_json_dumps is json.dumps