fix(ws): orderbook overhaul — fresh snapshots + O(1) apply_delta

kalshi/ws/orderbook.py

@@ -2,6 +2,7 @@
 from __future__ import annotations
 
 import logging
+from dataclasses import dataclass, field
 from decimal import Decimal
 
 from kalshi.models.markets import Orderbook, OrderbookLevel
@@ -13,6 +14,42 @@
 logger = logging.getLogger("kalshi.ws")
 
 
+@dataclass
+class _BookState:
+    """Internal mutable state for one ticker's orderbook.
+
+    Levels are stored price-indexed (``dict[Decimal, Decimal]``) so delta
+    application is O(1) regardless of book depth. The sorted
+    ``list[OrderbookLevel]`` exposed via :class:`Orderbook` is materialized
+    lazily in :meth:`to_orderbook` (O(n log n), but only when a snapshot is
+    actually emitted to a consumer).
+
+    Kept separate from the public :class:`Orderbook` model so the manager can
+    mutate freely without leaking changes to previously-handed-out snapshots.
+    """
+
+    ticker: str
+    yes: dict[Decimal, Decimal] = field(default_factory=dict)
+    no: dict[Decimal, Decimal] = field(default_factory=dict)
+
+    def to_orderbook(self) -> Orderbook:
+        """Build a fresh ``Orderbook`` snapshot from the current state.
+
+        Levels are emitted price-ascending to match the historical wire-order
+        contract that the previous list-backed implementation maintained via
+        ``list.sort`` after every insert.
+        """
+        yes_levels = [
+            OrderbookLevel(price=price, quantity=qty)
+            for price, qty in sorted(self.yes.items())
+        ]
+        no_levels = [
+            OrderbookLevel(price=price, quantity=qty)
+            for price, qty in sorted(self.no.items())
+        ]
+        return Orderbook(ticker=self.ticker, yes=yes_levels, no=no_levels)
+
+
 class OrderbookManager:
     """Maintains local orderbook state from WebSocket stream.
 
@@ -22,6 +59,10 @@ class OrderbookManager:
     for quantities; both parse directly into ``Decimal`` without any
     cents-to-dollars conversion.
 
+    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.
+
     Usage:
         mgr = OrderbookManager()
         book = mgr.apply_snapshot(snapshot_msg)  # Initialize
@@ -30,74 +71,67 @@ class OrderbookManager:
     """
 
     def __init__(self) -> None:
-        self._books: dict[str, Orderbook] = {}
+        self._books: dict[str, _BookState] = {}
 
     def apply_snapshot(self, msg: OrderbookSnapshotMessage) -> Orderbook:
         """Initialize (or reset) a book from a full snapshot."""
         ticker = msg.msg.market_ticker
-        yes_levels = [
-            OrderbookLevel(price=Decimal(p), quantity=Decimal(q))
-            for p, q in msg.msg.yes
-        ]
-        no_levels = [
-            OrderbookLevel(price=Decimal(p), quantity=Decimal(q))
-            for p, q in msg.msg.no
-        ]
-        book = Orderbook(ticker=ticker, yes=yes_levels, no=no_levels)
-        self._books[ticker] = book
+        yes_levels = {Decimal(p): Decimal(q) for p, q in msg.msg.yes}
+        no_levels = {Decimal(p): Decimal(q) for p, q in msg.msg.no}
+        state = _BookState(ticker=ticker, yes=yes_levels, no=no_levels)
+        self._books[ticker] = state
         logger.debug(
             "Orderbook snapshot: %s (%d yes, %d no levels)",
             ticker,
             len(yes_levels),
             len(no_levels),
         )
