diff --git a/headroom/proxy/prometheus_metrics.py b/headroom/proxy/prometheus_metrics.py index 889ed60be..8cd83846b 100644 --- a/headroom/proxy/prometheus_metrics.py +++ b/headroom/proxy/prometheus_metrics.py @@ -869,14 +869,23 @@ class PrometheusMetrics: # would hold the lock for the whole write instead of just the syscall. # ponytail: default thread pool, not a dedicated executor -- give it one # if a profile ever shows writers parked on flock saturating the pool. - if tokens_saved > 0 and not self._stateless: + # tool_search_deferral saves tool-SCHEMA tokens that never move the + # message-level tok_before/after, so a tool-heavy turn can have + # tokens_saved=0 while genuinely deferring thousands of tokens. Fold that + # component into the ledger delta the same way the PERF headline and + # perf/analyzer do (`headline_before = before + tool_saved`); otherwise + # `headroom savings` understates real compression 7-10x on tool-search + # sessions and drops deferral-only turns from the ledger entirely (#2795). + deferral_saved = max(0, int(tool_search_saved)) + ledger_saved = tokens_saved + deferral_saved + if ledger_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 + # understated the original by `ledger_saved`, inflating that # percentage (e.g. a real 40% reduction was reported as ~67%). # Reconstruct the original as forwarded + saved. await asyncio.to_thread( @@ -887,7 +896,7 @@ class PrometheusMetrics: # `tokens_saved` yields a mixed-ruler before/after (local 10->6 # with the provider reporting 8 would record 12->8). Use the # caller's local count when supplied. - tokens_before=ledger_input_tokens + tokens_saved, + tokens_before=ledger_input_tokens + ledger_saved, tokens_after=ledger_input_tokens, model=model, client=client or "proxy", diff --git a/tests/test_savings_ledger_before_forwarded.py b/tests/test_savings_ledger_before_forwarded.py index 52a10a133..72d158377 100644 --- a/tests/test_savings_ledger_before_forwarded.py +++ b/tests/test_savings_ledger_before_forwarded.py @@ -63,3 +63,69 @@ async def test_record_savings_event_uses_original_input_as_before( "source": "proxy", } ] + + +def _capture_ledger(monkeypatch: pytest.MonkeyPatch) -> list[dict[str, Any]]: + calls: list[dict[str, Any]] = [] + monkeypatch.setattr( + prometheus_metrics.savings_ledger, + "record_savings_event", + lambda **kwargs: calls.append(kwargs), + ) + return calls + + +@pytest.mark.asyncio +async def test_record_savings_event_includes_tool_search_deferral( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """tool_search_deferral savings must ride into the ledger delta so + `headroom savings` does not undercount tool-search sessions ~7-10x (#2795).""" + calls = _capture_ledger(monkeypatch) + metrics = prometheus_metrics.PrometheusMetrics( + savings_tracker=_FakeSavingsTracker(), + otel_metrics=_FakeOtelMetrics(), + ) + await metrics.record_request( + provider="anthropic", + model="claude-opus-4-6", + input_tokens=109844, # forwarded (post-compression) message count + output_tokens=25, + tokens_saved=1896, + tool_search_saved=13182, # deferred tool schemas never sent + latency_ms=10.0, + client="claude-code", + ) + + assert len(calls) == 1 + # saved = tokens_saved + tool_search_saved; before = forwarded + saved. + assert calls[0]["tokens_after"] == 109844 + assert calls[0]["tokens_before"] == 109844 + 1896 + 13182 # == 124922 + + +@pytest.mark.asyncio +async def test_record_savings_event_written_for_deferral_only_turn( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A tool-heavy turn can defer thousands of tool-schema tokens while + tokens_saved is 0 (deferral does not move the message-level count). It must + still be recorded, not dropped from the ledger (#2795).""" + calls = _capture_ledger(monkeypatch) + metrics = prometheus_metrics.PrometheusMetrics( + savings_tracker=_FakeSavingsTracker(), + otel_metrics=_FakeOtelMetrics(), + ) + await metrics.record_request( + provider="anthropic", + model="claude-opus-4-6", + input_tokens=50000, + output_tokens=25, + tokens_saved=0, + tool_search_saved=13182, + latency_ms=10.0, + client="claude-code", + ) + + assert len(calls) == 1 + assert calls[0]["tokens_before"] == 50000 + 13182 + assert calls[0]["tokens_after"] == 50000