From 09f752af1327fd8033de0ec864e9d7f2101a5d8b Mon Sep 17 00:00:00 2001 From: Jeff West Date: Sun, 17 May 2026 11:19:43 -0500 Subject: [PATCH 1/2] fix(ws): return fresh Orderbook snapshots from OrderbookManager OrderbookManager previously kept the public Orderbook model as its mutable backing store and handed the same instance back to every caller. apply_delta and apply_snapshot mutated book.yes / book.no in place, so consumers that retained a reference to a yielded book silently saw future updates leak in -- breaking the snapshot semantics that Pydantic users (and the iterator wrapper in client.py) expect. Separate internal mutable state (_BookState) from the emitted model. Each apply_*/get call now materializes a fresh Orderbook with copied level lists. OrderbookLevel itself is treated as immutable (we replace rather than mutate), so the levels can be shared by reference without risk. Add a regression test that pins both the identity invariant (each call returns a distinct Orderbook) and the value invariant (a previously returned book's state survives later deltas). Closes #85 Co-Authored-By: Claude Opus 4.7 (1M context) --- kalshi/ws/orderbook.py | 55 +++++++++++++++++++++++++++++++------- tests/ws/test_orderbook.py | 27 +++++++++++++++++++ 2 files changed, 72 insertions(+), 10 deletions(-) diff --git a/kalshi/ws/orderbook.py b/kalshi/ws/orderbook.py index 969501d..88d1ce7 100644 --- a/kalshi/ws/orderbook.py +++ b/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,29 @@ logger = logging.getLogger("kalshi.ws") +@dataclass +class _BookState: + """Internal mutable state for one ticker's orderbook. + + 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: list[OrderbookLevel] = field(default_factory=list) + no: list[OrderbookLevel] = field(default_factory=list) + + def to_orderbook(self) -> Orderbook: + """Build a fresh ``Orderbook`` snapshot from the current state. + + Lists are copied so subsequent mutations of internal state do not + affect the returned model. ``OrderbookLevel`` instances themselves + are immutable in practice (we replace rather than mutate them) so + they can be shared by reference. + """ + return Orderbook(ticker=self.ticker, yes=list(self.yes), no=list(self.no)) + + class OrderbookManager: """Maintains local orderbook state from WebSocket stream. @@ -22,6 +46,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,7 +58,7 @@ 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.""" @@ -43,15 +71,15 @@ def apply_snapshot(self, msg: OrderbookSnapshotMessage) -> Orderbook: 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 + 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. @@ -60,8 +88,8 @@ def apply_delta(self, msg: OrderbookDeltaMessage) -> Orderbook | None: (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 @@ -69,7 +97,7 @@ def apply_delta(self, msg: OrderbookDeltaMessage) -> Orderbook | None: delta = msg.msg.delta # Decimal via FixedPointCount side = msg.msg.side - levels = book.yes if side == "yes" else book.no + levels = state.yes if side == "yes" else state.no # Find existing level at this price existing_idx = -1 @@ -93,11 +121,18 @@ def apply_delta(self, msg: OrderbookDeltaMessage) -> Orderbook | None: # Keep sorted by price levels.sort(key=lambda lv: lv.price) - 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).""" diff --git a/tests/ws/test_orderbook.py b/tests/ws/test_orderbook.py index 33cb23a..58880ee 100644 --- a/tests/ws/test_orderbook.py +++ b/tests/ws/test_orderbook.py @@ -159,6 +159,33 @@ 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_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 From 1426f910d2e9b7acd8f9dfac826c2bf4c7e600ff Mon Sep 17 00:00:00 2001 From: Jeff West Date: Sun, 17 May 2026 11:23:40 -0500 Subject: [PATCH 2/2] perf(ws): make OrderbookManager.apply_delta O(1) via price-indexed dict The previous list-of-levels storage required a linear scan to locate a price (O(n)) and a re-sort after every new-level insert (O(n log n)). On shallow books that's fine, but deep books -- which multivariate markets surface at thousands of levels per side -- turn every WS delta into an O(n) operation on the hot path. Switch _BookState's internal storage to dict[Decimal, Decimal] keyed on price. apply_delta becomes O(1) for both the lookup and the update (get / set / del). The sorted list[OrderbookLevel] that consumers see is materialized lazily in _BookState.to_orderbook -- still O(n log n), but only when a snapshot is actually emitted (i.e. exactly once per delta, no extra work). Add a regression test that swaps in a counting dict and asserts the per-delta access count stays bounded as the book grows from 10 to 500 levels (proves O(1), not O(n)). Invariant-based, not wall-clock, so it is stable on slow CI. Closes #87 Co-Authored-By: Claude Opus 4.7 (1M context) --- kalshi/ws/orderbook.py | 67 ++++++++++++++++++------------------- tests/ws/test_orderbook.py | 68 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 101 insertions(+), 34 deletions(-) diff --git a/kalshi/ws/orderbook.py b/kalshi/ws/orderbook.py index 88d1ce7..16b5638 100644 --- a/kalshi/ws/orderbook.py +++ b/kalshi/ws/orderbook.py @@ -18,23 +18,36 @@ 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: list[OrderbookLevel] = field(default_factory=list) - no: list[OrderbookLevel] = field(default_factory=list) + 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. - Lists are copied so subsequent mutations of internal state do not - affect the returned model. ``OrderbookLevel`` instances themselves - are immutable in practice (we replace rather than mutate them) so - they can be shared by reference. + 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. """ - return Orderbook(ticker=self.ticker, yes=list(self.yes), no=list(self.no)) + 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: @@ -63,14 +76,8 @@ def __init__(self) -> None: 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 - ] + 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( @@ -84,6 +91,10 @@ def apply_snapshot(self, msg: OrderbookSnapshotMessage) -> 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). """ @@ -98,28 +109,16 @@ def apply_delta(self, msg: OrderbookDeltaMessage) -> Orderbook | None: side = msg.msg.side levels = state.yes if side == "yes" else state.no + existing_qty = levels.get(price) - # Find existing level at this price - existing_idx = -1 - for i, level in enumerate(levels): - if level.price == price: - existing_idx = i - break - - 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 state.to_orderbook() diff --git a/tests/ws/test_orderbook.py b/tests/ws/test_orderbook.py index 58880ee..42e40ea 100644 --- a/tests/ws/test_orderbook.py +++ b/tests/ws/test_orderbook.py @@ -186,6 +186,74 @@ def test_returned_book_is_snapshot_not_live_view(self) -> None: 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