feat(metrics): record per-extension token savings (#2371)

## What
Adds `PrometheusMetrics.record_extension_savings(key, saved)` so proxy
extensions can report the tokens they save, and surfaces the
per-extension totals in the `/stats` payload.

## Why
Proxy extensions that perform their own token reduction currently have
no supported way to report their savings to the metrics object — there
is no method for it, so that telemetry is silently dropped. This adds
the recording method and exposes the aggregate alongside the existing
per-strategy compression breakdown.

## How
- New `extension_savings: dict[str, int]` counter on
`PrometheusMetrics`, populated lazily per extension-supplied `key` (no
hardcoded list of extensions).
- `record_extension_savings(key, saved)` accumulates positive savings
per key, mirroring how `record_compression` aggregates
`tokens_saved_by_strategy` (lock-free `defaultdict(int)`, atomic under
the GIL for these key types); non-positive values are ignored.
- Cleared in `reset_runtime()` with the other in-memory counters.
- Surfaced in `/stats` as `extension_savings`, next to
`compressions_by_strategy` / `tokens_saved_by_strategy`. No new
Prometheus series.

## Behavior change
None to existing metrics.

## Testing
- Two focused tests in `tests/test_compression_observability.py`
(per-key accumulation incl. zero/negative ignored; surfaced in `/stats`
via `create_app`) → 2 passed (13 in file).
- `ruff check` / `ruff format` → clean; `mypy` → clean on changed
source.
This commit is contained in:
Tejas Chopra 2026-07-17 21:58:18 -07:00 committed by GitHub
parent a02073e332
commit 02eb90f243
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 72 additions and 0 deletions

View file

@ -117,6 +117,13 @@ class PrometheusMetrics:
self.compressions_by_strategy: dict[str, int] = defaultdict(int)
self.tokens_saved_by_strategy: dict[str, int] = defaultdict(int)
# Per-extension token savings, keyed by the extension-supplied
# ``key``. Populated lazily by ``record_extension_savings`` — no
# hardcoded list of extensions. Proxy extensions report the tokens
# they saved so per-extension contribution is observable via /stats,
# mirroring the per-strategy compression breakdown above.
self.extension_savings: dict[str, int] = defaultdict(int)
# Fail-open compression failures, keyed by reason ("timeout",
# "error"). The proxy fails open on any optimization error so the
# request still succeeds; without this counter the failure is only
@ -321,6 +328,7 @@ class PrometheusMetrics:
self.compressions_by_strategy.clear()
self.tokens_saved_by_strategy.clear()
self.extension_savings.clear()
with self._obs_counter_lock:
self.compression_failed_by_reason.clear()
self.kompress_size_gate_by_outcome.clear()
@ -483,6 +491,24 @@ class PrometheusMetrics:
if saved > 0:
self.tokens_saved_by_strategy[strategy] += saved
def record_extension_savings(self, key: str, saved: int) -> None:
"""Accumulate tokens saved by a proxy extension, keyed by ``key``.
Called by proxy extensions that perform their own token
reduction and want that contribution surfaced alongside the
built-in compression metrics. The per-extension totals are
exposed via /stats (``extension_savings``), mirroring how
``record_compression`` accumulates ``tokens_saved_by_strategy``.
Synchronous + lock-free: ``defaultdict(int)`` writes are atomic
under the GIL for these key types, matching ``record_compression``.
Non-positive ``saved`` values are ignored the metric never
records "negative savings".
"""
if saved > 0:
self.extension_savings[key] += saved
def record_compression_failed(self, reason: str) -> None:
"""Record one fail-open compression failure, bucketed by ``reason``.

View file

@ -3817,6 +3817,7 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
else {},
"compressions_by_strategy": dict(m.compressions_by_strategy),
"tokens_saved_by_strategy": dict(m.tokens_saved_by_strategy),
"extension_savings": dict(m.extension_savings),
"codex_ws": {
"units_total": m.codex_ws_units_total,
"units_modified_total": m.codex_ws_units_modified_total,

View file

@ -33,6 +33,7 @@ Coverage:
from __future__ import annotations
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any
import pytest
@ -273,6 +274,50 @@ def test_prometheus_metrics_accumulates_per_strategy_counters():
assert m.tokens_saved_by_strategy == {"smart_crusher": 210}
def test_prometheus_metrics_accumulates_extension_savings_per_key() -> None:
from headroom.proxy.prometheus_metrics import PrometheusMetrics
m = PrometheusMetrics()
m.record_extension_savings("tool_router", 120)
m.record_extension_savings("tool_router", 30)
m.record_extension_savings("skill_search", 45)
m.record_extension_savings("skill_search", 0) # no savings, ignored
m.record_extension_savings("noop_ext", -10) # negative, ignored
# Savings accumulate per key; non-positive values never create or
# bump an entry.
assert m.extension_savings == {"tool_router": 150, "skill_search": 45}
def test_extension_savings_surface_in_stats(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
from fastapi.testclient import TestClient
from headroom.proxy.server import ProxyConfig, create_app
monkeypatch.setenv("HEADROOM_SAVINGS_PATH", str(tmp_path / "proxy_savings.json"))
config = ProxyConfig(
cache_enabled=False,
rate_limit_enabled=False,
log_requests=False,
)
app = create_app(config)
with TestClient(app) as client:
proxy = app.state.proxy
proxy.metrics.record_extension_savings("tool_router", 200)
proxy.metrics.record_extension_savings("tool_router", 50)
proxy.metrics.record_extension_savings("skill_search", 75)
stats = client.get("/stats")
assert stats.status_code == 200
assert stats.json()["extension_savings"] == {
"tool_router": 250,
"skill_search": 75,
}
def test_prometheus_metrics_accumulates_codex_ws_unit_and_frame_counters():
from headroom.proxy.prometheus_metrics import PrometheusMetrics