From 3c81d40c9fd04d81990a1604a073dba3c6868b72 Mon Sep 17 00:00:00 2001 From: susheem-k Date: Sat, 12 Sep 2026 16:41:45 +0530 Subject: [PATCH] feat: wire Ledger onto LedgerBackend (precheck/apply_events) alongside Store (#118) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds a backend= mode to Ledger, additive alongside the existing store=/in-memory modes (mutually exclusive with store=) — the core Ledger/LedgerView piece of the remote-only rewrite's Phase 2 checklist. build_governor/ControlPlaneClient don't construct it yet; that swap, and migrating the ~18 Store-dependent tests, are the next increments (staged deliberately, per the epic's own Phase 2/3 split). - record(): batches the step + spent_add (when cost > 0) into one apply_events call sharing an idempotency seq, so the ack's totals cover the run-total cum spend without a separate read_state round trip for the common case. - admit/complete/mark_halted/clear_halt: one-shot apply_events writes. Their idempotency keys are random (uuid4), not the contract's deterministic {run_id}:...:{seq} recipe — that recipe needs a real run_id to stay collision-free across processes for a non-run-scoped segment key (e.g. an agent-dimension concurrency cap shared across runs), which isn't always available at these call sites. A random key is correct for a one-shot, non-retried write either way; the deterministic recipe only earns its keep once a buffered backend needs a retried flush to regenerate the same key (the existing # TODO(buffering) seam in ledger_backend.py). - cost_micros/budget_left/is_halted/inflight: read via backend.read_state (precheck). No per-call batching across detectors yet (one read per method call) — Governor grouping detectors by data_scope into a single precheck per governed moment is a separate, later optimization, not required for correctness given there's no client-side write buffer yet either (decision #6). - velocity/recent/window/step_count are unchanged: always Tier-1 LocalRunState, never a backend round trip — those are inherently per-process reads. tests/test_ledger_backend_mode.py: 9 new tests against FakeLedgerBackend, including two Ledger instances sharing one backend (the actual cross-process scenario this whole rewire is for) for halt visibility and inflight counting. Full suite: 246 passed, 16 skipped, 11 deselected. ruff check/format and mypy clean. Co-Authored-By: Claude Sonnet 5 --- CHANGELOG.md | 10 ++ src/tokenops/control/ledger.py | 230 +++++++++++++++++++++++++++--- tests/test_ledger_backend_mode.py | 163 +++++++++++++++++++++ 3 files changed, 381 insertions(+), 22 deletions(-) create mode 100644 tests/test_ledger_backend_mode.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 4c39282..439e1c9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -29,6 +29,16 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 migration) and round-tripped through `governance_config_for`; not yet consumed by `build_governor` — the Governor doesn't group detectors by scope until the `LedgerBackend` rewire lands. +- `Ledger(backend=...)` — routes every write through `LedgerBackend.apply_events` and + every spend/inflight/halt read through `read_state` (`precheck`), alongside the + existing `store=`/in-memory modes (mutually exclusive with `store`) (#118). + `step`/`spent_add` batch together per crossing so the ack's `totals` cover the run + total in one round trip; `admit`/`complete`/`halt_mark`/`halt_clear` are separate + one-shot writes. `velocity`/`recent`/`window`/`step_count` stay Tier-1-only (never a + backend round trip) — those are inherently local, per-process reads. Not yet wired + into `build_governor`/`ControlPlaneClient`; `tests/test_ledger_backend_mode.py` + exercises it directly against `FakeLedgerBackend`, including two `Ledger` instances + sharing one backend (the cross-process case). ### Changed diff --git a/src/tokenops/control/ledger.py b/src/tokenops/control/ledger.py index 883b053..6a65f55 100644 --- a/src/tokenops/control/ledger.py +++ b/src/tokenops/control/ledger.py @@ -44,6 +44,7 @@ from __future__ import annotations import threading +import uuid from collections import defaultdict from collections.abc import Callable, Sequence from dataclasses import dataclass, field @@ -55,6 +56,7 @@ Micros, Observation, ) +from tokenops.control.ledger_backend import LedgerBackend, LedgerEvent, PrecheckRequest if TYPE_CHECKING: from tokenops.control.store import Store @@ -135,6 +137,108 @@ def segment_key(attr: Attribution, budget: Budget) -> str | None: return segment_key_for(attr, budget.dimension, budget.tag_key) +def _run_id_from_segment_key(segment_key: str) -> str | None: + """Recover ``run_id`` from a ``run:``-dimension segment key. ``None`` for any other + dimension (user/agent/tenant/tag) — those don't carry a run_id at all. + + The plane's ``precheck`` only *uses* ``run_id`` to answer ``halt``/``window`` — a + ``spent`` or ``inflight`` lookup for a non-run segment is valid with an empty + ``run_id`` (see ``control_plane.store.Store.precheck``), so callers pass ``"" `` when + this returns ``None`` rather than needing the caller's own run context threaded in. + """ + prefix = "run:" + return segment_key[len(prefix) :] if segment_key.startswith(prefix) else None + + +# --------------------------------------------------------------------------- # +# Backend-mode event builders (see ``ledger_backend.LedgerEvent`` / the control-plane +# ``docs/api-contract.md`` §5-6 for the wire shape and idempotency-key recipe) +# --------------------------------------------------------------------------- # + + +def _new_idempotency_key() -> str: + """A fresh, globally-unique key for a one-shot (non-buffered, non-retried) write. + + ``admit``/``complete``/``halt_mark``/``halt_clear`` can be issued for a segment key + that carries no run_id (e.g. an ``agent``-dimension concurrency cap shared across + many runs and processes) — the contract's deterministic ``{run_id}:...:{seq}`` + recipe (``docs/api-contract.md`` §6) needs a real, globally-known run_id to stay + collision-free across processes, which isn't always available here. A random key is + always correct for a single fire-and-forget write; the deterministic recipe only + earns its keep once a buffered backend needs a retried flush to regenerate the same + key (tracked with the ``# TODO(buffering)`` seam in ``ledger_backend.py``). + """ + return uuid.uuid4().hex + + +def _inflight_event( + kind: Literal["admit", "complete"], run_id: str, segment_key: str +) -> LedgerEvent: + return { + "kind": kind, + "idempotency_key": _new_idempotency_key(), + "run_id": run_id, + "segment_key": segment_key, + } + + +def _halt_event( + kind: Literal["halt_mark", "halt_clear"], + run_id: str, + reason: str = "", + detector: str = "", +) -> LedgerEvent: + event: LedgerEvent = { + "kind": kind, + "idempotency_key": _new_idempotency_key(), + "run_id": run_id, + } + if reason: + event["reason"] = reason + if detector: + event["detector"] = detector + return event + + +def _step_event(obs: Observation, seq: int, cost: Micros) -> LedgerEvent: + event: LedgerEvent = { + "kind": "step", + "idempotency_key": f"{obs.attr.run_id}:{obs.attr.agent}:{seq}:step", + "ts": obs.ts, + "run_id": obs.attr.run_id, + "agent": obs.attr.agent, + "seq": seq, + "node_type": obs.node_type, + "boundary_id": obs.boundary_id, + "cost_micros": cost, + "tags": {**dict(obs.boundary_tags), **dict(obs.tags)}, + } + if obs.usage is not None: + event["usage"] = { + "input": obs.usage.input, + "output": obs.usage.output, + "cached": obs.usage.cached, + "reasoning": obs.usage.reasoning, + } + if obs.signature is not None: + event["tool_signature"] = obs.signature + if obs.result_hash is not None: + event["result_hash"] = obs.result_hash + return event + + +def _spent_add_event( + run_id: str, agent: str, seq: int, delta: Micros, targets: list[dict[str, str]] +) -> LedgerEvent: + return { + "kind": "spent_add", + "idempotency_key": f"{run_id}:{agent}:{seq}:spent_add", + "run_id": run_id, + "delta_micros": delta, + "targets": targets, + } + + # =========================================================================== # # Per-run ephemeral state # # =========================================================================== # @@ -154,6 +258,9 @@ class LocalRunState: halted: bool = False halt_reason: str | None = None parent_run: str | None = None + #: Last-known run-total cum_spent from a backend ack, so a zero-cost crossing's + #: BoundaryStep doesn't need its own read_state round trip. + cum_spent_cache: Micros = 0 # =========================================================================== # @@ -163,7 +270,17 @@ class LocalRunState: class Ledger: """Attribute + LedgerView. Per-process run state (window, local step count); spend, - inflight, and halt may be backed by :class:`Store` for cross-process A2A. + inflight, and halt may be backed by :class:`Store` (local SQLite) or a + :class:`~tokenops.control.ledger_backend.LedgerBackend` (the remote control plane) for + cross-process / cross-agent consistency. At most one of ``store``/``backend`` may be + given; neither means fully in-memory (single-process only, tests). + + ``backend`` mode is the target of the remote-only rewrite (tokenops#118): every write + goes through ``apply_events`` (one batch per crossing) and every spend/inflight/halt + read goes through ``read_state`` (``precheck``). There is no per-call batching across + detectors yet — each read is its own round trip — so a call with several ``global``- + scope policies costs more than one ``precheck``; that optimization (and buffering the + write side) are tracked follow-ups, not required for correctness. Thread-safe for concurrent use from multiple threads (see module docstring). """ @@ -174,9 +291,13 @@ def __init__( budgets: Sequence[Budget] = (), price: PriceFn | None = None, store: Store | None = None, + backend: LedgerBackend | None = None, ) -> None: + if store is not None and backend is not None: + raise ValueError("Ledger takes at most one of store= / backend=, not both") self._lock = threading.RLock() self._store = store + self._backend = backend self.runs: dict[str, LocalRunState] = {} self._spent: dict[tuple[str, str, str], Micros] = defaultdict(int) self._inflight: dict[str, int] = defaultdict(int) @@ -188,6 +309,17 @@ def _spent_key(self, budget_id: str, segment_key: str, period: str) -> tuple[str return (budget_id, segment_key, period) def _read_spent(self, budget_id: str, segment_key: str, period: str) -> Micros: + if self._backend is not None: + state = self._backend.read_state( + PrecheckRequest( + run_id=_run_id_from_segment_key(segment_key) or "", + budgets=[ + {"budget_id": budget_id, "segment_key": segment_key, "period": period} + ], + want=["spent"], + ) + ) + return state.spent.get(f"{budget_id}|{segment_key}|{period}", 0) if self._store is not None: return self._store.ledger_get_spent(budget_id, segment_key, period) return self._spent[self._spent_key(budget_id, segment_key, period)] @@ -199,6 +331,9 @@ def _write_spent_delta( period: str, delta: Micros, ) -> Micros: + # Backend mode batches spent_add into record()'s single apply_events call + # (it needs to share an idempotency key/seq with that crossing's step event), + # so it never reaches this in-memory/Store fallback path. if self._store is not None: return self._store.ledger_add_spent(budget_id, segment_key, period, delta) key = self._spent_key(budget_id, segment_key, period) @@ -224,7 +359,10 @@ def close_run(self, run_id: str) -> None: def admit(self, segment_key: str) -> None: """A call for this segment has started (concurrency).""" with self._lock: - if self._store is not None: + if self._backend is not None: + run_id = _run_id_from_segment_key(segment_key) or "" + self._backend.apply_events([_inflight_event("admit", run_id, segment_key)]) + elif self._store is not None: self._store.ledger_admit(segment_key) else: self._inflight[segment_key] += 1 @@ -233,7 +371,10 @@ def complete(self, segment_key: str) -> None: """A call for this segment has returned. Floored at 0 so a stray complete cannot drive the counter negative.""" with self._lock: - if self._store is not None: + if self._backend is not None: + run_id = _run_id_from_segment_key(segment_key) or "" + self._backend.apply_events([_inflight_event("complete", run_id, segment_key)]) + elif self._store is not None: self._store.ledger_complete(segment_key) else: self._inflight[segment_key] = max(0, self._inflight[segment_key] - 1) @@ -257,23 +398,48 @@ def record(self, obs: Observation) -> BoundaryStep: elif obs.node_type == "delegate": cost = obs.rolled_up_cost_micros # child run total rolls up into the parent - # One event can feed many accumulators (incl. the system run-total) — that is - # why spend lives in the map, not on the run object. A zero-cost crossing - # (tool call, delegate with no rollup) must not touch the cost ledger — it is - # a *step*, not spend. - if cost > 0: - for b in self._budgets: - sk = segment_key(obs.attr, b) - if sk is None: - continue - self._write_spent_delta(b.budget_id, sk, b.period, cost) - rs.steps += 1 - cum = self._read_spent( - RUN_TOTAL_BUDGET.budget_id, - f"run:{obs.attr.run_id}", - LIFETIME, - ) + + if self._backend is not None: + # One apply_events batch for both the step and the spend, so they share + # an idempotency seq and the ack's totals cover the run-total in one + # round trip (no separate read_state needed for the common case). + events: list[LedgerEvent] = [_step_event(obs, rs.steps, cost)] + if cost > 0: + targets = [] + for b in self._budgets: + sk = segment_key(obs.attr, b) + if sk is None: + continue + targets.append( + {"budget_id": b.budget_id, "segment_key": sk, "period": b.period} + ) + events.append( + _spent_add_event(obs.attr.run_id, obs.attr.agent, rs.steps, cost, targets) + ) + result = self._backend.apply_events(events) + if result.halted: + rs.halted = True + run_total_key = f"{RUN_TOTAL_BUDGET.budget_id}|run:{obs.attr.run_id}|{LIFETIME}" + if run_total_key in result.totals: + rs.cum_spent_cache = result.totals[run_total_key] + cum = rs.cum_spent_cache + else: + # One event can feed many accumulators (incl. the system run-total) — + # that is why spend lives in the map, not on the run object. A zero-cost + # crossing (tool call, delegate with no rollup) must not touch the cost + # ledger — it is a *step*, not spend. + if cost > 0: + for b in self._budgets: + sk = segment_key(obs.attr, b) + if sk is None: + continue + self._write_spent_delta(b.budget_id, sk, b.period, cost) + cum = self._read_spent( + RUN_TOTAL_BUDGET.budget_id, + f"run:{obs.attr.run_id}", + LIFETIME, + ) step = BoundaryStep( step=rs.steps, ts=obs.ts, @@ -297,7 +463,9 @@ def mark_halted(self, run_id: str, reason: str = "") -> None: rs = self.runs[run_id] = LocalRunState() rs.halted = True rs.halt_reason = reason or rs.halt_reason - if self._store is not None: + if self._backend is not None: + self._backend.apply_events([_halt_event("halt_mark", run_id, reason=reason)]) + elif self._store is not None: self._store.ledger_mark_halted(run_id, reason) def clear_halt(self, run_id: str) -> None: @@ -308,7 +476,9 @@ def clear_halt(self, run_id: str) -> None: if rs is not None: rs.halted = False rs.halt_reason = None - if self._store is not None: + if self._backend is not None: + self._backend.apply_events([_halt_event("halt_clear", run_id)]) + elif self._store is not None: self._store.ledger_clear_halt(run_id) # ---- read side (LedgerView) ------------------------------------------ # @@ -324,9 +494,19 @@ def step_count(self, run_id: str) -> int: def is_halted(self, run_id: str) -> bool: with self._lock: + rs = self.runs.get(run_id) + if rs and rs.halted: + return True # fast path: already known locally, no round trip + if self._backend is not None: + state = self._backend.read_state(PrecheckRequest(run_id=run_id, want=["halt"])) + if state.halted: + if rs is None: + rs = self.runs[run_id] = LocalRunState() + rs.halted = True + rs.halt_reason = state.halt_reason + return state.halted if self._store is not None and self._store.ledger_is_halted(run_id): return True - rs = self.runs.get(run_id) return bool(rs and rs.halted) def budget_left(self, budget_id: str, segment_key: str, period: str = "lifetime") -> Micros: @@ -341,6 +521,12 @@ def budget_left(self, budget_id: str, segment_key: str, period: str = "lifetime" def inflight(self, segment_key: str) -> int: with self._lock: + if self._backend is not None: + run_id = _run_id_from_segment_key(segment_key) or "" + state = self._backend.read_state( + PrecheckRequest(run_id=run_id, segment_keys=[segment_key], want=["inflight"]) + ) + return state.inflight.get(segment_key, 0) if self._store is not None: return self._store.ledger_inflight(segment_key) return self._inflight[segment_key] diff --git a/tests/test_ledger_backend_mode.py b/tests/test_ledger_backend_mode.py new file mode 100644 index 0000000..338c591 --- /dev/null +++ b/tests/test_ledger_backend_mode.py @@ -0,0 +1,163 @@ +"""Ledger(backend=...) — the LedgerBackend-routed write/read path (tokenops#118). + +Mirrors test_ledger.py's assertions but backed by FakeLedgerBackend instead of the +in-memory / Store path, proving record/admit/complete/mark_halted/clear_halt and every +LedgerView read produce the same answers when routed through apply_events/read_state. +""" + +from __future__ import annotations + +import pytest + +from conftest import make_attr, toy_price +from fakes import FakeLedgerBackend +from tokenops.control.core import Observation, Usage +from tokenops.control.ledger import RUN_TOTAL_BUDGET, UNLIMITED_LEFT, Budget, Ledger, segment_key + + +def _ledger(): + cap = Budget(budget_id="run_llm_cap", limit_micros=1_000_000, dimension="run") + backend = FakeLedgerBackend() + return Ledger(budgets=[cap], price=toy_price, backend=backend), cap, backend + + +def test_store_and_backend_are_mutually_exclusive(): + with pytest.raises(ValueError): + Ledger(store=object(), backend=FakeLedgerBackend()) + + +def test_pricing_and_cum_spent_via_backend(): + ledger, _cap, _backend = _ledger() + attr = make_attr() + ledger.open_run("run-1") + + step = ledger.record( + Observation( + attr=attr, + node_type="llm", + boundary_id="chat", + ts=1.0, + provider="openai", + model="gpt-4o-mini", + usage=Usage(input=820, output=45), + ) + ) + cost1 = 820 * 10 + 45 * 30 + assert step.step == 1 and step.cum_spent_micros == cost1 + assert ledger.cost_micros("run-1") == cost1 + + +def test_zero_cost_crossing_writes_no_spend_via_backend(): + ledger, _cap, backend = _ledger() + attr = make_attr() + ledger.open_run("run-1") + ledger.record( + Observation( + attr=attr, + node_type="tool", + boundary_id="search", + ts=1.0, + signature="sig", + result_hash="rh", + ) + ) + assert ledger.cost_micros("run-1") == 0 + # only a `step` event was applied — no spent_add touched the budget + assert backend._spent == {} + + +def test_budget_left_via_backend(): + ledger, cap, _backend = _ledger() + attr = make_attr() + ledger.open_run("run-1") + ledger.record( + Observation( + attr=attr, + node_type="llm", + boundary_id="chat", + ts=1.0, + provider="openai", + model="gpt-4o-mini", + usage=Usage(input=100, output=10), + ) + ) + spent = 100 * 10 + 10 * 30 + sk = segment_key(attr, cap) + assert ledger.budget_left(cap.budget_id, sk) == cap.limit_micros - spent + assert ledger.budget_left(RUN_TOTAL_BUDGET.budget_id, f"run:{attr.run_id}") == UNLIMITED_LEFT + assert ledger.budget_left("unknown-budget", sk) == 0 + + +def test_admit_complete_inflight_via_backend(): + ledger, _cap, _backend = _ledger() + ledger.open_run("run-1") + sk = "run:run-1" + assert ledger.inflight(sk) == 0 + ledger.admit(sk) + ledger.admit(sk) + assert ledger.inflight(sk) == 2 + ledger.complete(sk) + assert ledger.inflight(sk) == 1 + ledger.complete(sk) + ledger.complete(sk) # floored at 0, never negative + assert ledger.inflight(sk) == 0 + + +def test_mark_halted_and_clear_halt_via_backend(): + ledger, _cap, backend = _ledger() + ledger.open_run("run-1") + assert ledger.is_halted("run-1") is False + + ledger.mark_halted("run-1", "step_cap: 2") + assert ledger.is_halted("run-1") is True + assert backend._halt["run-1"] == (True, "step_cap: 2") + + ledger.clear_halt("run-1") + assert ledger.is_halted("run-1") is False + assert backend._halt["run-1"] == (False, None) + + +def test_is_halted_sees_a_halt_set_by_another_ledger_instance(): + """Two Ledgers sharing one backend (simulating two processes / a shared plane) — + a halt set on one must be visible from the other without local state.""" + backend = FakeLedgerBackend() + cap = Budget(budget_id="run_llm_cap", limit_micros=1_000_000, dimension="run") + ledger_a = Ledger(budgets=[cap], price=toy_price, backend=backend) + ledger_b = Ledger(budgets=[cap], price=toy_price, backend=backend) + + ledger_a.mark_halted("shared-run", "cost_budget: exhausted") + assert ledger_b.is_halted("shared-run") is True + + +def test_admit_complete_share_inflight_across_ledger_instances(): + backend = FakeLedgerBackend() + ledger_a = Ledger(backend=backend) + ledger_b = Ledger(backend=backend) + sk = "agent:research" # non-run-scoped segment key — no run_id embedded + ledger_a.admit(sk) + ledger_b.admit(sk) + assert ledger_a.inflight(sk) == 2 + assert ledger_b.inflight(sk) == 2 + + +def test_local_run_state_reads_stay_local_not_backend(): + """step_count/velocity/recent/window are the Tier-1 cache — per-process, never a + backend round trip. A second Ledger sharing the same backend does NOT see them.""" + backend = FakeLedgerBackend() + ledger_a = Ledger(price=toy_price, backend=backend) + ledger_b = Ledger(price=toy_price, backend=backend) + attr = make_attr() + ledger_a.open_run("run-1") + ledger_a.record( + Observation( + attr=attr, + node_type="llm", + boundary_id="chat", + ts=1.0, + provider="openai", + model="gpt-4o-mini", + usage=Usage(input=10, output=1), + ) + ) + assert ledger_a.step_count("run-1") == 1 + assert ledger_b.step_count("run-1") == 0 # no shared Tier-1 state, by design