perf(ws): adopt orderbook snapshot dicts by identity (#296)

kalshi/ws/client.py

@@ -643,6 +643,13 @@ async def subscribe_ticker(
     async def subscribe_orderbook_delta(
         self, *, tickers: list[str] | None = None, maxsize: int = 1000,
     ) -> AsyncIterator[OrderbookSnapshotMessage | OrderbookDeltaMessage]:
+        """Subscribe to ``orderbook_delta`` for the given tickers.
+
+        Note: ``OrderbookSnapshotMessage.msg.yes`` and ``.no`` are live dicts
+        owned by the :class:`OrderbookManager` after this dispatch — they
+        mutate on every delta. Use :meth:`orderbook` if you need an immutable
+        snapshot view.
+        """
         params: dict[str, Any] = {"send_initial_snapshot": True}
         if tickers:
             params["market_tickers"] = tickers

kalshi/ws/models/orderbook_delta.py

@@ -69,6 +69,11 @@ class OrderbookSnapshotPayload(BaseModel):
     ``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.
+
+    Aliasing note: when delivered through
+    :meth:`KalshiWSClient.subscribe_orderbook_delta`, ``yes`` / ``no`` are live
+    dicts owned by the :class:`OrderbookManager` after dispatch — they
+    mutate on every subsequent delta.
     """
 
     market_ticker: str

kalshi/ws/orderbook.py

@@ -89,6 +89,10 @@ class OrderbookManager:
     Each call to :meth:`apply_snapshot`, :meth:`apply_delta`, or :meth:`get`
     returns a fresh :class:`Orderbook` instance. Consumers may safely retain
     references; subsequent updates will not mutate previously-returned books.
+    Input ``OrderbookSnapshotMessage`` arguments are also safe to reuse on
+    the public API — see :meth:`apply_snapshot`. Raw
+    ``subscribe_orderbook_delta`` consumers, which bypass this wrapper, see
+    the live ``_BookState`` dicts; see :class:`OrderbookSnapshotPayload`.
 
     Per-sid tracking (:meth:`tickers_for_sid`, :meth:`remove_by_sid`) lets
     callers tear down all books owned by a subscription without enumerating
@@ -122,11 +126,14 @@ def _apply_snapshot_inplace(
         ticker can be torn down later via :meth:`remove_by_sid` — required
         for all-markets subscriptions that don't pin tickers up front.
         If omitted, defaults to ``msg.sid`` from the envelope.
+
+        Ownership note: ``msg.msg.yes`` / ``.no`` are adopted by identity
+        into the new ``_BookState``; callers must treat them as consumed.
         """
         ticker = msg.msg.market_ticker
         sid_val = sid if sid is not None else msg.sid
-        yes_levels = dict(msg.msg.yes)
-        no_levels = dict(msg.msg.no)
+        yes_levels = msg.msg.yes
+        no_levels = msg.msg.no
 
         # If a prior book existed under a different sid (e.g. server moved
         # the ticker between subs), unindex it from the old sid first so
@@ -193,14 +200,27 @@ def _apply_delta_inplace(self, msg: OrderbookDeltaMessage) -> bool:
     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
+        Public wrapper: mutates in place via :meth:`_apply_snapshot_inplace`
+        and materializes a fresh :class:`Orderbook`. The recv loop bypasses
+        this and calls :meth:`_apply_snapshot_inplace` directly to avoid the
         unused O(n log n) materialization (#199).
+
+        ``msg`` is safe to reuse after this call: the manager defensively
+        copies the adopted dicts into ``_BookState`` so subsequent delta
+        mutations do not leak back through ``msg.msg.yes`` / ``.no``. The
+        recv loop, which never holds ``msg`` past dispatch, skips this
+        defensive copy via the direct inplace path.
         """
         self._apply_snapshot_inplace(msg, sid=sid)
         ticker = msg.msg.market_ticker
         state = self._books[ticker]
+        # Defensive copy on the public path so the caller-supplied ``msg``
+        # keeps its original ``.yes`` / ``.no`` dicts; ``_apply_snapshot_inplace``
+        # adopted them by identity above, so we copy via ``msg.msg`` to make
+        # the "copy the caller's input" intent explicit. The recv-loop bypass
+        # keeps the identity-adopt perf win where it matters.
+        state.yes = dict(msg.msg.yes)
+        state.no = dict(msg.msg.no)
         return state.to_orderbook()
 
     def apply_delta(self, msg: OrderbookDeltaMessage) -> Orderbook | None:

tests/ws/test_orderbook.py

@@ -362,3 +362,146 @@ def test_public_apply_snapshot_returns_orderbook(self) -> None:
         assert book.ticker == "T"
         assert len(book.yes) == 1
         assert len(book.no) == 1
+
+class TestSnapshotIdentityAdoption:
+    """Regression for #296: ``_apply_snapshot_inplace`` must adopt the
+    validator-built ``dict[Decimal, Decimal]`` by identity (no rebuild).
+
+    The CHANGELOG v2.5.0 #263 claim is "snapshot apply collapses to a single
+    dict walk"; if a future refactor reintroduces a ``dict(...)`` copy, this
+    test fails immediately rather than silently regressing recv-loop perf.
+    """
+
+    def test_snapshot_adopts_validated_dicts_by_identity(self) -> None:
+        mgr = OrderbookManager()
+        # Build with Decimal-keyed dicts directly so the validator's
+        # identity fast-path engages (see ``_levels_to_dict``).
+        yes_in: dict[Decimal, Decimal] = {Decimal("0.50"): Decimal("100.00")}
+        no_in: dict[Decimal, Decimal] = {Decimal("0.45"): Decimal("150.00")}
+        msg = OrderbookSnapshotMessage.model_validate(
+            {
+                "type": "orderbook_snapshot",
+                "sid": 1,
+                "seq": 1,
+                "msg": {
+                    "market_ticker": "T",
+                    "market_id": "id",
+                    "yes": yes_in,
+                    "no": no_in,
+                },
+            }
+        )
+        # We don't assert state.yes is yes_in (Pydantic may copy the input during
+        # field binding), but we DO assert state.yes is msg.msg.yes: the SDK must
+        # adopt the already-validated dict without adding another dict() copy.
+        mgr._apply_snapshot_inplace(msg)
+        state = mgr._books["T"]
+        assert state.yes is msg.msg.yes
+        assert state.no is msg.msg.no
+
+    def test_delta_mutates_in_place_and_invalidates_cache(self) -> None:
+        mgr = OrderbookManager()
+        yes_in: dict[Decimal, Decimal] = {Decimal("0.50"): Decimal("100.00")}
+        msg = OrderbookSnapshotMessage.model_validate(
+            {
+                "type": "orderbook_snapshot",
+                "sid": 1,
+                "seq": 1,
+                "msg": {
+                    "market_ticker": "T",
+                    "market_id": "id",
+                    "yes": yes_in,
+                    "no": {},
+                },
+            }
+        )
+        mgr._apply_snapshot_inplace(msg)
+        state = mgr._books["T"]
+        adopted_yes = state.yes
+        # Materialize once to populate the #244 cache.
+        cached = state.to_orderbook()
+        assert state._cached is cached
+
+        applied = mgr._apply_delta_inplace(
+            make_delta(price="0.50", delta="25", side="yes")
+        )
+        assert applied is True
+        # In-place mutation: the dict instance is unchanged, the value moved.
+        assert state.yes is adopted_yes
+        assert state.yes[Decimal("0.50")] == Decimal("125.00")
+        # The cache was invalidated so the next read re-materializes.
+        assert state._cached is None
+
+    def test_subscriber_held_snapshot_dict_mutates_on_delta(self) -> None:
+        """Documents the live-dict contract for raw ``subscribe_orderbook_delta`` consumers."""
+        mgr = OrderbookManager()
+        yes_in: dict[Decimal, Decimal] = {Decimal("0.50"): Decimal("100.00")}
+        no_in: dict[Decimal, Decimal] = {Decimal("0.45"): Decimal("150.00")}
+        msg = OrderbookSnapshotMessage.model_validate(
+            {
+                "type": "orderbook_snapshot",
+                "sid": 1,
+                "seq": 1,
+                "msg": {
+                    "market_ticker": "T",
+                    "market_id": "id",
+                    "yes": yes_in,
+                    "no": no_in,
+                },
+            }
+        )
+        mgr._apply_snapshot_inplace(msg)
+        # A raw subscriber who stored ``msg.msg.yes`` after receiving the snapshot
+        # frame is holding the live ``_BookState`` dict; the next delta mutates it.
+        held_yes = msg.msg.yes
+        applied = mgr._apply_delta_inplace(
+            make_delta(price="0.60", delta="25", side="yes")
+        )
+        assert applied is True
+        assert Decimal("0.60") in held_yes
+        assert held_yes[Decimal("0.60")] == Decimal("25.00")
+        assert held_yes is mgr._books["T"].yes
+
+    def test_public_apply_snapshot_does_not_alias_input_msg(self) -> None:
+        """Public apply_snapshot must defensively copy so the caller's msg is not aliased."""
+        snapshot_payload = {
+            "type": "orderbook_snapshot",
+            "sid": 7,
+            "seq": 1,
+            "msg": {
+                "market_ticker": "PUB-T",
+                "market_id": "id",
+                "yes": {Decimal("0.30"): Decimal("5")},
+                "no": {Decimal("0.70"): Decimal("3")},
+            },
+        }
+        msg = OrderbookSnapshotMessage.model_validate(snapshot_payload)
+        held_yes = msg.msg.yes
+
+        mgr = OrderbookManager()
+        mgr.apply_snapshot(msg)
+
+        # Public path: state must NOT share identity with the caller's msg.
+        state = mgr._books["PUB-T"]
+        assert state.yes is not held_yes
+        assert state.no is not msg.msg.no
+        assert state.yes == {Decimal("0.30"): Decimal("5")}
+
+        # A subsequent delta mutates state but leaves the caller's msg dict untouched.
+        delta = OrderbookDeltaMessage.model_validate({
+            "type": "orderbook_delta",
+            "sid": 7,
+            "seq": 2,
+            "msg": {
+                "market_ticker": "PUB-T",
+                "market_id": "id",
+                "price": Decimal("0.30"),
+                "delta": Decimal("4"),
+                "side": "yes",
+            },
+        })
+        mgr.apply_delta(delta)
+        assert state.yes[Decimal("0.30")] == Decimal("9")
+        # Caller's held reference is untouched — the public-API safety contract.
+        assert held_yes[Decimal("0.30")] == Decimal("5")
+