From 02eb90f24318abdfb05438e873c8f2af7023ab91 Mon Sep 17 00:00:00 2001 From: Tejas Chopra Date: Fri, 17 Jul 2026 21:58:18 -0700 Subject: [PATCH] feat(metrics): record per-extension token savings (#2371) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## 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. --- headroom/proxy/prometheus_metrics.py | 26 ++++++++++++++ headroom/proxy/server.py | 1 + tests/test_compression_observability.py | 45 +++++++++++++++++++++++++ 3 files changed, 72 insertions(+) diff --git a/headroom/proxy/prometheus_metrics.py b/headroom/proxy/prometheus_metrics.py index 97006376c..ef423c22a 100644 --- a/headroom/proxy/prometheus_metrics.py +++ b/headroom/proxy/prometheus_metrics.py @@ -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``. diff --git a/headroom/proxy/server.py b/headroom/proxy/server.py index abecfe245..0226a7a79 100644 --- a/headroom/proxy/server.py +++ b/headroom/proxy/server.py @@ -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, diff --git a/tests/test_compression_observability.py b/tests/test_compression_observability.py index 3491a6ceb..d7c66d4dc 100644 --- a/tests/test_compression_observability.py +++ b/tests/test_compression_observability.py @@ -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