From 1199a6332d8e32f877e403d72c09cd6db7e47f77 Mon Sep 17 00:00:00 2001 From: RjyavardhanSingh Date: Thu, 17 Sep 2026 20:56:20 +0530 Subject: [PATCH] Record compaction token savings in ledger step events (#143) --- src/tokenops/control/boundary.py | 4 +- src/tokenops/control/core.py | 2 + src/tokenops/control/crossing.py | 10 ++++ src/tokenops/control/engine.py | 1 + src/tokenops/control/integration.py | 38 +++++++++++--- src/tokenops/control/ledger.py | 2 + src/tokenops/control/ledger_backend.py | 2 + tests/fakes.py | 1 + tests/test_context_compaction.py | 67 +++++++++++++++++++++++++ tests/test_policies_wrap_integration.py | 40 +++++++++++++++ 10 files changed, 158 insertions(+), 9 deletions(-) diff --git a/src/tokenops/control/boundary.py b/src/tokenops/control/boundary.py index 2e85eb7..9d93833 100644 --- a/src/tokenops/control/boundary.py +++ b/src/tokenops/control/boundary.py @@ -56,6 +56,7 @@ def observation_from_crossing( model: str = "", ts: float | None = None, extra_tags: Mapping[str, str] | None = None, + compaction: Mapping[str, int] | None = None, ) -> Observation: reg = require_registration() attr = _build_attribution(reg, service=service) @@ -102,7 +103,7 @@ def observation_from_crossing( raw_roll = getattr(result, "rolled_up_cost_micros", None) if raw_roll is None: raw_roll = input_state.get("rolled_up_cost_micros", 0) - rolled_up = int(raw_roll) if isinstance(raw_roll, (int, float, str)) else 0 + rolled_up = int(raw_roll) if isinstance(raw_roll, int | float | str) else 0 output = dict(result) if isinstance(result, dict) else {"result": str(result)} return Observation( @@ -119,6 +120,7 @@ def observation_from_crossing( result_hash=result_hash, rolled_up_cost_micros=rolled_up, boundary_tags=tags, + compaction=compaction, **_span_fields(service), ) diff --git a/src/tokenops/control/core.py b/src/tokenops/control/core.py index 6f36c67..18f3913 100644 --- a/src/tokenops/control/core.py +++ b/src/tokenops/control/core.py @@ -137,6 +137,8 @@ class Observation: result_hash: str | None = None # delegate only: child run total reported up the A2A hop rolled_up_cost_micros: Micros = 0 + # compaction only: tokens_before / tokens_after / tokens_saved + compaction: Mapping[str, int] | None = None @dataclass(frozen=True, kw_only=True) diff --git a/src/tokenops/control/crossing.py b/src/tokenops/control/crossing.py index 8185a30..5b762c1 100644 --- a/src/tokenops/control/crossing.py +++ b/src/tokenops/control/crossing.py @@ -151,6 +151,15 @@ def on_crossing( # MUTATE model overrides price correctly; fall back to bound governance. provider = str(state.get("provider") or (gov.provider if gov else "") or "") model = str(state.get("model") or (gov.model if gov else "") or "") + # Compaction metadata: carried from _compact_messages via controls.call + compaction = None + governor = getattr(gov, "governor", None) + controls = getattr(governor, "controls", None) + if controls is not None: + call = getattr(controls, "call", None) + comp = getattr(call, "compaction", None) if call is not None else None + if comp is not None: + compaction = dict(comp) obs = observation_from_crossing( boundary_id=boundary_id, kind=kind, @@ -159,6 +168,7 @@ def on_crossing( result=result, provider=provider, model=model, + compaction=compaction, ) emit_observation(obs) diff --git a/src/tokenops/control/engine.py b/src/tokenops/control/engine.py index b00a1a4..769148b 100644 --- a/src/tokenops/control/engine.py +++ b/src/tokenops/control/engine.py @@ -95,6 +95,7 @@ class _ResolvedCall: model_override: str | None = None max_output_tokens: int | None = None compact: bool = False # deep MUTATE: rewrite the outgoing messages (context_compaction) + compaction: dict[str, int] | None = None # tokens_before/after/saved from compaction @dataclass diff --git a/src/tokenops/control/integration.py b/src/tokenops/control/integration.py index 45bb1f4..b92e152 100644 --- a/src/tokenops/control/integration.py +++ b/src/tokenops/control/integration.py @@ -135,7 +135,7 @@ def on_step(ev) -> None: DispatchFn = Callable[..., object] -def _estimate_input_tokens(messages) -> int: +def _estimate_input_tokens(messages: object) -> int: return max(1, len(str(messages)) // 4) @@ -180,11 +180,21 @@ def consume_carry( return out -def _compact_messages(messages): +def _compact_messages( + messages: list[object], + *, + estimate: Callable[[object], int] | None = None, +) -> tuple[list[object], dict[str, int]]: """Deep context_compaction MUTATE: rewrite the outgoing messages — pin every system - message, drop duplicate non-system messages (deduped tool outputs / repeated context).""" - seen: set = set() - out: list = [] + message, drop duplicate non-system messages (deduped tool outputs / repeated context). + + Returns ``(compacted_messages, metadata)`` where *metadata* contains + ``tokens_before``, ``tokens_after`` and ``tokens_saved`` (all ints, ≥ 0). + If *estimate* is ``None`` token counts default to 0 (caller opts out of + measurement). + """ + seen: set[tuple[str | None, str]] = set() + out: list[object] = [] for msg in messages: role = msg.get("role") if isinstance(msg, dict) else None content = msg.get("content", "") if isinstance(msg, dict) else str(msg) @@ -199,7 +209,17 @@ def _compact_messages(messages): continue seen.add(key) out.append(msg) - return out + if estimate is not None: + before = estimate(messages) + after = estimate(out) + meta = { + "tokens_before": before, + "tokens_after": after, + "tokens_saved": max(0, before - after), + } + else: + meta = {"tokens_before": 0, "tokens_after": 0, "tokens_saved": 0} + return out, meta def wrap_complete( @@ -249,7 +269,8 @@ def governed(p: str, m: str, messages) -> object: use_model = controls.call.model_override or m messages = consume_carry(controls, messages) if controls.call.compact: # deep prompt compaction - messages = _compact_messages(messages) + messages, _cmeta = _compact_messages(messages, estimate=estimate) + controls.call.compaction = _cmeta governor.ledger.admit(seg) try: @@ -388,7 +409,8 @@ def governed(p: str, m: str, messages) -> object: use_model = controls.call.model_override or m messages = consume_carry(controls, messages) if controls.call.compact: # deep prompt compaction - messages = _compact_messages(messages) + messages, _cmeta = _compact_messages(messages, estimate=estimate) + controls.call.compaction = _cmeta governor.ledger.admit(seg) try: diff --git a/src/tokenops/control/ledger.py b/src/tokenops/control/ledger.py index 6a65f55..c25bfb9 100644 --- a/src/tokenops/control/ledger.py +++ b/src/tokenops/control/ledger.py @@ -224,6 +224,8 @@ def _step_event(obs: Observation, seq: int, cost: Micros) -> LedgerEvent: event["tool_signature"] = obs.signature if obs.result_hash is not None: event["result_hash"] = obs.result_hash + if obs.compaction is not None: + event["compaction"] = dict(obs.compaction) return event diff --git a/src/tokenops/control/ledger_backend.py b/src/tokenops/control/ledger_backend.py index 664f257..07d49a4 100644 --- a/src/tokenops/control/ledger_backend.py +++ b/src/tokenops/control/ledger_backend.py @@ -63,6 +63,8 @@ class LedgerEvent(TypedDict, total=False): tags: dict[str, str] tool_signature: str result_hash: str + # compaction + compaction: dict[str, int] # halt_mark reason: str detector: str diff --git a/tests/fakes.py b/tests/fakes.py index c734917..046151e 100644 --- a/tests/fakes.py +++ b/tests/fakes.py @@ -162,6 +162,7 @@ def _apply_one(self, ev: LedgerEvent) -> None: "tool_signature", "result_hash", "ts", + "compaction", ) if ev.get(k) is not None } diff --git a/tests/test_context_compaction.py b/tests/test_context_compaction.py index aad030d..fdf9a5c 100644 --- a/tests/test_context_compaction.py +++ b/tests/test_context_compaction.py @@ -6,6 +6,7 @@ from tokenops.control import ActionKind, CallRequest, Usage from tokenops.control.context import reset_current_controls, set_current_controls from tokenops.control.engine import ApplyControls +from tokenops.control.integration import _compact_messages from tokenops.control.policies import context_compaction @@ -52,3 +53,69 @@ def test_no_hook_is_telemetry_only(): sig = det.pre_call(_req(10_000), FakeView()) # No controls in context → compaction not supported → ALLOW assert pol.decide(sig, FakeView()).kind is ActionKind.ALLOW + + +# --------------------------------------------------------------------------- # +# compaction token-recording tests (issue #143) # +# --------------------------------------------------------------------------- # + +_MSGS = [ + {"role": "system", "content": "You are helpful."}, + {"role": "user", "content": "Hello"}, + {"role": "assistant", "content": "Hi there!"}, + {"role": "user", "content": "Hello"}, # duplicate — will be dropped +] + + +def _estimate(messages): + return len(str(messages)) // 4 + + +def test_compact_messages_returns_metadata_with_estimate(): + compacted, meta = _compact_messages(_MSGS, estimate=_estimate) + assert len(compacted) < len(_MSGS) + assert meta["tokens_before"] > 0 + assert meta["tokens_after"] > 0 + assert meta["tokens_saved"] == meta["tokens_before"] - meta["tokens_after"] + assert meta["tokens_saved"] >= 0 + + +def test_compact_messages_zeroed_metadata_without_estimate(): + compacted, meta = _compact_messages(_MSGS) + assert len(compacted) < len(_MSGS) + assert meta == {"tokens_before": 0, "tokens_after": 0, "tokens_saved": 0} + + +def test_compact_messages_drops_duplicates_preserves_system(): + compacted, _ = _compact_messages(_MSGS, estimate=_estimate) + roles_and_content = [(m.get("role"), m.get("content")) for m in compacted] + # System message always kept + assert ("system", "You are helpful.") in roles_and_content + # Duplicate user message dropped + assert roles_and_content.count(("user", "Hello")) == 1 + + +def test_observation_carry_compaction_when_present(): + from tokenops.control.core import Attribution, Observation + + obs = Observation( + attr=Attribution(user="u", agent="a", run_id="r"), + node_type="llm", + boundary_id="a.chat", + ts=1.0, + compaction={"tokens_before": 100, "tokens_after": 60, "tokens_saved": 40}, + ) + assert obs.compaction is not None + assert obs.compaction["tokens_saved"] == 40 + + +def test_observation_no_compaction_when_none(): + from tokenops.control.core import Attribution, Observation + + obs = Observation( + attr=Attribution(user="u", agent="a", run_id="r"), + node_type="llm", + boundary_id="a.chat", + ts=1.0, + ) + assert obs.compaction is None diff --git a/tests/test_policies_wrap_integration.py b/tests/test_policies_wrap_integration.py index 8134567..30b73ac 100644 --- a/tests/test_policies_wrap_integration.py +++ b/tests/test_policies_wrap_integration.py @@ -375,6 +375,46 @@ def run(): assert any(m.get("content") == "unique" for m in sent) +def test_it_context_compaction_records_tokens_in_ledger_event(): + """Compaction metadata (tokens_before / tokens_after / tokens_saved) flows through + wrap_complete → crossing hook → Observation → step event in the ledger.""" + from fakes import FakeLedgerBackend + + backend = FakeLedgerBackend() + controls = ApplyControls() + gov = Governor(Ledger(price=toy_price, backend=backend), controls) + gov.register(*context_compaction.build(ctx_max=10)) # tiny ctx → always trips + attr = _attr("r-cc-meta") + gov.ledger.open_run("r-cc-meta") + dispatch, calls = _dispatch(inp=80, out=20) + governed = _governed(gov, attr, dispatch, run_id="r-cc-meta") + msgs = [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "dup"}, + {"role": "user", "content": "dup"}, + {"role": "user", "content": "unique"}, + ] + + def run(): + governed("openai", "gpt-4o-mini", msgs) + + _with_scope(gov, attr, "r-cc-meta", run) + + # Verify the step event in the backend carries compaction metadata + state = backend._run_state.get("r-cc-meta") + assert state is not None + recent = state["recent"] + assert len(recent) >= 1 + llm_step = recent[-1] + assert llm_step.get("node_type") == "llm" + comp = llm_step.get("compaction") + assert comp is not None, "compaction metadata missing from step event" + assert comp["tokens_before"] > 0 + assert comp["tokens_after"] > 0 + assert comp["tokens_saved"] >= 0 + assert comp["tokens_saved"] == comp["tokens_before"] - comp["tokens_after"] + + def test_it_output_runaway_retries_then_succeeds(): controls = ApplyControls() gov = Governor(Ledger(price=toy_price), controls)