headroomlabs-ai
diff --git a/‎headroom/proxy/savings_tracker.py‎
Lines changed: 24 additions & 13 deletions b/‎headroom/proxy/savings_tracker.py‎
Lines changed: 24 additions & 13 deletions
diff --git a/‎tests/test_proxy_savings_history.py‎
Lines changed: 145 additions & 5 deletions b/‎tests/test_proxy_savings_history.py‎
Lines changed: 145 additions & 5 deletions
@@ -34,6 +34,7 @@
 DEFAULT_MAX_HISTORY_AGE_DAYS = 365
 DEFAULT_MAX_RESPONSE_HISTORY_POINTS = 500
 DEFAULT_DISPLAY_SESSION_INACTIVITY_MINUTES = 60
+DEFAULT_FALLBACK_INPUT_COST_PER_TOKEN = 3.0 / 1_000_000
 
 LITELLM_AVAILABLE = importlib.util.find_spec("litellm") is not None
 litellm: Any | None = None
@@ -189,18 +190,20 @@ def _resolve_litellm_model(model: str) -> str:
 def _estimate_compression_savings_usd(model: str, tokens_saved: int) -> float:
     """Estimate compression savings in USD from saved input tokens."""
     litellm = _get_litellm_module()
-    if tokens_saved <= 0 or litellm is None:
+    if tokens_saved <= 0:
         return 0.0
+    if litellm is None:
+        return float(tokens_saved) * float(DEFAULT_FALLBACK_INPUT_COST_PER_TOKEN)
 
     try:
         resolved = _resolve_litellm_model(model)
         info = litellm.model_cost.get(resolved, {})
         input_cost_per_token = info.get("input_cost_per_token")
         if not input_cost_per_token:
-            return 0.0
+            raise RuntimeError("input cost unavailable")
         return float(tokens_saved) * float(input_cost_per_token)
     except Exception:
