fix(savings): record pre-compression original as ledger before, not forwarded count (#2176)

## Description

`headroom savings` overstates the proxy reduction percentage because the
durable ledger is written with the wrong `before` value.

In `PrometheusMetrics.record_request` the proxy appends a savings event:

```python
if tokens_saved > 0 and not self._stateless:
    savings_ledger.record_savings_event(
        tokens_before=input_tokens,
        tokens_after=max(input_tokens - tokens_saved, 0),
        ...
    )
```

But `input_tokens` here is the optimized, **post-compression** count
that was actually forwarded, not the original. `emit_request_outcome`
(the single funnel that calls `record_request`) passes
`input_tokens=outcome.optimized_tokens`.

The ledger derives the reported reduction as `saved / before`
(`savings_ledger._Bucket.savings_percent`), with `saved = max(before -
after, 0)`. Passing the forwarded count as `before` (and `before -
saved` as `after`) keeps `saved` correct but understates `before` by
`tokens_saved`, so the percentage is inflated:

- original input 1000 tokens, forwarded 600, saved 400 → true reduction
40%.
- recorded as `before=600, after=200` → `400 / 600` = **66.7%** on the
dashboard.

So `headroom savings` (which aggregates this ledger across restarts and
processes) reports a reduction percent well above what actually happened
for all proxy traffic.

## Fix

Reconstruct the original as forwarded + saved:

```python
tokens_before=input_tokens + tokens_saved,   # the pre-compression original
tokens_after=input_tokens,                   # what we forwarded
```

`saved` (= `before - after` = `tokens_saved`) and the stored `cost_usd`
(derived from `saved`) are unchanged; only the `before`/`after` labels
are corrected, so the reduction percent becomes honest.

Closes #

## Type of Change

- [x] Bug fix (non-breaking change that fixes an issue)
- [ ] New feature (non-breaking change that adds functionality)
- [ ] Breaking change (fix or feature that would cause existing
functionality to change)
- [ ] Documentation update
- [ ] Performance improvement
- [ ] Code refactoring (no functional changes)

## Changes Made

- `headroom/proxy/prometheus_metrics.py`: pass
`tokens_before=input_tokens + tokens_saved` and
`tokens_after=input_tokens` to `record_savings_event`, with a comment
explaining that `input_tokens` is the forwarded count.
- `tests/test_savings_ledger_before_forwarded.py`: new regression guard
on the call shape.
- `CHANGELOG.md`: Bug Fixes entry.

## Testing

- [ ] Unit tests pass (`pytest`)
- [x] Linting passes (`ruff check .`)
- [x] Type checking passes (`mypy headroom`)
- [x] New tests added for new functionality
- [ ] Manual testing performed

### Test Output

```text
$ uvx ruff@0.15.17 check headroom/proxy/prometheus_metrics.py tests/test_savings_ledger_before_forwarded.py
All checks passed!
$ uvx mypy@1.20.2 --ignore-missing-imports headroom/proxy/prometheus_metrics.py
Success: no issues found in 1 source file
```

## Real Behavior Proof

- Environment: Windows 11, Python 3.12, `uvx ruff@0.15.17` / `uvx
mypy@1.20.2`. A full `pytest` OOM-kills this box (ML stack import), so I
reproduced the math with a dependency-free script modelling the ledger's
own `saved = before - after` and `saved / before * 100`, and left the
full pytest to CI.
- Exact command / steps: fed a request with original=1000,
forwarded=600, saved=400 through the OLD call shape
(`before=input_tokens`, `after=input_tokens-saved`) and the NEW shape
(`before=input_tokens+saved`, `after=input_tokens`).
- Observed result: OLD → `before=600, after=200`, reported 66.7%; NEW →
`before=1000, after=600`, reported 40.0% (the true reduction). `saved`
is 400 in both, so the cost figure is unaffected.
- Not tested: a live proxy end-to-end run; full local `pytest` deferred
to CI (OOM).

## 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
- [ ] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective or that my
feature works
- [ ] New and existing unit tests pass locally with my changes
- [x] I have updated the CHANGELOG.md if applicable

## Additional Notes

The "unit tests pass locally" box is unchecked because the full suite
imports the ML stack, which I can't run here. The regression test
asserts on the source of `record_request` (the enclosing module imports
the ML stack, so it executes the assertion against the source text
rather than calling the method); the behavioural verification is the
standalone proof above.

---------

Co-authored-by: JerrettDavis <mxjerrett@gmail.com>
Co-authored-by: Tejas Chopra <chopratejas@gmail.com>
This commit is contained in:
Abhay Singh 2026-07-14 21:51:36 +05:30 committed by GitHub
parent f723925be7
commit 195ed90ced
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 74 additions and 2 deletions

View file

@ -116,6 +116,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
### Bug Fixes
* **proxy:** run a cold-start fast pass before background-compression deferral so byte-identical freeze doesn't lock sessions to the uncompressed transcript. Since #1850, a session's provider-cached prefix is frozen in whatever form its cold start forwarded; deferring the WHOLE pipeline (`HEADROOM_BACKGROUND_COMPRESSION=1`, frozen=0, ≥50k tokens) therefore cached the raw transcript and forfeited the session's compression savings for its lifetime — including sub-second lossless wins like `read_lifecycle` stale-read drops, observed in the field as sessions permanently stuck at 0 savings. The deferral branch now runs the pipeline synchronously with the new `skip_kompress=True` kwarg (everything except the Kompress ML stage — the only stage that can blow the request budget per #1171) under a bounded fast-pass budget (`HEADROOM_COLD_START_FAST_PASS_TIMEOUT_SECONDS`, default 10s), forwards the pruned form, and defers only Kompress to the background job (tagged `deferred:kompress_background`). Fail-open: on fast-pass timeout/error the request forwards uncompressed exactly as before. Units routed to Kompress under `skip_kompress` take the same fallback as when the model isn't ready.
* **savings:** record the pre-compression original as the ledger `before`, not the forwarded count. In `PrometheusMetrics.record_request` the durable savings ledger was written with `tokens_before=input_tokens`, but `input_tokens` there is the optimized (post-compression) count that was forwarded (`emit_request_outcome` passes `outcome.optimized_tokens`). `headroom savings` derives the reported reduction percent as saved / before, so understating `before` by `tokens_saved` inflated it — a real 40% reduction (1000 → 600) was reported as ~67% (400 / 600). The event now records `tokens_before=input_tokens + tokens_saved` (the reconstructed original) and `tokens_after=input_tokens` (the forwarded count); the `saved` and cost figures are unchanged.
* **proxy/gemini:** forward a non-JSON upstream error body with its real status instead of a synthetic 502. In `handle_gemini_generate_content` the token-extraction `except (KeyError, TypeError, AttributeError)` guarding `response.json()` omitted `json.JSONDecodeError` / `ValueError`, so a non-JSON body (an HTML/empty error page from an overloaded Google/Vertex/Copilot frontend, common on 5xx/429) escaped to the outer `except Exception` and was returned as a generic 502 — discarding the real upstream `status_code` and body and defeating the client's retry/backoff. The except now catches the JSON/ValueError family, matching the all-non-text sibling branch, so the true status and body are forwarded verbatim.
* **init/codex:** don't overwrite the user's `hooks.json`. `_ensure_codex_hooks` wrote a fresh payload containing only Headroom's two hooks, wholesale-replacing `~/.codex/hooks.json` — so any user-managed Codex hooks (and other top-level keys) were silently destroyed on `headroom init codex`. It now read-merges: existing entries are preserved, Headroom's are deduped on the `headroom-init-codex` marker and appended, matching `_ensure_claude_hooks` / `_ensure_copilot_hooks`.
* **ccr:** detect `read_lifecycle` stale/superseded markers in the retrieve-tool injector so they stay redeemable. Those markers (`[Read content stale: … Retrieve original: hash=<hash>]`) store the original bytes in the CCR store under a valid hash, but none of `CCRToolInjector`'s patterns matched them — every pattern required the word "compressed" or the `<<ccr:` form. So on a frozen-prefix turn (both `read_lifecycle` and prefix freezing are on by default) the injector reported no compressed content, the `headroom_retrieve` tool was not injected, and the model was handed a marker advertising `Retrieve original: hash=X` with no tool to redeem it — silent data loss for stale reads, where retrieval is the only way to recover the original-at-read-time content (the exact case the #1006 guard exists to prevent). Added a pattern matching the load-bearing `Retrieve original: hash=` phrase, aligning the injector with the sibling `read_maturation` marker that was already (incidentally) detected.

View file

@ -703,9 +703,18 @@ class PrometheusMetrics:
# (claude-code, codex, cursor, ...); it falls back to "proxy" only
# when the harness is unidentified.
if tokens_saved > 0 and not self._stateless:
# `input_tokens` here is the optimized (post-compression) count
# that was actually forwarded — see emit_request_outcome, which
# passes `input_tokens=outcome.optimized_tokens`. The ledger's
# `before` is the pre-compression original and `after` is what we
# forwarded, and `headroom savings` derives the reduction percent
# as saved / before. Passing the forwarded count as `before`
# understated the original by `tokens_saved`, inflating that
# percentage (e.g. a real 40% reduction was reported as ~67%).
# Reconstruct the original as forwarded + saved.
savings_ledger.record_savings_event(
tokens_before=input_tokens,
tokens_after=max(input_tokens - tokens_saved, 0),
tokens_before=input_tokens + tokens_saved,
tokens_after=input_tokens,
model=model,
client=client or "proxy",
source="proxy",

View file

@ -0,0 +1,62 @@
"""Proxy savings-ledger events record the original input as ``before``."""
from __future__ import annotations
from typing import Any
import pytest
from headroom.proxy import prometheus_metrics
class _FakeSavingsTracker:
def snapshot(self) -> dict[str, dict[str, int | float]]:
return {"lifetime": {"total_input_tokens": 0, "total_input_cost_usd": 0.0}}
def record_request(self, **kwargs: Any) -> None:
pass
class _FakeOtelMetrics:
def record_proxy_request(self, **kwargs: Any) -> None:
pass
@pytest.mark.asyncio
async def test_record_savings_event_uses_original_input_as_before(
monkeypatch: pytest.MonkeyPatch,
) -> None:
calls: list[dict[str, Any]] = []
def record_savings_event(**kwargs: Any) -> None:
calls.append(kwargs)
monkeypatch.setattr(
prometheus_metrics.savings_ledger,
"record_savings_event",
record_savings_event,
)
metrics = prometheus_metrics.PrometheusMetrics(
savings_tracker=_FakeSavingsTracker(),
otel_metrics=_FakeOtelMetrics(),
)
await metrics.record_request(
provider="anthropic",
model="claude-opus-4-6",
input_tokens=600,
output_tokens=25,
tokens_saved=400,
latency_ms=10.0,
client="claude-code",
)
assert calls == [
{
"tokens_before": 1000,
"tokens_after": 600,
"model": "claude-opus-4-6",
"client": "claude-code",
"source": "proxy",
}
]