From 15513b81f47147b5a8a52318cab30a0e7dda53bd Mon Sep 17 00:00:00 2001 From: Michael Lam Date: Fri, 15 May 2026 17:17:48 -0700 Subject: [PATCH] fix: harden OpenRouter cost snapshots --- api/providers.py | 47 +++++++++++-------- tests/test_provider_cost_history.py | 71 +++++++++++++++++++++++++++++ 2 files changed, 100 insertions(+), 18 deletions(-) diff --git a/api/providers.py b/api/providers.py index b3ddb4d1..804e112a 100644 --- a/api/providers.py +++ b/api/providers.py @@ -1445,6 +1445,7 @@ def _provider_is_oauth(provider_id: str) -> bool: _COST_SNAPSHOTS_DIR_NAME = "cost-snapshots" _COST_SNAPSHOT_MAX_DAYS = 365 # hard cap to prevent unbounded growth +_COST_SNAPSHOT_LOCK = threading.Lock() def _cost_snapshots_dir() -> Path: @@ -1558,23 +1559,28 @@ def _append_cost_snapshot(provider: str, usage: int | float | None, limit: int | repeated calls within the same day are idempotent. """ today = datetime.now(timezone.utc).strftime("%Y-%m-%d") - snapshots = _read_cost_snapshots(provider) - # Update or append today's entry - updated = False - for entry in snapshots: - if entry["date"] == today: - entry["used"] = usage - entry["limit"] = limit - updated = True - break - if not updated: - snapshots.append({"date": today, "used": usage, "limit": limit}) - snapshots.sort(key=lambda e: e["date"]) - # Cap to _COST_SNAPSHOT_MAX_DAYS entries (keep most recent) - if len(snapshots) > _COST_SNAPSHOT_MAX_DAYS: - snapshots = snapshots[-_COST_SNAPSHOT_MAX_DAYS:] - _write_cost_snapshots(provider, snapshots) - return snapshots + # Serialize the read-modify-write cycle. The atomic os.replace in + # _write_cost_snapshots protects the file write itself, but without this + # lock two concurrent requests can both read the same old snapshot list and + # race to replace it with stale data. + with _COST_SNAPSHOT_LOCK: + snapshots = _read_cost_snapshots(provider) + # Update or append today's entry + updated = False + for entry in snapshots: + if entry["date"] == today: + entry["used"] = usage + entry["limit"] = limit + updated = True + break + if not updated: + snapshots.append({"date": today, "used": usage, "limit": limit}) + snapshots.sort(key=lambda e: e["date"]) + # Cap to _COST_SNAPSHOT_MAX_DAYS entries (keep most recent) + if len(snapshots) > _COST_SNAPSHOT_MAX_DAYS: + snapshots = snapshots[-_COST_SNAPSHOT_MAX_DAYS:] + _write_cost_snapshots(provider, snapshots) + return snapshots def _compute_deltas(snapshots: list[dict[str, Any]], window_days: int) -> list[dict[str, Any]]: @@ -1583,7 +1589,10 @@ def _compute_deltas(snapshots: list[dict[str, Any]], window_days: int) -> list[d Each snapshot carries cumulative ``used``; the delta for a day is the difference between that day's cumulative value and the previous day's. The oldest day in the window has ``delta=None`` (no - previous baseline). + previous baseline). If the cumulative value drops, treat that day + as the start of a fresh series (for example after an API-key rotation) + and use the current value as that day's delta instead of emitting a + negative spend bar. """ # Take only the last *window_days* entries window = snapshots[-window_days:] if len(snapshots) > window_days else list(snapshots) @@ -1592,6 +1601,8 @@ def _compute_deltas(snapshots: list[dict[str, Any]], window_days: int) -> list[d delta = None if i > 0 and entry.get("used") is not None and window[i - 1].get("used") is not None: delta = float(entry["used"]) - float(window[i - 1]["used"]) + if delta < 0: + delta = float(entry["used"]) # Rounding: avoid -0.0 and tiny floating-point noise if abs(delta) < 1e-9: delta = 0.0 diff --git a/tests/test_provider_cost_history.py b/tests/test_provider_cost_history.py index 7ddf96c9..54bd9971 100644 --- a/tests/test_provider_cost_history.py +++ b/tests/test_provider_cost_history.py @@ -189,6 +189,77 @@ def test_openrouter_cost_history_deltas_from_cumulative(monkeypatch, tmp_path): assert snaps[2]["delta"] == 2.5 +def test_openrouter_cost_history_reset_uses_fresh_series_delta(monkeypatch, tmp_path): + """A lower cumulative value starts a fresh series instead of a negative bar.""" + monkeypatch.setattr(profiles, "get_active_hermes_home", lambda: tmp_path) + monkeypatch.delenv("OPENROUTER_API_KEY", raising=False) + (tmp_path / ".env").write_text("OPENROUTER_API_KEY=test-or-key\n", encoding="utf-8") + old_cfg, old_mtime = _with_config(model={"provider": "openrouter"}) + + import api.providers as providers + + snap_dir = tmp_path / "cost-snapshots" + snap_dir.mkdir(parents=True, exist_ok=True) + historical = { + "provider": "openrouter", + "snapshots": [ + {"date": "2030-04-13", "used": 9.0, "limit": 20}, + {"date": "2030-04-14", "used": 12.0, "limit": 20}, + ], + } + (snap_dir / "openrouter.json").write_text(json.dumps(historical), encoding="utf-8") + + def fake_urlopen(req, timeout): + # Simulate key rotation or provider reset: cumulative usage dropped. + payload = {"data": {"usage": 1.25, "limit": 20, "label": "Credits"}} + return _FakeResponse(json.dumps(payload).encode("utf-8")) + + monkeypatch.setattr(providers.urllib.request, "urlopen", fake_urlopen) + monkeypatch.setattr(providers, "datetime", type("DT", (), { + "now": staticmethod(lambda tz=None: datetime(2030, 4, 15, 12, 0, 0, tzinfo=tz or timezone.utc)), + "strftime": datetime.strftime, + })) + + try: + result = providers.get_provider_cost_history("openrouter", days=7) + finally: + _restore_config(old_cfg, old_mtime) + + assert result["ok"] is True + assert result["snapshots"][-1]["date"] == "2030-04-15" + assert result["snapshots"][-1]["used"] == 1.25 + assert result["snapshots"][-1]["delta"] == 1.25 + assert all(snap["delta"] is None or snap["delta"] >= 0 for snap in result["snapshots"]) + + +def test_cost_snapshot_append_uses_lock(monkeypatch, tmp_path): + """Snapshot append serializes the read-modify-write critical section.""" + monkeypatch.setattr(profiles, "get_active_hermes_home", lambda: tmp_path) + + import api.providers as providers + + entered = {"count": 0} + + class RecordingLock: + def __enter__(self): + entered["count"] += 1 + return self + + def __exit__(self, *exc): + return False + + monkeypatch.setattr(providers, "_COST_SNAPSHOT_LOCK", RecordingLock()) + monkeypatch.setattr(providers, "datetime", type("DT", (), { + "now": staticmethod(lambda tz=None: datetime(2030, 4, 15, 12, 0, 0, tzinfo=tz or timezone.utc)), + "strftime": datetime.strftime, + })) + + snapshots = providers._append_cost_snapshot("openrouter", 4.0, 20.0) + + assert entered["count"] == 1 + assert snapshots == [{"date": "2030-04-15", "used": 4.0, "limit": 20.0}] + + # ── Missing credentials ───────────────────────────────────────────────────────