From 1cc53c9c92cd4dffaf048dc806cb8c570bdb86b6 Mon Sep 17 00:00:00 2001 From: Tejas Chopra Date: Fri, 24 Jul 2026 14:54:43 -0700 Subject: [PATCH] fix(proxy/perf): tokenizer-consistent token accounting + surface tool-schema savings (#2542) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description Follow-up to #2520 (turn-hook message-fold accounting). While validating that PR on live Claude Code traffic, two accounting defects surfaced: 1. **Impossible/misleading token deltas.** The handler and the compression pipeline use *different* token estimators — the handler's `EstimatingTokenCounter(3.5)` (or real tiktoken on OpenAI) vs `content_router`'s adaptive `EstimatingTokenCounter()`. Cross-assigning `original_tokens` (handler) against `optimized_tokens = result.tokens_after` (pipeline) put the two endpoints on different scales, producing **`tok_after > tok_before` on 101/783 PERF lines** and phantom savings on `transforms=none` lines. It also made the turn-hook recount fire on the *scale difference* rather than a real fold, emitting a **spurious `turn_hook` tag with `tok_saved=0`**. 2. **Tool-schema savings were invisible.** Tool deferral (`defer_loading`) and turn-hook tool shrink save thousands of tool-schema tokens, but `tok_before`/`tok_after` count **messages only** — so a tool-heavy turn logged `tok_saved=0` while genuinely saving ~29k tool-schema tokens, reading as "no compression." Closes # ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) - [x] Performance improvement - [ ] Breaking change - [ ] Documentation update - [ ] Code refactoring (no functional changes) ## Changes Made - **Tier B.1 — tokenizer-consistent accounting.** On the Anthropic and OpenAI-chat paths, recount **both** endpoints with the **same** tokenizer (pre-compression snapshot vs final outbound messages) right before the outcome is recorded. This puts `tok_before`/`tok_after` on one scale (fixes the inflated + phantom lines), and it subsumes any turn-hook fold. `turn_hook` is now attributed **only** when the hook itself reduced tokens (same-tokenizer pre vs post), not when a recount merely normalized a scale difference. OpenAI preserves its existing tool-schema delta folding. - **Surface tool-schema savings.** New `tool_saved=` field on the PERF line (summed from `tool_search_deferred_tokens` + `turn_hook_tools_saved_tokens` tags) and a separate `Tool saved` line in `headroom perf`. Additive + backward-compatible (key=value parse; old lines default to 0). `tok_saved` still means message savings, so ratios and calibrated thresholds are unaffected. - The OpenAI **Responses** path already accumulates per-transform deltas (each consistent within its own transform), so it isn't exposed to the cross-scale subtraction bug and needs no change. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check` + `ruff format --check`) - [x] Type checking passes (`mypy`) - [x] Manual testing performed ### Test Output ```text $ ruff check headroom/proxy/handlers/anthropic.py headroom/proxy/handlers/openai.py \ headroom/proxy/outcome.py headroom/perf/analyzer.py All checks passed! $ ruff format --check 4 files already formatted $ mypy Success: no issues found in 4 source files $ pytest tests/test_turn_hooks.py tests/test_openai_chat_turn_hooks.py \ tests/test_openai_responses_context_compaction.py -q 26 passed in 14.34s ``` ## Real Behavior Proof - **Environment:** local proxy `headroom proxy --port 8793 --proxy-extension lossless_guard`, with `HEADROOM_LOSSLESS_GUARD_LOSSY=1`, `HEADROOM_TOOL_SEARCH=1`, model `claude-haiku-4-5` (Anthropic path = the one Claude Code uses). - **Observed, before vs after this PR:** - foldable tool_result: `tok_before=607 tok_after=178 tok_saved=429 transforms=turn_hook` (real fold, correctly attributed) - plain multi-turn (nothing foldable): `tok_before=3700 tok_after=3700 tok_saved=0 transforms=none` — **no inflation, no spurious `turn_hook`** (before this PR: same request showed a spurious `turn_hook`) - tool-heavy (12 tools): `tok_saved=0 tool_saved=1794` → `headroom perf` shows `Total saved: … (messages)` **and** `Tool saved: 1,794 tokens (tool schemas, deferral)` (before: the 1,794 was invisible) - **Not tested:** OpenAI chat/Responses paths verified by unit test, not a live client run (local setup routes Claude Code through the Anthropic handler only). ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## Checklist - [x] My code follows the project's style guidelines - [x] I have performed a self-review of my code - [x] I have commented my code, particularly in hard-to-understand areas - [ ] Documentation changes — N/A (internal accounting; `tool_saved` is self-describing in `headroom perf`) - [x] My changes generate no new warnings - [x] New and existing unit tests pass locally with my changes - [x] I did **not** edit `CHANGELOG.md` ## Additional Notes - Behavior is unchanged when no turn hook is registered for the *attribution* tag; the consistency recount runs unconditionally so pure-OSS installs also get correct before/after (it only ever makes the two endpoints comparable — it never fabricates savings). - Follow-up (separate PR, intentionally not here): swap the char-estimator for a real BPE (tiktoken `o200k_base`) for private-tokenizer models like Claude. That's the "Tier B.2" accuracy upgrade; it shifts absolute numbers ~10–20% and touches calibrated thresholds, so it needs its own recalibration pass. --- headroom/perf/analyzer.py | 11 ++++++- headroom/proxy/handlers/anthropic.py | 44 ++++++++++++++++++---------- headroom/proxy/handlers/openai.py | 31 +++++++++++++++----- headroom/proxy/outcome.py | 10 +++++++ 4 files changed, 72 insertions(+), 24 deletions(-) diff --git a/headroom/perf/analyzer.py b/headroom/perf/analyzer.py index 39e676c9b..b5f8487f3 100644 --- a/headroom/perf/analyzer.py +++ b/headroom/perf/analyzer.py @@ -146,6 +146,7 @@ class PerfRecord: tokens_before: int = 0 tokens_after: int = 0 tokens_saved: int = 0 + tool_saved: int = 0 cache_read: int = 0 cache_write: int = 0 cache_hit_pct: int = 0 @@ -351,6 +352,7 @@ def parse_log_files(last_n_hours: float = 168.0) -> PerfReport: tokens_before=int(kv.get("tok_before", 0)), tokens_after=int(kv.get("tok_after", 0)), tokens_saved=int(kv.get("tok_saved", 0)), + tool_saved=int(kv.get("tool_saved", 0)), cache_read=int(kv.get("cache_read", 0)), cache_write=int(kv.get("cache_write", 0)), cache_hit_pct=int(kv.get("cache_hit_pct", 0)), @@ -537,11 +539,17 @@ def format_report(report: PerfReport) -> str: total_before = sum(r.tokens_before for r in records) total_after = sum(r.tokens_after for r in records) total_saved = sum(r.tokens_saved for r in records) + total_tool_saved = sum(r.tool_saved for r in records) pct = (total_saved / total_before * 100) if total_before > 0 else 0 lines.append(f"Requests: {len(records)}") lines.append(f"Tokens: {total_before:,} -> {total_after:,} ({pct:.1f}% reduction)") - lines.append(f"Total saved: {total_saved:,} tokens") + lines.append(f"Total saved: {total_saved:,} tokens (messages)") + # Tool-schema savings (deferral + turn-hook tool shrink) are counted apart + # from message compression — messages never include tool bytes — so surface + # them explicitly instead of hiding a tool-heavy turn's win behind tok_saved=0. + if total_tool_saved > 0: + lines.append(f"Tool saved: {total_tool_saved:,} tokens (tool schemas, deferral)") lines.append("") # Per-model breakdown with list prices @@ -786,6 +794,7 @@ PERF_RECORD_FIELDS = [ "tokens_before", "tokens_after", "tokens_saved", + "tool_saved", "cache_read", "cache_write", "cache_hit_pct", diff --git a/headroom/proxy/handlers/anthropic.py b/headroom/proxy/handlers/anthropic.py index 1162fd2c2..05dd6fdd8 100644 --- a/headroom/proxy/handlers/anthropic.py +++ b/headroom/proxy/handlers/anthropic.py @@ -2385,6 +2385,7 @@ class AnthropicHandlerMixin: run_request_hooks, ) + _pre_hook_tokens: int | None = None if registered_turn_hooks(): _req_ctx = TurnContext( provider="anthropic", @@ -2393,6 +2394,13 @@ class AnthropicHandlerMixin: tools=body.get("tools"), config=self.config, ) + # Snapshot BEFORE the hook (same tokenizer) so we can tell whether the + # hook itself folded — comparing against the pipeline's optimized_tokens + # instead conflates a real fold with a cross-estimator delta (see below). + try: + _pre_hook_tokens = tokenizer.count_messages(optimized_messages) + except Exception: + _pre_hook_tokens = None run_request_hooks(_req_ctx) if _req_ctx.messages is not optimized_messages: optimized_messages = _req_ctx.messages @@ -2400,21 +2408,27 @@ class AnthropicHandlerMixin: if _req_ctx.tools is not body.get("tools"): tools = _req_ctx.tools body["tools"] = tools - # Turn hooks (e.g. lossless-guard) fold messages AFTER the pipeline's - # token accounting, and may mutate them IN PLACE (identity unchanged), - # so their savings were invisible to the PERF line / `headroom perf` - # (record_compression /stats already counts them). Re-count regardless - # of replace-vs-in-place so original->optimized reflects the fold too. - # tokenizer is initialized → count_messages is a pure CPU call here. - # Only ever lowers optimized_tokens. - try: - _hooked_tokens = tokenizer.count_messages(optimized_messages) - if _hooked_tokens < optimized_tokens: - optimized_tokens = _hooked_tokens - tokens_saved = max(0, original_tokens - optimized_tokens) - transforms_applied.append("turn_hook") - except Exception: - logger.debug("turn-hook token re-count skipped", exc_info=True) + + # Consistency: report tok_before/tok_after with ONE tokenizer. The pipeline + # and the handler use different token estimators, and cache-mode branches + # can leave original_tokens (handler, line ~1049) and optimized_tokens + # (pipeline, result.tokens_after) on different scales — which produced + # impossible tok_after>tok_before deltas and masked real savings. Recount + # BOTH endpoints (pre-compression snapshot vs final outbound messages) with + # the handler tokenizer so the delta is meaningful. This also captures any + # turn-hook fold (optimized_messages is post-hook). Runs unconditionally. + try: + _orig_snapshot = original_client_messages # noqa: F821 (bound at request start) + original_tokens = tokenizer.count_messages(_orig_snapshot) + optimized_tokens = tokenizer.count_messages(optimized_messages) + tokens_saved = max(0, original_tokens - optimized_tokens) + # Attribute the fold to the hook ONLY when the hook itself reduced + # tokens (same-tokenizer pre vs post) — not when the recount above + # merely normalized a cross-estimator scale difference. + if _pre_hook_tokens is not None and optimized_tokens < _pre_hook_tokens: + transforms_applied.append("turn_hook") + except Exception: + logger.debug("consistency token re-count skipped", exc_info=True) # Output shaping (opt-in via HEADROOM_OUTPUT_SHAPER): verbosity # steering appended to the system-prompt tail + effort routing on diff --git a/headroom/proxy/handlers/openai.py b/headroom/proxy/handlers/openai.py index 93c5147e5..b80d6a6bc 100644 --- a/headroom/proxy/handlers/openai.py +++ b/headroom/proxy/handlers/openai.py @@ -3453,6 +3453,12 @@ class OpenAIHandlerMixin: body["tools"] = tools if presend_event.headers is not None: headers = presend_event.headers + # Consistency: recount BOTH endpoints with the provider tokenizer. An upstream + # branch may have left original_tokens in the pipeline's char-estimator scale + # (result.tokens_before), which mismatches optimized_tokens (provider tokenizer) + # and yields impossible tok_after>tok_before. Recount original from the + # pre-compression snapshot so the message delta is on one scale. + original_tokens = tokenizer.count_messages(original_client_messages) optimized_tokens = tokenizer.count_messages(body["messages"]) if tool_tokens_before_compaction > 0: try: @@ -3462,7 +3468,7 @@ class OpenAIHandlerMixin: if 0 < tool_tokens_after_compaction < tool_tokens_before_compaction: original_tokens += tool_tokens_before_compaction optimized_tokens += tool_tokens_after_compaction - tokens_saved = original_tokens - optimized_tokens + tokens_saved = max(0, original_tokens - optimized_tokens) # Turn hooks (opt-in extensions): a registered hook may rewrite the # outbound tools/messages before we send. Buffered requests only — a @@ -3489,6 +3495,13 @@ class OpenAIHandlerMixin: tools=_th_tools_before, config=self.config, ) + # Snapshot messages BEFORE the hook (same tokenizer) so we can tell whether + # the hook itself folded — comparing against optimized_tokens instead + # conflates a real fold with a cross-estimator scale delta. + try: + _th_msg_before: int | None = tokenizer.count_messages(body["messages"]) + except Exception: + _th_msg_before = None run_request_hooks(_th_ctx) # A hook may either replace ctx.messages/ctx.tools or mutate them in # place (the contract allows both). Use object identity only to decide @@ -3500,15 +3513,17 @@ class OpenAIHandlerMixin: if _th_ctx.tools is not _th_tools_before: tools = _th_ctx.tools body["tools"] = tools - # Message folds land AFTER the accounting above, and a hook may mutate - # messages IN PLACE (identity unchanged), so re-count regardless or - # `headroom perf` sees 0 for them (record_compression /stats already - # does). tokenizer is initialized → pure CPU. Only lowers the count. + # Recount messages after the hook (it may fold in place), preserving the + # tool-schema delta already folded into the headline above and keeping the + # scale consistent with original_tokens (both provider tokenizer). try: _th_msg_after = tokenizer.count_messages(body["messages"]) - if _th_msg_after < optimized_tokens: - optimized_tokens = _th_msg_after - tokens_saved = max(0, original_tokens - optimized_tokens) + optimized_tokens = _th_msg_after + if 0 < tool_tokens_after_compaction < tool_tokens_before_compaction: + optimized_tokens += tool_tokens_after_compaction + tokens_saved = max(0, original_tokens - optimized_tokens) + # Attribute to the hook ONLY when the hook itself reduced tokens. + if _th_msg_before is not None and _th_msg_after < _th_msg_before: transforms_applied.append("turn_hook") except Exception: logger.debug("turn-hook token re-count skipped", exc_info=True) diff --git a/headroom/proxy/outcome.py b/headroom/proxy/outcome.py index 1fec166d6..41a87dc4a 100644 --- a/headroom/proxy/outcome.py +++ b/headroom/proxy/outcome.py @@ -470,11 +470,21 @@ async def emit_request_outcome(handler: Any, outcome: RequestOutcome) -> None: # line unchanged, and gives ``headroom perf --client X`` # parsers a clean key to filter on. client_part = f" client={outcome.client}" if outcome.client else "" + # Tool-schema savings are tracked separately from message compression: tool + # deferral (defer_loading) and turn-hook tool shrink don't move tok_before/after + # (those count messages only), so a tool-heavy turn shows tok_saved=0 while + # genuinely saving thousands of tool-schema tokens. Surface it as its own field + # so `headroom perf` / log readers see the whole picture. + _tags = outcome.tags or {} + tool_saved = int(_tags.get("tool_search_deferred_tokens", 0) or 0) + int( + _tags.get("turn_hook_tools_saved_tokens", 0) or 0 + ) logger.info( f"[{outcome.request_id}] PERF " f"model={outcome.model} msgs={outcome.num_messages} " f"tok_before={outcome.original_tokens} tok_after={outcome.optimized_tokens} " f"tok_saved={outcome.tokens_saved} " + f"tool_saved={tool_saved} " f"cache_read={outcome.cache_read_tokens} cache_write={outcome.cache_write_tokens} " f"cache_hit_pct={outcome.cache_hit_pct} " f"opt_ms={outcome.overhead_ms:.0f} "