-        return 0.0
+        return float(tokens_saved) * float(DEFAULT_FALLBACK_INPUT_COST_PER_TOKEN)
 
 
 def _estimate_input_cost_usd(
@@ -220,23 +223,31 @@ def _estimate_input_cost_usd(
     cache_read = _coerce_int(cache_read_tokens)
     cache_write = _coerce_int(cache_write_tokens)
     uncached = _coerce_int(uncached_input_tokens)
-    litellm = _get_litellm_module()
-    # Gate on tokens actually sent. Providers like Anthropic report cache
-    # reads/writes separately from `input_tokens` (the uncached portion), so a
-    # fully prefix-cached request has input_tokens == 0 while cache_read > 0.
-    # Bailing on `input_tokens <= 0` alone dropped the real cache-read cost,
-    # leaving days with compression savings but zero recorded spend.
-    if total_input_tokens + cache_read + cache_write + uncached <= 0 or litellm is None:
+
+    # Prefer the breakdown when callers supply segmented token counts.
+    # Never add `input_tokens` on top of the breakdown to avoid double-counting.
+    use_breakdown = (cache_read + cache_write + uncached) > 0
+    chargeable_tokens = (
+        (cache_read + cache_write + uncached) if use_breakdown else total_input_tokens
+    )
+    if chargeable_tokens <= 0:
         return 0.0
 
+    litellm = _get_litellm_module()
+    # Keep exact provider pricing authoritative when available.
+    # `litellm` can be present but lack an entry for the resolved model,
+    # in which case we fall back to a blended rate instead of zeroing usage.
+    if litellm is None:
+        return float(chargeable_tokens) * float(DEFAULT_FALLBACK_INPUT_COST_PER_TOKEN)
+
     try:
         resolved = _resolve_litellm_model(model)
         info = litellm.model_cost.get(resolved, {})
         input_cost_per_token = info.get("input_cost_per_token")
         if not input_cost_per_token:
-            return 0.0
+            raise RuntimeError("input cost unavailable")
 
-        if cache_read + cache_write + uncached > 0:
+        if use_breakdown:
             cache_read_cost = info.get(
                 "cache_read_input_token_cost",
                 input_cost_per_token,
@@ -253,7 +264,7 @@ def _estimate_input_cost_usd(
 
         return float(total_input_tokens) * float(input_cost_per_token)
     except Exception:
-        return 0.0
+        return float(chargeable_tokens) * float(DEFAULT_FALLBACK_INPUT_COST_PER_TOKEN)
 
 
 def _normalize_history_entry(entry: Any) -> dict[str, Any] | None:
 
@@ -327,20 +327,128 @@ def fake_cost_per_token(*, model, prompt_tokens, completion_tokens):
     ) == pytest.approx(0.2)
 
     fake_litellm.model_cost = {}
-    assert savings_tracker_module._estimate_compression_savings_usd("gpt-4o", 100) == 0.0
-    assert savings_tracker_module._estimate_input_cost_usd("gpt-4o", 100) == 0.0
+    assert savings_tracker_module._estimate_compression_savings_usd("gpt-4o", 100) == pytest.approx(
+        100 * savings_tracker_module.DEFAULT_FALLBACK_INPUT_COST_PER_TOKEN
+    )
+    assert savings_tracker_module._estimate_input_cost_usd("gpt-4o", 100) == pytest.approx(
+        100 * savings_tracker_module.DEFAULT_FALLBACK_INPUT_COST_PER_TOKEN
+    )
 
     monkeypatch.setattr(
         fake_litellm,
         "cost_per_token",
         lambda **kwargs: (_ for _ in ()).throw(RuntimeError("boom")),
     )
     assert savings_tracker_module._resolve_litellm_model("mystery-model") == "mystery-model"
-    assert savings_tracker_module._estimate_compression_savings_usd("mystery-model", 100) == 0.0
+    assert savings_tracker_module._estimate_compression_savings_usd(
+        "mystery-model", 100
+    ) == pytest.approx(100 * savings_tracker_module.DEFAULT_FALLBACK_INPUT_COST_PER_TOKEN)
+    assert savings_tracker_module._estimate_input_cost_usd("mystery-model", 100) == pytest.approx(
+        100 * savings_tracker_module.DEFAULT_FALLBACK_INPUT_COST_PER_TOKEN
+    )
+    # Explicitly force the unavailable path for the whole tracker.
+    monkeypatch.setattr(savings_tracker_module, "LITELLM_AVAILABLE", False)
+    assert savings_tracker_module._estimate_compression_savings_usd("gpt-4o", 100) == pytest.approx(
+        100 * savings_tracker_module.DEFAULT_FALLBACK_INPUT_COST_PER_TOKEN
+    )
+    assert savings_tracker_module._estimate_input_cost_usd("gpt-4o", 100) == pytest.approx(
+        100 * savings_tracker_module.DEFAULT_FALLBACK_INPUT_COST_PER_TOKEN
+    )
+
+
+def test_fallback_request_pricing_stays_nonzero_with_litellm_unavailable_and_preserves_historic_zeros(
+    tmp_path, monkeypatch
+):
+    # Legacy proxy_savings rows can legitimately store zero-dollar values.
+    savings_path = tmp_path / "proxy_savings.json"
+    savings_path.write_text(
+        json.dumps(
+            {
+                "schema_version": 3,
+                "lifetime": {
+                    "requests": 1,
+                    "tokens_saved": 10,
+                    "compression_savings_usd": 0.0,
+                    "total_input_tokens": 120,
+                    "total_input_cost_usd": 0.0,
+                },
+                "display_session": {},
+                "history": [
+                    {
+                        "timestamp": "2026-03-27T09:00:00Z",
+                        "provider": "openai",
+                        "model": "gpt-4o",
+                        "total_tokens_saved": 10,
+                        "compression_savings_usd": 0.0,
+                        "total_input_tokens": 120,
+                        "total_input_cost_usd": 0.0,
+                    }
+                ],
+                "projects": {
+                    "fallback-demo": {
+                        "requests": 1,
+                        "tokens_saved": 10,
+                        "compression_savings_usd": 0.0,
+                        "total_input_tokens": 120,
+                        "total_input_cost_usd": 0.0,
+                        "last_activity_at": "2026-03-27T09:00:00Z",
+                    }
+                },
+            }
+        ),
+        encoding="utf-8",
+    )
+
+    tracker = SavingsTracker(path=str(savings_path))
+    initial_snapshot = tracker.snapshot()
+    assert initial_snapshot["lifetime"]["compression_savings_usd"] == 0.0
+    assert initial_snapshot["display_session"]["compression_savings_usd"] == 0.0
+    assert initial_snapshot["projects"]["fallback-demo"]["compression_savings_usd"] == 0.0
+    assert initial_snapshot["history"][-1]["compression_savings_usd"] == 0.0
 
     monkeypatch.setattr(savings_tracker_module, "LITELLM_AVAILABLE", False)
-    assert savings_tracker_module._estimate_compression_savings_usd("gpt-4o", 100) == 0.0
-    assert savings_tracker_module._estimate_input_cost_usd("gpt-4o", 100) == 0.0
+    monkeypatch.setattr(savings_tracker_module, "litellm", None)
+    assert tracker.record_request(
+        model="gpt-4o",
+        input_tokens=100,
+        tokens_saved=50,
+        project="fallback-demo",
+        timestamp="2026-03-27T09:10:00Z",
+    )
+    monkeypatch.setattr(
+        savings_tracker_module,
+        "_utc_now",
+        lambda: datetime(2026, 3, 27, 9, 10, 30, tzinfo=timezone.utc),
+    )
+
+    snapshot = tracker.snapshot()
+    expected_savings_fallback = 50 * savings_tracker_module.DEFAULT_FALLBACK_INPUT_COST_PER_TOKEN
+    expected_input_fallback = 100 * savings_tracker_module.DEFAULT_FALLBACK_INPUT_COST_PER_TOKEN
+    assert snapshot["lifetime"]["compression_savings_usd"] == pytest.approx(
+        expected_savings_fallback
+    )
+    assert snapshot["lifetime"]["total_input_cost_usd"] == pytest.approx(expected_input_fallback)
+    assert snapshot["display_session"]["compression_savings_usd"] == pytest.approx(
+        expected_savings_fallback
+    )
+    assert snapshot["display_session"]["total_input_cost_usd"] == pytest.approx(
+        expected_input_fallback
+    )
+    assert snapshot["projects"]["fallback-demo"]["compression_savings_usd"] == pytest.approx(
+        expected_savings_fallback
+    )
+    assert snapshot["projects"]["fallback-demo"]["total_input_cost_usd"] == pytest.approx(
+        expected_input_fallback
+    )
+    assert snapshot["history"][-1]["compression_savings_usd"] == pytest.approx(
+        expected_savings_fallback
+    )
+
+    persisted = json.loads(savings_path.read_text(encoding="utf-8"))
+    assert persisted["history"][0]["compression_savings_usd"] == 0.0
+    assert persisted["history"][-1]["compression_savings_usd"] == pytest.approx(
+        expected_savings_fallback
+    )
 
 
 def test_input_cost_counts_cache_reads_when_uncached_input_is_zero(monkeypatch):
@@ -374,6 +482,38 @@ def fake_cost_per_token(*, model, prompt_tokens, completion_tokens):
     assert cost == pytest.approx(0.3)
 
 
+def test_fallback_input_cost_uses_breakdown_sum_not_input_tokens_when_litellm_unavailable(
+    monkeypatch,
+):
+    # Regression: when both `input_tokens` and a nonzero cache breakdown are
+    # present and LiteLLM is unavailable, the fallback must price only the
+    # breakdown sum — never input_tokens + breakdown_sum — to avoid
+    # double-counting the tokens that the breakdown already covers.
+    monkeypatch.setattr(savings_tracker_module, "LITELLM_AVAILABLE", False)
+    monkeypatch.setattr(savings_tracker_module, "litellm", None)
+
+    input_tokens = 1000
+    cache_read = 200
+    cache_write = 100
+    uncached = 300
+    breakdown_sum = cache_read + cache_write + uncached  # 600
+
+    result = savings_tracker_module._estimate_input_cost_usd(
+        "gpt-4o",
+        input_tokens,
+        cache_read_tokens=cache_read,
+        cache_write_tokens=cache_write,
+        uncached_input_tokens=uncached,
+    )
+
+    fallback_rate = savings_tracker_module.DEFAULT_FALLBACK_INPUT_COST_PER_TOKEN
+    expected = breakdown_sum * fallback_rate
+    double_counted = (input_tokens + breakdown_sum) * fallback_rate
+
+    assert result == pytest.approx(expected)
+    assert result != pytest.approx(double_counted)
+
+
 def test_display_session_rolls_after_inactivity_and_counts_zero_savings_requests(
     tmp_path, monkeypatch
 ):