From 12a38d31808283c935a4d15f297376d2b159e567 Mon Sep 17 00:00:00 2001 From: GUOHAO LIU <94768569+lennney@users.noreply.github.com> Date: Mon, 13 Jul 2026 21:37:12 +0800 Subject: [PATCH] feat(proxy): persist per-model savings breakdown in proxy_savings.json (#2055) ## Description Persist per-model savings breakdown in `proxy_savings.json` so per-model stats survive proxy restarts (Closes #1913). Previously only the in-memory Prometheus metrics kept per-model data, which reset on restart. Add `by_model` dict keyed by normalized model name, each entry following the lifetime aggregate shape with a derived `savings_percent`. ## Type of Change - [ ] Bug fix (non-breaking change that fixes an issue) - [x] New feature (non-breaking change that adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [ ] Performance improvement - [ ] Code refactoring (no functional changes) ## Changes Made - `headroom/proxy/savings_tracker.py`: add `_empty_by_model_entry()` and `_normalize_by_model()` helpers; add `_record_by_model_locked()` and `_by_model_snapshot_locked()` methods to SavingsTracker; include `by_model` in `_default_state()`, `_sanitize_state()`, `snapshot()`, `stats_preview()`, and `history_response()`; update `record_request()` and `record_compression_savings()` to accumulate per-model counters ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check`) - [x] Manual testing performed ### Test Output ```text $ uv run pytest tests/test_proxy_savings_history.py -x -q 39 passed in 11.26s $ uv run ruff check headroom/proxy/savings_tracker.py All checks passed! ``` ## Real Behavior Proof - Environment: Linux, headroom main @ 868b88bc - Exact command / steps: (1) apply patch, (2) `uv run pytest tests/test_proxy_savings_history.py -x -q`, (3) `uv run ruff check headroom/proxy/savings_tracker.py` - Observed result: All 39 tests pass, ruff clean, Python AST parse OK - Not tested: End-to-end with live proxy serving /stats and /stats-history to verify by_model appears in API response with correct per-model data ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review --------- Co-authored-by: lennney Co-authored-by: JerrettDavis Co-authored-by: Tejas Chopra --- headroom/proxy/savings_tracker.py | 100 ++++++++++++++++++++++++++++ tests/test_proxy_savings_history.py | 31 +++++++++ 2 files changed, 131 insertions(+) diff --git a/headroom/proxy/savings_tracker.py b/headroom/proxy/savings_tracker.py index 560cfad36..0b45d9ccc 100644 --- a/headroom/proxy/savings_tracker.py +++ b/headroom/proxy/savings_tracker.py @@ -372,6 +372,16 @@ def _empty_display_session() -> dict[str, Any]: } +def _empty_by_model_entry() -> dict[str, Any]: + return { + "requests": 0, + "tokens_saved": 0, + "compression_savings_usd": 0.0, + "total_input_tokens": 0, + "total_input_cost_usd": 0.0, + } + + def _empty_project_entry() -> dict[str, Any]: return { "requests": 0, @@ -416,6 +426,28 @@ def _normalize_projects(raw: Any) -> dict[str, dict[str, Any]]: return projects +def _normalize_by_model(raw: Any) -> dict[str, dict[str, Any]]: + if not isinstance(raw, dict): + return {} + result: dict[str, dict[str, Any]] = {} + for model_name, entry in raw.items(): + model = _normalize_model(model_name) + if not isinstance(entry, dict): + continue + normalized = _empty_by_model_entry() + normalized["requests"] = _coerce_int(entry.get("requests")) + normalized["tokens_saved"] = _coerce_int(entry.get("tokens_saved")) + normalized["compression_savings_usd"] = round( + _coerce_float(entry.get("compression_savings_usd")), 6 + ) + normalized["total_input_tokens"] = _coerce_int(entry.get("total_input_tokens")) + normalized["total_input_cost_usd"] = round( + _coerce_float(entry.get("total_input_cost_usd")), 6 + ) + result[model] = normalized + return result + + def _normalize_display_session(entry: Any) -> dict[str, Any]: if not isinstance(entry, dict): return _empty_display_session() @@ -553,6 +585,12 @@ class SavingsTracker: 6, ) + self._record_by_model_locked( + model, + tokens_saved_delta=delta_tokens, + savings_usd_delta=delta_usd, + ) + self._state["history"].append( { "timestamp": _to_utc_iso(timestamp_dt), @@ -687,6 +725,15 @@ class SavingsTracker: if session.get("started_at") is None: session["started_at"] = session["last_activity_at"] + self._record_by_model_locked( + model, + requests_delta=1, + tokens_saved_delta=delta_tokens_saved, + savings_usd_delta=delta_savings_usd, + input_tokens_delta=delta_input_tokens, + input_cost_usd_delta=delta_input_cost_usd, + ) + self._record_project_locked( project, timestamp_dt=timestamp_dt, @@ -756,6 +803,34 @@ class SavingsTracker: ) del projects[evict] + def _record_by_model_locked( + self, + model: str, + *, + requests_delta: int = 0, + tokens_saved_delta: int = 0, + savings_usd_delta: float = 0.0, + input_tokens_delta: int = 0, + input_cost_usd_delta: float = 0.0, + ) -> None: + """Accumulate per-model savings. Caller must hold ``self._lock``. + + Lazy-inits ``by_model`` so existing state files without the key work + without migration. + """ + by_model: dict[str, dict[str, Any]] = self._state.setdefault("by_model", {}) + key = _normalize_model(model) + entry = by_model.setdefault(key, _empty_by_model_entry()) + entry["requests"] += max(requests_delta, 0) + entry["tokens_saved"] += max(tokens_saved_delta, 0) + entry["compression_savings_usd"] = round( + entry["compression_savings_usd"] + max(savings_usd_delta, 0.0), 6 + ) + entry["total_input_tokens"] += max(input_tokens_delta, 0) + entry["total_input_cost_usd"] = round( + entry["total_input_cost_usd"] + max(input_cost_usd_delta, 0.0), 6 + ) + def _projects_snapshot_locked(self) -> dict[str, dict[str, Any]]: """Per-project stats with a derived ``savings_percent``, sorted by savings.""" projects = self._state.get("projects", {}) @@ -775,6 +850,25 @@ class SavingsTracker: result[name] = view return result + def _by_model_snapshot_locked(self) -> dict[str, dict[str, Any]]: + """Per-model stats ranked by savings.""" + by_model = self._state.get("by_model", {}) + ranked = sorted( + by_model.items(), + key=lambda item: item[1]["tokens_saved"], + reverse=True, + ) + result: dict[str, dict[str, Any]] = {} + for model, entry in ranked: + view = dict(entry) + total_before = entry["tokens_saved"] + entry["total_input_tokens"] + view["savings_percent"] = round( + (entry["tokens_saved"] / total_before * 100) if total_before > 0 else 0.0, + 2, + ) + result[model] = view + return result + def stats_preview(self, recent_points: int = 20) -> dict[str, Any]: """Return a compact preview for `/stats`.""" snapshot = self.snapshot() @@ -789,6 +883,7 @@ class SavingsTracker: "retention": snapshot["retention"], "projects": snapshot["projects"], "projects_limit": DEFAULT_MAX_PROJECTS, + "by_model": snapshot["by_model"], } def history_response(self, history_mode: str = "compact") -> dict[str, Any]: @@ -818,6 +913,7 @@ class SavingsTracker: }, "retention": snapshot["retention"], "projects": snapshot["projects"], + "by_model": snapshot["by_model"], "history_summary": { "mode": history_mode, "stored_points": len(raw_history), @@ -882,6 +978,7 @@ class SavingsTracker: "max_response_history_points": self._max_response_history_points, }, "projects": self._projects_snapshot_locked(), + "by_model": self._by_model_snapshot_locked(), } def _default_state(self) -> dict[str, Any]: @@ -899,6 +996,7 @@ class SavingsTracker: "display_session": _empty_display_session(), "history": [], "projects": {}, + "by_model": {}, } def _load_state(self) -> dict[str, Any]: @@ -978,6 +1076,7 @@ class SavingsTracker: "display_session": _normalize_display_session(raw.get("display_session")), "history": normalized_history, "projects": _normalize_projects(raw.get("projects")), + "by_model": _normalize_by_model(raw.get("by_model")), } if normalized_history: @@ -1093,6 +1192,7 @@ class SavingsTracker: "display_session": self._state["display_session"], "history": self._state["history"], "projects": self._state.get("projects", {}), + "by_model": self._state.get("by_model", {}), } json_data = json.dumps(payload, indent=2) diff --git a/tests/test_proxy_savings_history.py b/tests/test_proxy_savings_history.py index 117fd1ee6..2d1bcdbc0 100644 --- a/tests/test_proxy_savings_history.py +++ b/tests/test_proxy_savings_history.py @@ -1577,6 +1577,37 @@ def test_cache_read_savings_accumulate_and_survive_restart(tmp_path, monkeypatch assert reloaded.history_response()["lifetime"]["cache_read_tokens"] == 1_600_000 +def test_by_model_savings_accumulate_and_survive_restart(tmp_path): + path = tmp_path / "proxy_savings.json" + tracker = SavingsTracker(path=str(path)) + + tracker.record_request( + model="gpt-4o", + input_tokens=100, + tokens_saved=40, + timestamp="2026-07-01T09:00:00Z", + ) + tracker.record_request( + model="claude-sonnet-4-6", + input_tokens=300, + tokens_saved=60, + timestamp="2026-07-01T09:01:00Z", + ) + + persisted = json.loads(path.read_text(encoding="utf-8")) + assert set(persisted["by_model"]) == {"gpt-4o", "claude-sonnet-4-6"} + assert persisted["by_model"]["gpt-4o"]["tokens_saved"] == 40 + assert persisted["by_model"]["gpt-4o"]["total_input_tokens"] == 100 + + reloaded = SavingsTracker(path=str(path)) + stats_by_model = reloaded.stats_preview()["by_model"] + assert set(stats_by_model) == {"gpt-4o", "claude-sonnet-4-6"} + assert stats_by_model["claude-sonnet-4-6"]["tokens_saved"] == 60 + assert stats_by_model["claude-sonnet-4-6"]["total_input_tokens"] == 300 + assert stats_by_model["claude-sonnet-4-6"]["savings_percent"] == 16.67 + assert reloaded.history_response()["by_model"] == stats_by_model + + def test_v3_state_without_cache_fields_loads_clean_and_saves_v4(tmp_path): path = tmp_path / "proxy_savings.json" path.write_text(