diff --git a/headroom/proxy/handlers/openai.py b/headroom/proxy/handlers/openai.py index 8dff6c160..8baa8610f 100644 --- a/headroom/proxy/handlers/openai.py +++ b/headroom/proxy/handlers/openai.py @@ -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, diff --git a/tests/test_codex_ws_savings_deferral.py b/tests/test_codex_ws_savings_deferral.py new file mode 100644 index 000000000..6f067d3d6 --- /dev/null +++ b/tests/test_codex_ws_savings_deferral.py @@ -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." + )