-        return book
+        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.
+
         Returns the updated Orderbook, or None if no book exists for this ticker
         (delta arrived before snapshot -- should not happen in normal flow).
         """
         ticker = msg.msg.market_ticker
-        book = self._books.get(ticker)
-        if book is None:
+        state = self._books.get(ticker)
+        if state is None:
             logger.warning("Delta for unknown ticker %s (no snapshot yet)", ticker)
             return None
 
         price = msg.msg.price  # Decimal via DollarDecimal
         delta = msg.msg.delta  # Decimal via FixedPointCount
         side = msg.msg.side
 
-        levels = book.yes if side == "yes" else book.no
-
-        # Find existing level at this price
-        existing_idx = -1
-        for i, level in enumerate(levels):
-            if level.price == price:
-                existing_idx = i
-                break
+        levels = state.yes if side == "yes" else state.no
+        existing_qty = levels.get(price)
 
-        if existing_idx >= 0:
-            existing = levels[existing_idx]
-            new_qty = existing.quantity + delta
+        if existing_qty is not None:
+            new_qty = existing_qty + delta
             if new_qty <= 0:
-                # Remove the level
-                levels.pop(existing_idx)
+                del levels[price]
             else:
-                levels[existing_idx] = OrderbookLevel(price=price, quantity=new_qty)
-        else:
-            if delta > 0:
-                # Add new level
-                levels.append(OrderbookLevel(price=price, quantity=delta))
-                # Keep sorted by price
-                levels.sort(key=lambda lv: lv.price)
+                levels[price] = new_qty
+        elif delta > 0:
+            levels[price] = delta
 
-        return book
+        return state.to_orderbook()
 
     def get(self, ticker: str) -> Orderbook | None:
-        """Get current book state (non-blocking)."""
-        return self._books.get(ticker)
+        """Get current book state (non-blocking).
+
+        Returns a fresh :class:`Orderbook` snapshot; the caller is free to
+        retain it without seeing future mutations leak in.
+        """
+        state = self._books.get(ticker)
+        if state is None:
+            return None
+        return state.to_orderbook()
 
     def remove(self, ticker: str) -> None:
         """Remove a book (e.g., on unsubscribe)."""

tests/ws/test_orderbook.py

@@ -159,6 +159,101 @@ def test_negative_delta_partial(self) -> None:
         assert book is not None
         assert book.yes[0].quantity == Decimal("70")  # 100 - 30
 
+    def test_returned_book_is_snapshot_not_live_view(self) -> None:
+        """Regression for #85: a book handed to a consumer must not mutate
+        when subsequent deltas are applied. Pin both the identity invariant
+        (each call returns a distinct ``Orderbook``) and the value invariant
+        (the previously-returned book's state is preserved verbatim).
+        """
+        mgr = OrderbookManager()
+        snap_book = mgr.apply_snapshot(make_snapshot(yes=[["0.50", "100"]]))
+
+        # Apply a delta that would have mutated the level in-place under the
+        # old implementation (quantity 100 -> 150).
+        delta_book = mgr.apply_delta(make_delta(price="0.50", delta="50", side="yes"))
+        assert delta_book is not None
+
+        # Distinct instances.
+        assert snap_book is not delta_book
+        # The first-handed-out book is frozen at its emission-time state.
+        assert snap_book.yes[0].quantity == Decimal("100")
+        assert delta_book.yes[0].quantity == Decimal("150")
+
+        # And ``get()`` is also a snapshot, not a live view.
+        get_book = mgr.get("T")
+        assert get_book is not None
+        assert get_book is not delta_book
+        mgr.apply_delta(make_delta(price="0.50", delta="25", side="yes"))
+        assert get_book.yes[0].quantity == Decimal("150")  # unchanged
+
+    def test_apply_delta_is_constant_time_in_book_depth(self) -> None:
+        """Regression for #87: apply_delta must touch at most a constant
+        number of levels per delta, not scan the whole book.
+
+        We swap the internal price->qty mapping for a counting dict that
+        records every key access (``get`` / ``__getitem__`` /
+        ``__setitem__`` / ``__delitem__`` / ``__contains__``). The old
+        list-scan implementation iterated the level container once per
+        level (O(n) accesses); the new dict-backed implementation does
+        O(1) accesses regardless of book depth.
+
+        This is an *invariant* test (access count grows like O(1), not
+        O(n)), not a wall-clock budget -- stable on slow CI.
+        """
+        from kalshi.ws.orderbook import _BookState
+
+        class CountingDict(dict[Decimal, Decimal]):
+            def __init__(self, *a: object, **kw: object) -> None:
+                super().__init__(*a, **kw)  # type: ignore[arg-type]
+                self.accesses = 0
+
+            def get(self, key: object, default: object = None) -> object:  # type: ignore[override]
+                self.accesses += 1
+                return super().get(key, default)  # type: ignore[arg-type]
+
+            def __getitem__(self, key: Decimal) -> Decimal:
+                self.accesses += 1
+                return super().__getitem__(key)
+
+            def __setitem__(self, key: Decimal, value: Decimal) -> None:
+                self.accesses += 1
+                super().__setitem__(key, value)
+
+            def __delitem__(self, key: Decimal) -> None:
+                self.accesses += 1
+                super().__delitem__(key)
+
+            def __contains__(self, key: object) -> bool:
+                self.accesses += 1
+                return super().__contains__(key)
+
+        # Two probe sizes -- a small book and a large book. If the
+        # implementation is O(n), the second number is ~50x the first.
+        # If O(1), they're identical.
+        def measure(n: int) -> int:
+            mgr = OrderbookManager()
+            mgr.apply_snapshot(
+                make_snapshot(yes=[[f"0.{i:04d}", "100"] for i in range(1, n + 1)])
+            )
+            counting = CountingDict(mgr._books["T"].yes)
+            mgr._books["T"] = _BookState(
+                ticker="T", yes=counting, no=mgr._books["T"].no
+            )
+            mgr.apply_delta(
+                make_delta(price=f"0.{n // 2:04d}", delta="1", side="yes")
+            )
+            return counting.accesses
+
+        small = measure(10)
+        large = measure(500)
+
+        # Allow generous slack but reject anything that scales with n.
+        # O(1) -> equal counts; O(n) -> large is ~50x small.
+        assert large <= small + 2, (
+            f"apply_delta access count grew from {small} (n=10) to "
+            f"{large} (n=500) -- expected O(1), looks O(n)"
+        )
+
     def test_fractional_delta_matches_wire_format(self) -> None:
         """Live capture on demo shows delta_fp arrives as ``"1.00"`` (with
         trailing .00), not bare ``"1"``. Lock in that Decimal arithmetic