mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
fix(proxy/openai): don't record Codex WS savings without input accounting (#2493)
## Description
On the Codex WS Responses path (`handle_openai_responses_ws`),
`tokens_saved` accumulates at compression time (our own token count),
while input tokens only arrive with a usage frame on
`response.completed`. A turn that is compressed but never completes —
cancelled mid-response (Esc in Codex), or an upstream error before the
usage frame — records `tokens_saved > 0` with `input_tokens == 0`
through the outcome funnel.
That writes a savings-with-zero-spend checkpoint into the savings
tracker: `compression_savings_usd` advances while `total_input_tokens` /
`total_input_cost_usd` stay flat. `/stats-history` then serves daily
buckets with `compression_savings_usd_delta > 0` and
`total_input_tokens_delta == 0`, which savings dashboards flag as a
data-integrity anomaly ("graph shows compression savings but zero tokens
spent on recent day(s)").
Both WS record sites have the hazard:
- the per-turn metrics closure (`_record_ws_response_metrics`) records
per-field-clamped deltas, so a usage-less turn contributes a
savings-only outcome;
- the session-end residual flush records `residual_tokens_saved` with
`residual_input_tokens` possibly 0.
## Type of Change
- [x] Bug fix (non-breaking change that fixes an issue)
## Changes Made
- `headroom/proxy/handlers/openai.py`:
- New module-level pure helper `_deferrable_savings_delta(input_delta,
saved_delta)` — returns 0 when `saved_delta > 0` with `input_delta <=
0`, passes everything else through unchanged.
- Per-turn metrics closure: gate `saved_delta` through the helper, and
advance `ws_recorded_tokens_saved_total += saved_delta` (previously `=
tokens_saved`) so deferred savings stay pending and ride along with the
next usage-carrying turn instead of being silently dropped.
- Session-end residual flush: gate `residual_tokens_saved` through the
same helper — savings that never saw a usage frame by session close are
dropped rather than recorded against zero spend (the spend for those
turns is genuinely unknown).
- `tests/test_codex_ws_savings_deferral.py`: truth-table test for the
helper; a bookkeeping walk asserting deferred savings land with the next
usage-carrying turn; and a source-level regression guard for the
closure-internal wiring (same idiom as
`test_codex_ws_compression_scheduler.py`, since the WS closures have no
unit harness yet).
## Testing
- [x] Unit tests pass (`pytest`)
- [x] Linting passes (`ruff check .`)
- [x] Type checking passes (`mypy headroom`)
- [x] New tests added for new functionality
- [x] Manual testing performed (real-behavior script below)
### Test Output
```text
$ uv run --frozen --extra dev pytest tests/test_codex_ws_savings_deferral.py tests/test_openai_codex_ws_lifecycle.py tests/test_codex_ws_compression_scheduler.py tests/test_openai_codex_ws_timings.py tests/test_proxy_savings_history.py
83 passed, 1 skipped (pre-existing pending-harness skip)
$ ruff check headroom/proxy/handlers/openai.py tests/test_codex_ws_savings_deferral.py
All checks passed!
$ uv run --frozen --extra dev mypy headroom/proxy/handlers/openai.py
Success: no issues found in 1 source file
```
## Real Behavior Proof
- Environment: macOS arm64 (Darwin 24.6.0), Python 3.12, this branch
checked out in the repo, run via `uv run --frozen python`.
- Exact command / steps: `uv run --frozen python rbp_demo.py` — a
real-behavior script exercising the REAL `SavingsTracker` (persistence +
`/stats-history` rollup via `history_response()`) and the REAL
`_deferrable_savings_delta` from this branch, no mocks. Scenario: turn 1
compressed (1200 saved) then cancelled before its usage frame (recorded
on day 1), turn 2 compressed (600 more) and completed with
`input_tokens=40000` (day 2); "BEFORE" records what the unfixed handler
emitted, "AFTER" walks the fixed bookkeeping. Additionally, a real
production `~/.headroom/proxy_savings.json` (5000 checkpoints, live
proxy in daily Claude Code + Codex use) was scanned for consecutive
checkpoint pairs where `compression_savings_usd` grew while
`total_input_tokens` stayed flat — one such pair was present
(`provider=openai, model=gpt-5.4-mini`, a Codex WS turn), exactly the
shape this PR removes at the source.
- Observed result: the unfixed recording produces a day-1
`/stats-history` bucket with `compression_savings_usd_delta > 0` and
`total_input_tokens_delta == 0` (the flagged anomaly); the fixed
bookkeeping produces no such bucket and preserves the full 1800 tokens
of savings, paired with the usage-carrying turn. Full output:
```text
BEFORE (unfixed recording): [{'tokens_saved': 1200, 'compression_savings_usd_delta': 0.0009, 'total_input_tokens_delta': 0},
{'tokens_saved': 600, 'compression_savings_usd_delta': 0.00045, 'total_input_tokens_delta': 40000}]
AFTER (fixed recording): [{'tokens_saved': 1800, 'compression_savings_usd_delta': 0.00135, 'total_input_tokens_delta': 40000}]
desync bucket present before fix: True
desync bucket present after fix: False
total savings preserved after fix: True
```
- Not tested: a live end-to-end WS session against the real OpenAI
upstream with a mid-response cancel (needs a real Codex client +
billable upstream). The per-turn/residual closure wiring is covered by
the source-level regression guard instead, per the pending-harness note
in `test_codex_ws_compression_scheduler.py`.
## 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 own code
- [x] I have commented my code, particularly in hard-to-understand areas
- [ ] I have made corresponding changes to the documentation (N/A —
internal accounting fix, no user-facing docs affected)
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective or that my
feature works
- [x] New and existing unit tests pass locally with my changes
- [x] I did **not** edit `CHANGELOG.md` — it is generated by
release-please from the Conventional Commit PR title (a CI guard
enforces this)
## Additional Notes
- Sessions that end on a cancelled turn under-report savings slightly
(the deferred savings are dropped at close because their spend is
genuinely unknown). This is the honest trade-off: the alternative —
recording savings against zero spend — is the desync this PR removes.
- `attempted_input_tokens` is intentionally not gated: a cancelled turn
still records its attempted delta, keeping funnel-drop visibility.
🤖 Generated with [Claude Code](https://claude.com/claude-code)
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
parent
5a0a5a79cd
commit
2195ba7d91
2 changed files with 135 additions and 3 deletions
|
|
@ -1249,6 +1249,26 @@ def _infer_openai_cache_write_tokens(input_tokens: int, cache_read_tokens: int)
|
|||
return max(input_tokens - cache_read_tokens, 0)
|
||||
|
||||
|
||||
def _deferrable_savings_delta(input_delta: int, saved_delta: int) -> int:
|
||||
"""Gate a WS Responses turn's compression-savings delta on input accounting.
|
||||
|
||||
``tokens_saved`` accumulates at compression time (our own token count),
|
||||
while input tokens only arrive with a usage frame on
|
||||
``response.completed``. A turn that was compressed but never completed
|
||||
(cancelled mid-response, upstream error) therefore shows
|
||||
``saved_delta > 0`` with ``input_delta == 0`` — and recording that pair
|
||||
writes a savings-with-zero-spend checkpoint into the savings tracker,
|
||||
desyncing every downstream funnel (dashboards flag "savings but no
|
||||
spend"). Return 0 for that case so the caller defers the savings until a
|
||||
usage-carrying turn; non-positive deltas pass through unchanged so the
|
||||
recorded-total bookkeeping keeps its normal resync behaviour.
|
||||
"""
|
||||
|
||||
if saved_delta > 0 and input_delta <= 0:
|
||||
return 0
|
||||
return saved_delta
|
||||
|
||||
|
||||
def _extract_responses_usage(event: dict[str, Any]) -> tuple[int, int, int, int, int]:
|
||||
"""Return input/output/cache usage from a Responses event.
|
||||
|
||||
|
|
@ -6978,7 +6998,16 @@ class OpenAIHandlerMixin:
|
|||
ws_uncached_input_tokens_total
|
||||
- ws_recorded_uncached_input_tokens_total
|
||||
)
|
||||
saved_delta = tokens_saved - ws_recorded_tokens_saved_total
|
||||
# Usage-less turn (cancelled/failed before
|
||||
# response.completed): defer its savings — see
|
||||
# _deferrable_savings_delta. The recorded total
|
||||
# advances by the recorded delta below, so
|
||||
# deferred savings stay pending and land with the
|
||||
# next usage-carrying turn.
|
||||
saved_delta = _deferrable_savings_delta(
|
||||
input_delta,
|
||||
tokens_saved - ws_recorded_tokens_saved_total,
|
||||
)
|
||||
attempted_delta = (
|
||||
attempted_input_tokens_total
|
||||
- ws_recorded_attempted_input_tokens_total
|
||||
|
|
@ -7091,7 +7120,13 @@ class OpenAIHandlerMixin:
|
|||
ws_recorded_cache_read_tokens_total = ws_cache_read_tokens_total
|
||||
ws_recorded_cache_write_tokens_total = ws_cache_write_tokens_total
|
||||
ws_recorded_uncached_input_tokens_total = ws_uncached_input_tokens_total
|
||||
ws_recorded_tokens_saved_total = tokens_saved
|
||||
# Advance by the recorded delta, not to the live
|
||||
# total: when the usage-less guard above zeroed
|
||||
# saved_delta, the un-recorded savings must stay
|
||||
# pending for the next usage-carrying turn.
|
||||
# Equivalent to `= tokens_saved` whenever the
|
||||
# delta was recorded as computed.
|
||||
ws_recorded_tokens_saved_total += saved_delta
|
||||
ws_recorded_attempted_input_tokens_total = attempted_input_tokens_total
|
||||
ws_recorded_overhead_ms_total = _current_ws_overhead_ms()
|
||||
ws_recorded_compression_timing_totals.update(
|
||||
|
|
@ -7552,7 +7587,15 @@ class OpenAIHandlerMixin:
|
|||
0,
|
||||
ws_uncached_input_tokens_total - ws_recorded_uncached_input_tokens_total,
|
||||
)
|
||||
residual_tokens_saved = max(0, tokens_saved - ws_recorded_tokens_saved_total)
|
||||
# Savings deferred from usage-less turns (per-turn guard in
|
||||
# _record_ws_response_metrics) land here when the session closes
|
||||
# before another usage frame arrives. With no residual input there
|
||||
# is no spend to pair them with — drop them rather than write the
|
||||
# savings-with-zero-spend checkpoint the per-turn guard prevents.
|
||||
residual_tokens_saved = _deferrable_savings_delta(
|
||||
residual_input_tokens,
|
||||
max(0, tokens_saved - ws_recorded_tokens_saved_total),
|
||||
)
|
||||
residual_attempted_input_tokens = max(
|
||||
0,
|
||||
attempted_input_tokens_total - ws_recorded_attempted_input_tokens_total,
|
||||
|
|
|
|||
89
tests/test_codex_ws_savings_deferral.py
Normal file
89
tests/test_codex_ws_savings_deferral.py
Normal file
|
|
@ -0,0 +1,89 @@
|
|||
"""Codex WS Responses: never record compression savings without input accounting.
|
||||
|
||||
``tokens_saved`` accumulates at compression time (our own count); input tokens
|
||||
only arrive with a usage frame on ``response.completed``. A cancelled or failed
|
||||
turn therefore produces a savings delta with no input delta — recording that
|
||||
pair writes a savings-with-zero-spend checkpoint into the savings tracker and
|
||||
desyncs every downstream funnel (dashboards flag "compression savings but zero
|
||||
tokens spent").
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from pathlib import Path
|
||||
|
||||
from headroom.proxy.handlers.openai import _deferrable_savings_delta
|
||||
|
||||
OPENAI_HANDLER = Path(__file__).parent.parent / "headroom" / "proxy" / "handlers" / "openai.py"
|
||||
|
||||
|
||||
def test_deferrable_savings_delta_gates_on_input() -> None:
|
||||
# Normal turn: usage arrived, savings recorded as computed.
|
||||
assert _deferrable_savings_delta(500, 120) == 120
|
||||
# Usage-less turn (cancelled/failed): savings deferred.
|
||||
assert _deferrable_savings_delta(0, 120) == 0
|
||||
assert _deferrable_savings_delta(-1, 120) == 0
|
||||
# Non-positive savings pass through unchanged regardless of input, so the
|
||||
# recorded-total bookkeeping keeps its normal resync behaviour.
|
||||
assert _deferrable_savings_delta(0, 0) == 0
|
||||
assert _deferrable_savings_delta(500, 0) == 0
|
||||
assert _deferrable_savings_delta(0, -5) == -5
|
||||
assert _deferrable_savings_delta(500, -5) == -5
|
||||
|
||||
|
||||
def test_deferred_savings_land_with_next_usage_turn() -> None:
|
||||
"""Walk the recorded-total bookkeeping the handler performs per turn.
|
||||
|
||||
The handler computes ``saved_delta = _deferrable_savings_delta(input_delta,
|
||||
tokens_saved - recorded)`` and then advances ``recorded += saved_delta``.
|
||||
A deferred turn must leave the savings pending so they ride along with the
|
||||
next usage-carrying turn instead of being dropped.
|
||||
"""
|
||||
|
||||
recorded = 0
|
||||
|
||||
# Turn 1: compressed (100 saved) but cancelled before any usage frame.
|
||||
tokens_saved = 100
|
||||
delta = _deferrable_savings_delta(0, tokens_saved - recorded)
|
||||
assert delta == 0 # nothing recorded...
|
||||
recorded += delta
|
||||
assert recorded == 0 # ...and the 100 stays pending.
|
||||
|
||||
# Turn 2: compressed (50 more saved) and completed with usage.
|
||||
tokens_saved = 150
|
||||
delta = _deferrable_savings_delta(4_000, tokens_saved - recorded)
|
||||
assert delta == 150 # turn 2's 50 plus the deferred 100.
|
||||
recorded += delta
|
||||
assert recorded == tokens_saved
|
||||
|
||||
|
||||
def test_per_turn_and_residual_sites_use_the_gate() -> None:
|
||||
"""Source-level guard for the closure-internal wiring.
|
||||
|
||||
The per-turn metrics closure and the session-end residual flush both live
|
||||
inside ``handle_openai_responses_ws`` and cannot be reached by unit tests
|
||||
(see the pending-harness note in test_codex_ws_compression_scheduler.py),
|
||||
so guard the wiring in source: both sites must gate their savings delta
|
||||
through ``_deferrable_savings_delta``, and the recorded-savings total must
|
||||
advance by the recorded delta (``+= saved_delta``) — a naked
|
||||
``= tokens_saved`` assignment would silently drop deferred savings.
|
||||
"""
|
||||
|
||||
source = OPENAI_HANDLER.read_text()
|
||||
assert source.count("_deferrable_savings_delta(") >= 3, (
|
||||
"Expected the per-turn WS metrics closure and the session-end "
|
||||
"residual flush to both gate savings through "
|
||||
"_deferrable_savings_delta (plus its def). A savings delta recorded "
|
||||
"without input accounting writes a savings-with-zero-spend "
|
||||
"checkpoint into the savings tracker."
|
||||
)
|
||||
assert re.search(r"ws_recorded_tokens_saved_total\s*\+=\s*saved_delta", source), (
|
||||
"The recorded-savings total must advance by the recorded delta so "
|
||||
"savings deferred from usage-less turns stay pending for the next "
|
||||
"usage-carrying turn."
|
||||
)
|
||||
assert not re.search(r"ws_recorded_tokens_saved_total\s*=\s*tokens_saved\b", source), (
|
||||
"Naked `ws_recorded_tokens_saved_total = tokens_saved` reintroduced: "
|
||||
"this silently drops savings deferred from usage-less turns."
|
||||
)
|
||||
Loading…
Add table
Add a link
Reference in a new issue