feat(proxy/savings): aggregate tool-schema savings into Metrics + all reporting sinks (#2546)

## Description

Companion to #2545 (the "sources" double-count fix) — this fixes the
"sinks" half found in the same savings audit: **tool-schema / deferral
savings were never aggregated into `Metrics`**. They lived only in
per-request log tags, so every sink that reads `metrics.*` silently
dropped them, and one CLI mode disagreed with another.

Confirmed sinks that under-reported:
- **Session-summary printout** — `Tokens saved:` is message-only; a
24K-tool-deferral turn printed `0`.
- **`cost.py` session summary** (feeds `/stats.summary`) —
`total_tokens_saved_with_rtk` etc. were message+CLI only.
- **`/stats` `all_layers_tokens_saved`** — the advertised "total"
excluded the `tool_search` layer it enumerates in `by_layer`.
- **`headroom perf --format json/csv`** — omitted `tool_saved` while the
**text** output of the same command showed it.

Closes #

## Type of Change
- [x] Bug fix (non-breaking) / observability correctness

## Changes Made
- `PrometheusMetrics.tool_search_saved_total` — new counter, accumulated
in `record_request` from a new `tool_search_saved` arg;
`emit_request_outcome` fills it from the `tool_search_deferred_tokens` +
`turn_hook_tools_saved_tokens` tags. **One source of truth.**
- Fed into: session summary (`Tool schemas deferred:` line), `cost.py`
summary (new `tool_schema_tokens_saved` +
`total_tokens_saved_all_layers`; existing fields unchanged for
back-compat), `/stats` `all_layers` total, and `build_perf_summary`
(`tool_saved`).
- Kept **distinct** from `tokens_saved_total` (message compression) —
tool bytes never move `tok_before/after`, so it's a separate layer, not
a merge (no double-count).

## Testing
- [x] `ruff` + `ruff format --check` + `mypy` clean
- [x] Regression tests + existing suites pass

### Test Output
```text
pytest tests/test_savings_tool_search_aggregation.py tests/test_cli_perf_format.py -q → 18 passed
pytest tests/test_cli_perf_format.py test_proxy_savings_history.py test_dashboard_token_savings.py
      test_bundled_tools_savings.py test_openai_chat_turn_hooks.py → 68 passed, 2 skipped
mypy (metrics/outcome/cost/analyzer) → clean
```

## Real Behavior Proof
- Standalone: `record_request(tool_search_saved=1500)` then `(…=800)` →
`metrics.tool_search_saved_total == 2300`, `tokens_saved_total == 200`
(message stays separate); `build_perf_summary` over records with
`tool_saved` 5000+3000 → `tool_saved == 8000`.

## Checklist
- [x] Self-reviewed; no new warnings; tests pass; did **not** edit
`CHANGELOG.md`

## Additional Notes
Together, #2545 (record once) + this (surface every layer) make savings
correct **and** complete end-to-end across `/stats`, the dashboard,
`headroom perf`, the session summary, and cost/budget. The `/stats`
`by_layer.tool_search` and dashboard card already showed the layer
(windowed, from the log scan); this makes the lifetime/metrics-based
sinks agree.
This commit is contained in:
Tejas Chopra 2026-07-24 20:40:46 -07:00 committed by GitHub
parent 0845b26ee6
commit 9f1ffefe83
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 104 additions and 1 deletions

View file

@ -1038,6 +1038,7 @@ def build_perf_summary(report: PerfReport) -> dict:
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)
total_cr = sum(r.cache_read for r in records)
total_cw = sum(r.cache_write for r in records)
@ -1094,6 +1095,7 @@ def build_perf_summary(report: PerfReport) -> dict:
"total_tokens_before": total_before,
"total_tokens_after": total_after,
"tokens_saved": total_saved,
"tool_saved": total_tool_saved,
"savings_pct": _pct(total_saved, total_before),
"cache_read_tokens": total_cr,
"cache_write_tokens": total_cw,

View file

