akadmin/headroom
1
0
Fork 0
mirror of https://github.com/headroomlabs-ai/headroom.git synced 2026-08-27 14:17:10 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions
headroom/tests/test_proxy_cache_ttl_metrics.py
Abhay Singh b4f807f21a
fix(proxy/cost): price cache savings by most-used model, not first-seen (#2023) 
## Description

`build_prefix_cache_stats` (`headroom/proxy/cost.py`) values each
provider's cache-read savings
using a single "base input price per token". It derives that price by
scanning
`cost_tracker._tokens_sent_by_model` and **breaking on the first**
provider-matching model that
has a price — even though the comment says "most-used model":

```python
# Get the base input price per token for the most-used model on this provider
input_price_per_token = None
if cost_tracker:
    for model_name in cost_tracker._tokens_sent_by_model:   # insertion order, NOT usage order
        ...
        if is_match:
            price_per_1m = cost_tracker._get_list_price(model_name)
            if price_per_1m:
                input_price_per_token = price_per_1m / 1_000_000
                break                                        # first match wins
```

`_tokens_sent_by_model` is insertion-ordered, so the price used depends
on which model was
*recorded first*, not on usage volume. A Claude Code session sends both
Sonnet (main loop) and
Haiku (titles/subagents). If Haiku ($0.80/M) was seen before Sonnet
($3/M), **all** of the
provider's cache-read savings are priced at Haiku's rate — understating
the dashboard's cache
savings by ~3.75×. Reverse the order and it overstates.

Closes: no issue filed — found while auditing the cache-savings pricing.

## Fix

Pick the provider-matching, priced model with the **highest token
volume** instead of breaking
on the first match:

```python
best_tokens = -1
for model_name, tokens_sent in cost_tracker._tokens_sent_by_model.items():
    if is_match and tokens_sent > best_tokens:
        price_per_1m = cost_tracker._get_list_price(model_name)
        if price_per_1m:
            input_price_per_token = price_per_1m / 1_000_000
            best_tokens = tokens_sent
```

## Type of Change

- [x] Bug fix (non-breaking change that fixes an issue)

## Changes Made

- `headroom/proxy/cost.py`: select the highest-volume provider-matching
model (with a known price) rather than the first-recorded one.
- `tests/test_proxy_cache_ttl_metrics.py`: add
`test_prefix_cache_stats_prices_by_most_used_model` using real distinct
per-model prices. (The existing cache-stats tests monkeypatch
`_get_list_price` to a constant `100.0`, which masked the
model-selection logic — hence the bug slipped through.)

## Testing

- [x] New regression test added
(`tests/test_proxy_cache_ttl_metrics.py`)
- [x] Linting/formatting clean — run with the CI-pinned `ruff==0.15.17`
- [ ] Full `pytest` deferred to CI (local-OOM reason below).

```text
$ uvx ruff@0.15.17 check headroom/proxy/cost.py tests/test_proxy_cache_ttl_metrics.py
All checks passed!
```

## Real Behavior Proof

- Environment: Windows 11, Python 3.10, headroom from this branch.
Importing `headroom` pulls in the torch/transformers stack and a full
`pytest` gets OOM-killed on this box, so I verified the selection logic
with a dependency-free script and left the full pytest to CI.
- Exact command / steps: ran a `{haiku: 500, sonnet: 50000}` token map
(Haiku recorded first, Sonnet the higher volume) through both the old
first-match and new highest-volume selection with real prices.
- Observed result: the old logic picks Haiku's $0.80/M (first-inserted);
the new logic picks Sonnet's $3/M (highest volume) and is
insertion-order independent:

```text
OLD picks Haiku price: 0.80/M  (first-inserted)
NEW picks Sonnet price: 3.00/M  (highest volume)
  -> old understates the input price by 3.75x (3.75x)
NEW is insertion-order independent
COST MOST-USED-MODEL FIX VERIFIED
```

- Not tested: rendering the live dashboard (needs the running app). The
fix is confined to the price-selection loop and the new test drives
`build_prefix_cache_stats` directly. Full local `pytest` deferred to CI
(OOM, per above).

## Review Readiness

- [x] I have performed a self-review
- [x] This PR is ready for human review

## Checklist

- [x] My code follows the project's style guidelines
- [x] I have performed a self-review of my code
- [x] I have commented my code, particularly in hard-to-understand areas
- [ ] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective or that my
feature works
- [ ] New and existing unit tests pass locally with my changes — ran
lint + a standalone logic check; full pytest deferred to CI (local OOM,
disclosed above)
- [x] I have updated the CHANGELOG.md if applicable

## Additional Notes

- No new dependencies; a single-loop change plus a test with realistic
prices.
- @JerrettDavis tagging you — this skews the dashboard's per-provider
cache-savings dollar figure by the ratio between a provider's models
(≈3.75× for Sonnet/Haiku), so it seemed worth surfacing. Thanks!

---------

Co-authored-by: JerrettDavis <mxjerrett@gmail.com>
Co-authored-by: Tejas Chopra <chopratejas@gmail.com>
2026-07-13 09:37:58 -04:00

500 lines
18 KiB
Python
Raw Blame History

"""Tests for observed Anthropic cache TTL bucket metrics."""
from __future__ import annotations
import asyncio
import pytest
from headroom.observability import reset_headroom_tracing, reset_otel_metrics
from headroom.proxy.cost import CostTracker, build_prefix_cache_stats
from headroom.proxy.prometheus_metrics import PrometheusMetrics
from headroom.proxy.savings_tracker import SavingsTracker
def test_prometheus_metrics_tracks_observed_ttl_buckets() -> None:
metrics = PrometheusMetrics()
asyncio.run(
metrics.record_request(
provider="anthropic",
model="claude-opus-4-6",
input_tokens=100,
output_tokens=20,
tokens_saved=5,
latency_ms=10.0,
cache_read_tokens=40,
cache_write_tokens=60,
cache_write_5m_tokens=10,
cache_write_1h_tokens=50,
)
)
stats = metrics.cache_by_provider["anthropic"]
assert stats["cache_write_5m_tokens"] == 10
assert stats["cache_write_1h_tokens"] == 50
assert stats["cache_write_5m_requests"] == 1
assert stats["cache_write_1h_requests"] == 1
def test_cost_tracker_exposes_observed_ttl_buckets_per_model() -> None:
tracker = CostTracker()
tracker.record_tokens(
"claude-opus-4-6",
tokens_saved=10,
tokens_sent=90,
cache_read_tokens=40,
cache_write_tokens=60,
cache_write_5m_tokens=10,
cache_write_1h_tokens=50,
uncached_tokens=20,
)
stats = tracker.stats()
assert stats["cache_write_5m_tokens"] == 10
assert stats["cache_write_1h_tokens"] == 50
assert stats["per_model"]["claude-opus-4-6"]["cache_write_5m_tokens"] == 10
assert stats["per_model"]["claude-opus-4-6"]["cache_write_1h_tokens"] == 50
def test_prefix_cache_stats_include_observed_ttl_mix() -> None:
metrics = PrometheusMetrics()
provider_stats = metrics.cache_by_provider["anthropic"]
provider_stats["requests"] = 2
provider_stats["hit_requests"] = 1
provider_stats["cache_read_tokens"] = 40
provider_stats["cache_write_tokens"] = 60
provider_stats["cache_write_5m_tokens"] = 15
provider_stats["cache_write_1h_tokens"] = 45
provider_stats["cache_write_5m_requests"] = 1
provider_stats["cache_write_1h_requests"] = 1
stats = build_prefix_cache_stats(metrics, None)
anthropic = stats["by_provider"]["anthropic"]
assert anthropic["observed_ttl_buckets"]["5m"]["tokens"] == 15
assert anthropic["observed_ttl_buckets"]["1h"]["tokens"] == 45
assert anthropic["observed_ttl_mix"]["5m_pct"] == 25.0
assert anthropic["observed_ttl_mix"]["1h_pct"] == 75.0
assert stats["totals"]["observed_ttl_buckets"]["5m"]["tokens"] == 15
assert stats["totals"]["observed_ttl_buckets"]["1h"]["tokens"] == 45
def test_prefix_cache_stats_subtracts_write_premium_from_provider_net_savings(
monkeypatch: pytest.MonkeyPatch,
) -> None:
metrics = PrometheusMetrics()
metrics.cache_by_provider["anthropic"].update(
{
"requests": 2,
"hit_requests": 1,
"cache_read_tokens": 40,
"cache_write_tokens": 60,
"cache_write_5m_tokens": 60,
"cache_write_1h_tokens": 0,
"cache_write_5m_requests": 1,
"cache_write_1h_requests": 0,
}
)
tracker = CostTracker()
tracker._tokens_sent_by_model.update({"claude-opus-4-6": 1})
monkeypatch.setattr(CostTracker, "_get_list_price", lambda _self, _model: 100.0)
stats = build_prefix_cache_stats(metrics, tracker)
anthropic = stats["by_provider"]["anthropic"]
assert anthropic["savings_usd"] == 0.0036
assert anthropic["write_premium_usd"] == 0.0015
assert anthropic["net_savings_usd"] == 0.0021
def test_prefix_cache_stats_prices_by_most_used_model(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""Cache-read savings must be valued with the highest-volume model's price,
not whichever model was recorded first. A Claude Code session sends Haiku
(titles) and Sonnet (main loop); pricing the savings at Haiku's rate because
it was inserted first understates the dashboard figure ~3.75x."""
prices = {"claude-haiku-4-5": 0.80, "claude-sonnet-4-5": 3.00}
monkeypatch.setattr(CostTracker, "_get_list_price", lambda _self, m: prices.get(m))
def _savings(tokens_by_model: dict[str, int]) -> float:
metrics = PrometheusMetrics()
# Use a large read count so the reported savings_usd (rounded to 4 dp)
# stays exact and the price ratio is not lost to rounding.
metrics.cache_by_provider["anthropic"].update(
{
"requests": 1,
"hit_requests": 1,
"cache_read_tokens": 1_000_000,
"cache_write_tokens": 0,
"cache_write_5m_tokens": 0,
"cache_write_1h_tokens": 0,
"cache_write_5m_requests": 0,
"cache_write_1h_requests": 0,
}
)
tracker = CostTracker()
tracker._tokens_sent_by_model.update(tokens_by_model)
stats = build_prefix_cache_stats(metrics, tracker)
return stats["by_provider"]["anthropic"]["savings_usd"]
# Haiku recorded first, but Sonnet carries the higher token volume.
haiku_first = _savings({"claude-haiku-4-5": 500, "claude-sonnet-4-5": 50_000})
sonnet_only = _savings({"claude-sonnet-4-5": 50_000})
haiku_only = _savings({"claude-haiku-4-5": 500})
# Priced by Sonnet regardless of insertion order, not by first-seen Haiku.
assert haiku_first == sonnet_only
assert haiku_first > haiku_only
assert haiku_only == pytest.approx(sonnet_only * 0.80 / 3.00)
def test_prefix_cache_stats_subtracts_write_premium_from_total_net_savings(
monkeypatch: pytest.MonkeyPatch,
) -> None:
metrics = PrometheusMetrics()
metrics.cache_by_provider["anthropic"].update(
{
"requests": 2,
"hit_requests": 1,
"cache_read_tokens": 40,
"cache_write_tokens": 60,
"cache_write_5m_tokens": 60,
"cache_write_1h_tokens": 0,
"cache_write_5m_requests": 1,
"cache_write_1h_requests": 0,
}
)
metrics.cache_by_provider["openai"].update(
{
"requests": 1,
"hit_requests": 1,
"cache_read_tokens": 20,
"cache_write_tokens": 10,
"cache_write_5m_tokens": 0,
"cache_write_1h_tokens": 10,
"cache_write_5m_requests": 0,
"cache_write_1h_requests": 1,
}
)
tracker = CostTracker()
tracker._tokens_sent_by_model.update({"claude-opus-4-6": 1, "gpt-4o": 1})
monkeypatch.setattr(CostTracker, "_get_list_price", lambda _self, _model: 100.0)
stats = build_prefix_cache_stats(metrics, tracker)
openai = stats["by_provider"]["openai"]
assert openai["write_premium_usd"] == 0.0
assert openai["net_savings_usd"] == openai["savings_usd"]
assert stats["totals"]["savings_usd"] == 0.0046
assert stats["totals"]["write_premium_usd"] == 0.0015
assert stats["totals"]["net_savings_usd"] == 0.0031
def test_prefix_cache_stats_keeps_net_equal_without_write_premium(
monkeypatch: pytest.MonkeyPatch,
) -> None:
metrics = PrometheusMetrics()
metrics.cache_by_provider["openai"].update(
{
"requests": 1,
"hit_requests": 1,
"cache_read_tokens": 20,
"cache_write_tokens": 0,
"cache_write_5m_tokens": 0,
"cache_write_1h_tokens": 0,
"cache_write_5m_requests": 0,
"cache_write_1h_requests": 0,
}
)
tracker = CostTracker()
tracker._tokens_sent_by_model.update({"gpt-4o": 1})
monkeypatch.setattr(CostTracker, "_get_list_price", lambda _self, _model: 100.0)
stats = build_prefix_cache_stats(metrics, tracker)
openai = stats["by_provider"]["openai"]
assert openai["savings_usd"] == 0.001
assert openai["write_premium_usd"] == 0.0
assert openai["net_savings_usd"] == openai["savings_usd"]
def test_prometheus_metrics_export_includes_extended_fields(tmp_path) -> None:
metrics = PrometheusMetrics(
savings_tracker=SavingsTracker(path=str(tmp_path / "proxy_savings.json"))
)
asyncio.run(
metrics.record_request(
provider="anthropic",
model="claude-opus-4-6",
input_tokens=100,
output_tokens=20,
tokens_saved=5,
latency_ms=12.5,
overhead_ms=3.0,
ttfb_ms=9.0,
pipeline_timing={"router": 4.5},
waste_signals={"json_bloat": 7},
cache_read_tokens=40,
cache_write_tokens=60,
cache_write_5m_tokens=10,
cache_write_1h_tokens=50,
uncached_input_tokens=20,
)
)
asyncio.run(metrics.record_cache_bust(11))
exported = asyncio.run(metrics.export())
assert "headroom_requests_total 1" in exported
assert "headroom_tokens_saved_total 5" in exported
assert "headroom_persistent_savings_requests_total 1" in exported
assert "headroom_persistent_savings_tokens_saved_total 5" in exported
assert "headroom_persistent_savings_input_tokens_total 100" in exported
assert "headroom_latency_ms_count 1" in exported
assert 'headroom_transform_timing_ms_sum{transform="router"} 4.5' in exported
assert 'headroom_waste_signal_tokens_total{signal="json_bloat"} 7' in exported
assert 'headroom_cache_write_ttl_tokens_total{provider="anthropic",ttl="5m"} 10' in exported
assert 'headroom_provider_cache_hit_requests_total{provider="anthropic"} 1' in exported
assert "headroom_cache_bust_tokens_lost_total 11" in exported
def test_prometheus_export_includes_persistent_savings_after_restart(tmp_path) -> None:
savings_path = tmp_path / "proxy_savings.json"
metrics = PrometheusMetrics(savings_tracker=SavingsTracker(path=str(savings_path)))
asyncio.run(
metrics.record_request(
provider="openai",
model="gpt-4o",
input_tokens=120,
output_tokens=20,
tokens_saved=40,
latency_ms=12.5,
)
)
reloaded = PrometheusMetrics(savings_tracker=SavingsTracker(path=str(savings_path)))
exported = asyncio.run(reloaded.export())
assert "headroom_tokens_saved_total 0" in exported
assert "headroom_requests_total 0" in exported
assert "headroom_persistent_savings_requests_total 1" in exported
assert "headroom_persistent_savings_tokens_saved_total 40" in exported
assert "headroom_persistent_savings_input_tokens_total 120" in exported
def test_streaming_parser_extracts_anthropic_ttl_bucket_usage() -> None:
from headroom.proxy.server import HeadroomProxy, ProxyConfig
proxy = HeadroomProxy(
ProxyConfig(
optimize=False,
cache_enabled=False,
rate_limit_enabled=False,
cost_tracking_enabled=False,
log_requests=False,
ccr_inject_tool=False,
ccr_handle_responses=False,
ccr_context_tracking=False,
)
)
chunk = (
b'data: {"type":"message_start","message":{"usage":{"input_tokens":12,'
b'"cache_read_input_tokens":3,"cache_creation_input_tokens":9,'
b'"cache_creation":{"ephemeral_5m_input_tokens":4,"ephemeral_1h_input_tokens":5}}}}\n\n'
)
usage = proxy._parse_sse_usage(chunk, "anthropic")
assert usage is not None
assert usage["cache_creation_ephemeral_5m_input_tokens"] == 4
assert usage["cache_creation_ephemeral_1h_input_tokens"] == 5
def test_stats_endpoint_reports_observed_ttl_buckets() -> None:
pytest.importorskip("fastapi")
from fastapi.testclient import TestClient
from headroom.proxy.server import ProxyConfig, create_app
app = create_app(
ProxyConfig(
optimize=False,
cache_enabled=False,
rate_limit_enabled=False,
cost_tracking_enabled=False,
log_requests=False,
ccr_inject_tool=False,
ccr_handle_responses=False,
ccr_context_tracking=False,
)
)
proxy = app.state.proxy
provider_stats = proxy.metrics.cache_by_provider["anthropic"]
provider_stats["requests"] = 1
provider_stats["hit_requests"] = 1
provider_stats["cache_read_tokens"] = 30
provider_stats["cache_write_tokens"] = 70
provider_stats["cache_write_5m_tokens"] = 20
provider_stats["cache_write_1h_tokens"] = 50
provider_stats["cache_write_5m_requests"] = 1
provider_stats["cache_write_1h_requests"] = 1
with TestClient(app) as client:
response = client.get("/stats")
assert response.status_code == 200
prefix_cache = response.json()["prefix_cache"]
anthropic = prefix_cache["by_provider"]["anthropic"]
assert anthropic["observed_ttl_buckets"]["5m"]["tokens"] == 20
assert anthropic["observed_ttl_buckets"]["1h"]["tokens"] == 50
assert prefix_cache["totals"]["observed_ttl_mix"]["active_buckets"] == ["5m", "1h"]
def test_stats_endpoint_reports_otel_configuration(monkeypatch: pytest.MonkeyPatch) -> None:
pytest.importorskip("fastapi")
from fastapi.testclient import TestClient
from headroom.proxy.server import ProxyConfig, create_app
reset_otel_metrics()
monkeypatch.setenv("HEADROOM_OTEL_METRICS_ENABLED", "1")
monkeypatch.setenv("HEADROOM_OTEL_METRICS_EXPORTER", "console")
app = create_app(
ProxyConfig(
optimize=False,
cache_enabled=False,
rate_limit_enabled=False,
cost_tracking_enabled=False,
log_requests=False,
ccr_inject_tool=False,
ccr_handle_responses=False,
ccr_context_tracking=False,
)
)
with TestClient(app) as client:
response = client.get("/stats")
assert response.status_code == 200
otel = response.json()["otel"]
assert otel["configured"] is True
assert otel["enabled"] is True
assert otel["service_name"] == "headroom-proxy"
assert otel["exporter"] == "console"
def test_stats_endpoint_reports_langfuse_configuration(monkeypatch: pytest.MonkeyPatch) -> None:
pytest.importorskip("fastapi")
from fastapi.testclient import TestClient
from headroom.proxy.server import ProxyConfig, create_app
reset_headroom_tracing()
monkeypatch.setenv("HEADROOM_LANGFUSE_ENABLED", "1")
monkeypatch.setenv("LANGFUSE_PUBLIC_KEY", "pk-lf-test")
monkeypatch.setenv("LANGFUSE_SECRET_KEY", "sk-lf-test")
monkeypatch.setenv("LANGFUSE_BASE_URL", "https://cloud.langfuse.com")
app = create_app(
ProxyConfig(
optimize=False,
cache_enabled=False,
rate_limit_enabled=False,
cost_tracking_enabled=False,
log_requests=False,
ccr_inject_tool=False,
ccr_handle_responses=False,
ccr_context_tracking=False,
)
)
with TestClient(app) as client:
response = client.get("/stats")
assert response.status_code == 200
langfuse = response.json()["langfuse"]
assert langfuse["configured"] is True
assert langfuse["enabled"] is True
assert langfuse["service_name"] == "headroom-proxy"
assert langfuse["endpoint"] == "https://cloud.langfuse.com/api/public/otel/v1/traces"
# --- Cache-miss attribution (#1313) ---
def test_record_cache_miss_attribution_buckets_by_provider_and_reason() -> None:
metrics = PrometheusMetrics()
asyncio.run(metrics.record_cache_miss_attribution("anthropic", "ttl_expiry"))
asyncio.run(metrics.record_cache_miss_attribution("anthropic", "ttl_expiry"))
asyncio.run(metrics.record_cache_miss_attribution("anthropic", "prefix_change"))
asyncio.run(metrics.record_cache_miss_attribution("anthropic", "unknown"))
buckets = metrics.cache_miss_attribution_by_provider["anthropic"]
assert buckets["ttl_expiry"] == 2
assert buckets["prefix_change"] == 1
assert buckets["unknown"] == 1
def test_prefix_cache_stats_include_miss_attribution() -> None:
metrics = PrometheusMetrics()
asyncio.run(metrics.record_cache_miss_attribution("anthropic", "ttl_expiry"))
asyncio.run(metrics.record_cache_miss_attribution("anthropic", "ttl_expiry"))
asyncio.run(metrics.record_cache_miss_attribution("anthropic", "prefix_change"))
asyncio.run(metrics.record_cache_miss_attribution("anthropic", "unknown"))
stats = build_prefix_cache_stats(metrics, None)
ma = stats["miss_attribution"]
assert ma["totals"]["ttl_expiry"] == 2
assert ma["totals"]["prefix_change"] == 1
assert ma["totals"]["unknown"] == 1
assert ma["totals"]["total"] == 4
# Percentages are over attributed (non-unknown) misses: 2 / 3, 1 / 3.
assert ma["totals"]["ttl_expiry_pct"] == 66.7
assert ma["totals"]["prefix_change_pct"] == 33.3
assert ma["by_provider"]["anthropic"]["total"] == 4
def test_prefix_cache_stats_miss_attribution_empty_when_no_misses() -> None:
metrics = PrometheusMetrics()
stats = build_prefix_cache_stats(metrics, None)
ma = stats["miss_attribution"]
assert ma["totals"]["total"] == 0
assert ma["totals"]["ttl_expiry_pct"] == 0.0
assert ma["by_provider"] == {}
def test_prometheus_export_includes_miss_attribution() -> None:
metrics = PrometheusMetrics()
asyncio.run(metrics.record_cache_miss_attribution("anthropic", "ttl_expiry"))
asyncio.run(metrics.record_cache_miss_attribution("anthropic", "prefix_change"))
exported = asyncio.run(metrics.export())
assert (
'headroom_cache_miss_attribution_total{provider="anthropic",reason="ttl_expiry"} 1'
in exported
)
assert (
'headroom_cache_miss_attribution_total{provider="anthropic",reason="prefix_change"} 1'
in exported
)
def test_reset_runtime_clears_miss_attribution() -> None:
metrics = PrometheusMetrics()
asyncio.run(metrics.record_cache_miss_attribution("anthropic", "ttl_expiry"))
asyncio.run(metrics.reset_runtime())
assert dict(metrics.cache_miss_attribution_by_provider) == {}
Reference in a new issue View git blame Copy permalink