@ -593,6 +593,15 @@ def build_session_summary(
"rtk_tokens_avoided": cli_tokens_avoided,
"total_tokens_saved_with_rtk": metrics.tokens_saved_total + cli_tokens_avoided,
"total_tokens_before_with_rtk": total_tokens_before,
# Tool-schema deferral / turn-hook tool shrink, tracked apart from
# message compression. New fields (existing ones stay message+CLI only
# for backward compat) so consumers can see the full picture.
"tool_schema_tokens_saved": getattr(metrics, "tool_search_saved_total", 0),
"total_tokens_saved_all_layers": (
metrics.tokens_saved_total
+ cli_tokens_avoided
+ getattr(metrics, "tool_search_saved_total", 0)
),
},
"uncompressed_requests": {k: v for k, v in uncompressed_reasons.items() if v > 0},
"cost": {

View file

@ -388,6 +388,14 @@ async def emit_request_outcome(handler: Any, outcome: RequestOutcome) -> None:
# HTTP middleware / WS accept captured from ``X-Headroom-Project``.
project = outcome.project or get_current_project()
# Tool-schema savings (deferral + turn-hook tool shrink) live in per-request
# tags and never move tok_before/after; aggregate them into Metrics so the
# session summary / cost summary / all-layers total can surface the layer.
_otags = outcome.tags or {}
tool_search_saved = int(_otags.get("tool_search_deferred_tokens", 0) or 0) + int(
_otags.get("turn_hook_tools_saved_tokens", 0) or 0
)
# 1. Prometheus / SavingsTracker.
await handler.metrics.record_request(
provider=outcome.provider,
@ -410,6 +418,7 @@ async def emit_request_outcome(handler: Any, outcome: RequestOutcome) -> None:
output_tokens_saved=output_tokens_saved_est,
project=project,
client=outcome.client,
tool_search_saved=tool_search_saved,
)
# 2. Cost tracker (optional).

View file

@ -97,6 +97,11 @@ class PrometheusMetrics:
self.tokens_input_total = 0
self.tokens_output_total = 0
self.tokens_saved_total = 0
# Tool-schema savings (deferral + turn-hook tool shrink), aggregated from
# per-request tags. Tracked apart from tokens_saved_total (which is message
# compression only — tool bytes never move tok_before/after) so every sink
# can surface the tool-schema layer instead of silently dropping it.
self.tool_search_saved_total = 0
# Sum of tokens we actually attempted to compress across the
# session: extracted units that passed all gates + tool-schema
# tokens we ran compaction against. Excludes prefix-frozen
@ -324,6 +329,7 @@ class PrometheusMetrics:
self.tokens_input_total = 0
self.tokens_output_total = 0
self.tokens_saved_total = 0
self.tool_search_saved_total = 0
self.attempted_input_tokens_total = 0
self.compressions_by_strategy.clear()
@ -683,6 +689,7 @@ class PrometheusMetrics:
output_tokens_saved: int = 0,
project: str | None = None,
client: str | None = None,
tool_search_saved: int = 0,
):
"""Record metrics for a request."""
# Post-guard invariant (all providers): Headroom never forwards a request
@ -709,6 +716,7 @@ class PrometheusMetrics:
self.tokens_input_total += input_tokens
self.tokens_output_total += output_tokens
self.tokens_saved_total += tokens_saved
self.tool_search_saved_total += max(0, int(tool_search_saved))
# See the attribute definition for why this is the right
# denominator for the active-compression ratio.
self.attempted_input_tokens_total += max(0, int(attempted_input_tokens))

View file

@ -1897,6 +1897,10 @@ class HeadroomProxy(
logger.info(f"Input tokens: {m.tokens_input_total:,}")
logger.info(f"Output tokens: {m.tokens_output_total:,}")
logger.info(f"Tokens saved: {m.tokens_saved_total:,}")
if m.tool_search_saved_total > 0:
# Tool-schema deferral / turn-hook tool shrink — counted apart from
# message compression (tool bytes never move tok_before/after).
logger.info(f"Tool schemas deferred: {m.tool_search_saved_total:,}")
# Active-compression ratio: savings as a fraction of what we
# *attempted* to compress (extracted units + tool schema),
# NOT the whole request. The full-request denominator is
@ -3647,7 +3651,11 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
# compression and the configured context tool both remove tokens before
# they reach model context, so dashboard-facing savings combines them.
proxy_compression_tokens = m.tokens_saved_total
all_layers_tokens_saved = proxy_compression_tokens + cli_tokens_avoided
# "All layers" must include tool-schema deferral (the tool_search layer
# enumerated in by_layer below) — otherwise the advertised total omits it.
all_layers_tokens_saved = (
proxy_compression_tokens + cli_tokens_avoided + m.tool_search_saved_total
)
total_tokens_before = m.tokens_input_total + all_layers_tokens_saved
proxy_total_before_compression = m.tokens_input_total + proxy_compression_tokens
# `attempted_input_tokens` is the compressible-only denominator

View file

@ -0,0 +1,67 @@
"""Tool-search / deferral savings must aggregate into Metrics and surface in the
reporting sinks not live only in per-request tags (which every sink reading
metrics.* structurally missed: session summary, cost summary, all-layers total,
`headroom perf --json`)."""
from __future__ import annotations
import asyncio
from headroom.perf.analyzer import PerfRecord, PerfReport, build_perf_summary
from headroom.proxy.prometheus_metrics import PrometheusMetrics
def test_metrics_accumulates_tool_search_saved_apart_from_message() -> None:
m = PrometheusMetrics()
async def go() -> None:
await m.record_request(
provider="anthropic",
model="claude-x",
input_tokens=100,
output_tokens=10,
tokens_saved=0,
latency_ms=1.0,
tool_search_saved=1500,
)
await m.record_request(
provider="anthropic",
model="claude-x",
input_tokens=100,
output_tokens=10,
tokens_saved=200,
latency_ms=1.0,
tool_search_saved=800,
)
asyncio.run(go())
assert m.tokens_saved_total == 200 # message compression only
assert m.tool_search_saved_total == 2300 # tool-schema layer, aggregated
def test_build_perf_summary_includes_tool_saved() -> None:
report = PerfReport(
perf_records=[
PerfRecord(
timestamp="t",
request_id="r1",
model="m",
tokens_before=1000,
tokens_after=900,
tokens_saved=100,
tool_saved=5000,
),
PerfRecord(
timestamp="t",
request_id="r2",
model="m",
tokens_before=500,
tokens_after=500,
tokens_saved=0,
tool_saved=3000,
),
]
)
summary = build_perf_summary(report)
assert summary["tokens_saved"] == 100 # message
assert summary["tool_saved"] == 8000 # tool-schema surfaced in json/csv sink