mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-10 14:27:00 -04:00
feat(savings): durable savings ledger + headroom savings command (#1127)
## Description
Adds a durable, cross-process savings ledger and a `headroom savings`
CLI that shows cost avoided plus Today / Last 7 days / All time
breakdowns by model and client. Unlike `headroom_stats` (a per-session,
in-memory snapshot), the ledger is on disk and survives proxy and agent
restarts, and is safe across the many MCP processes Headroom spawns.
## 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
- Add `headroom/savings_ledger.py`: append-only, `fcntl`-locked JSONL
ledger at `~/.headroom/savings_events.jsonl`, safe across concurrent
writers (main MCP server, each subagent, and the proxy), aggregated on
read so totals survive restarts.
- litellm list pricing for known models; blended `$3/1M` input-token
fallback for `model="unknown"` (MCP compressions do not know the
upstream model). Self-pruning: events past the 365-day retention window
are dropped on read and the file is compacted once large.
- Add `headroom savings` CLI (`headroom/cli/savings.py`) with `--json`,
`--days N`, and `--reset` flags.
- Proxy client attribution: `record_request` accepts `client` and
threads `outcome.client` into the ledger, so proxy events record the
real harness (claude-code, codex, cursor, …) from the existing
`classify_client()` detection, falling back to `"proxy"` only when
unidentified.
- MCP compress hook records the client (from `clientInfo.name`) and
tokens saved after each `headroom_compress`; `HEADROOM_MCP_CLIENT` /
`HEADROOM_MCP_MODEL` env overrides.
- Add the `savings_events_path()` helper +
`HEADROOM_SAVINGS_EVENTS_PATH` env in `headroom/paths.py`, the docs page
`docs/content/docs/savings.mdx`, and 15 tests.
## Testing
- [x] Unit tests pass (`pytest`)
- [x] Linting passes (`ruff check .`)
- [x] Type checking passes (`mypy headroom`)
- [x] New tests added for new functionality
- [x] Manual testing performed
### Test Output
The single warning is a pre-existing, repo-wide
`StarletteDeprecationWarning` from
`fastapi.testclient` (the venv has `httpx`, not `httpx2`); it is
unrelated to this
change and fires in every proxy test that spins up a `TestClient`.
```text
$ .venv/bin/python -m pytest tests/test_savings_ledger.py -q
............... [100%]
15 passed, 1 warning in 5.17s
# warning: fastapi/testclient.py StarletteDeprecationWarning (httpx vs httpx2) — third-party, pre-existing
$ .venv/bin/ruff check headroom/savings_ledger.py headroom/cli/savings.py \
headroom/ccr/mcp_server.py headroom/proxy/prometheus_metrics.py \
headroom/proxy/outcome.py headroom/paths.py tests/test_savings_ledger.py
All checks passed!
$ .venv/bin/mypy headroom/savings_ledger.py headroom/cli/savings.py \
headroom/ccr/mcp_server.py headroom/proxy/prometheus_metrics.py \
headroom/proxy/outcome.py headroom/paths.py
Success: no issues found in 6 source files
```
## Real Behavior Proof
- Environment: macOS (Darwin 25.5.0), Python 3.13.13, editable install
of this branch, proxy running on :8787
- Exact command / steps: route live agent + proxy traffic through
Headroom, then run `headroom savings`
- Observed result: distinct Today / Last 7 days / All time windows with
per-model and per-client breakdowns, as below
- Not tested: Windows runtime (no `fcntl`; the ledger falls back to
best-effort append)
```text
Today ██░░░░░░░░░░░░░░ 11.3% saved 472,870 / 4,193,288 tokens $1.5920
Last 7 days ██░░░░░░░░░░░░░░ 11.9% saved 505,170 / 4,244,288 tokens $1.7385
All time ██░░░░░░░░░░░░░░ 13.0% saved 566,170 / 4,339,288 tokens $1.9815
Cost avoided per model:
claude-sonnet-4-6 $1.2200
claude-opus-4-8 $0.6685
gpt-5.5 $0.0840
claude-haiku-4-5 $0.0090
Savings by client:
claude-code 60 calls · 524,970 tokens saved
cursor 2 calls · 16,800 tokens saved
codex 3 calls · 24,400 tokens saved
```
## 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
- [x] 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
- [x] New and existing unit tests pass locally with my changes
- [x] I have updated the CHANGELOG.md if applicable
## Additional Notes
The one pytest warning is a third-party `StarletteDeprecationWarning`
from `fastapi.testclient` (pre-existing, repo-wide); not introduced
here. CHANGELOG.md not updated.
This commit is contained in:
Ztkent
GitHub
parent
f39858c233
commit
978ffa0a6a
10 changed files with 921 additions and 0 deletions
|
|
@ -39,6 +39,7 @@
|
|||
"configuration",
|
||||
"filesystem-contract",
|
||||
"---Observability---",
|
||||
"savings",
|
||||
"metrics",
|
||||
"simulation",
|
||||
"---API Reference---",
|
||||
|
|
|
|||
62
docs/content/docs/savings.mdx
Normal file
62
docs/content/docs/savings.mdx
Normal file
|
|
@ -0,0 +1,62 @@
|
|||
---
|
||||
title: Savings Tracking
|
||||
description: Durable, over-time compression savings — cost avoided plus Today / Last 7 days / All time and per-model / per-client breakdowns via `headroom savings`.
|
||||
---
|
||||
|
||||
`headroom savings` shows how much Headroom has saved you over time — cost avoided, token counts, and breakdowns by model and client. Unlike `headroom_stats` (a single in-memory session snapshot), it reads a **durable ledger** that survives proxy and agent restarts.
|
||||
|
||||
## Usage
|
||||
|
||||
```bash
|
||||
headroom savings # human-readable summary
|
||||
headroom savings --json # machine-readable report
|
||||
headroom savings --days 30 # restrict the lookback/retention window
|
||||
headroom savings --reset # delete the ledger and start fresh
|
||||
```
|
||||
|
||||
### Example
|
||||
|
||||
```text
|
||||
Today ███████████░░░░░ 67.9% saved 19,000 / 28,000 tokens $0.0850
|
||||
Last 7 days ███████████░░░░░ 67.1% saved 47,000 / 70,000 tokens $0.2250
|
||||
All time ██████████░░░░░░ 65.0% saved 78,000 / 120,000 tokens $0.2680
|
||||
|
||||
Cost avoided per model:
|
||||
claude-opus-4-8 $0.1750
|
||||
gpt-5.5 $0.0350
|
||||
unknown $0.0330
|
||||
claude-haiku-4-5 $0.0250
|
||||
|
||||
Savings by client:
|
||||
claude-code 4 calls · 60,000 tokens saved
|
||||
codex 2 calls · 18,000 tokens saved
|
||||
```
|
||||
|
||||
## How it works
|
||||
|
||||
Every compression appends one line to an **append-only, file-locked event ledger** at `~/.headroom/savings_events.jsonl`, and `headroom savings` aggregates it on read. This design is:
|
||||
|
||||
- **Durable** — the ledger is on disk, so totals survive proxy and agent restarts.
|
||||
- **Accurate under concurrency** — Headroom's MCP server runs as multiple processes (the main agent plus each subagent), and the proxy is a separate process. An append-only, locked log lets every writer contribute without the lost-update races a single shared mutable file would suffer.
|
||||
- **Self-pruning** — events older than the retention window (365 days by default) are dropped on read, and the file is compacted once it grows large.
|
||||
|
||||
Both compression paths feed the same ledger:
|
||||
|
||||
- **MCP tool** — each `headroom_compress` call records its client (the MCP client name) and tokens saved.
|
||||
- **Proxy** — each request records its real upstream model, so cost is priced accurately.
|
||||
|
||||
### Cost basis
|
||||
|
||||
Cost avoided is the dollar value of the saved **input** tokens. Headroom uses [litellm](https://headroom-docs.vercel.app/docs/litellm) list pricing where the model is known (proxy traffic). MCP-tool compressions don't know the agent's upstream model, so they record `model="unknown"` and fall back to a blended per-token rate rather than reporting `$0`.
|
||||
|
||||
## Configuration
|
||||
|
||||
| Variable | Purpose |
|
||||
| --- | --- |
|
||||
| `HEADROOM_SAVINGS_EVENTS_PATH` | Override the ledger location (default `~/.headroom/savings_events.jsonl`). |
|
||||
| `HEADROOM_MCP_CLIENT` | Override the client label recorded by the MCP tool path. |
|
||||
| `HEADROOM_MCP_MODEL` | Optional model hint so MCP-tool compressions price against a known model instead of the blended fallback. |
|
||||
|
||||
<Callout type="info">
|
||||
`headroom savings` is distinct from `headroom_stats` (a per-session, in-memory snapshot) and from the proxy's live `/stats` endpoint (backed by `proxy_savings.json`). The savings ledger is the durable, cross-process source of truth.
|
||||
</Callout>
|
||||
|
|
@ -33,6 +33,7 @@ from pathlib import Path
|
|||
from typing import Any
|
||||
|
||||
from headroom import paths as _paths
|
||||
from headroom import savings_ledger
|
||||
|
||||
# fcntl is Unix-only; on Windows we skip file locking (stats are best-effort).
|
||||
# Keep the module typed as Any so Windows mypy runs don't try to resolve Unix-only attrs.
|
||||
|
|
@ -662,8 +663,49 @@ class HeadroomMCPServer:
|
|||
loop = asyncio.get_running_loop()
|
||||
result = await loop.run_in_executor(None, self._compress_content, content)
|
||||
|
||||
# Record durably so `headroom savings` reflects this compression across
|
||||
# restarts. Best-effort: never let savings bookkeeping break the tool.
|
||||
try:
|
||||
self._record_savings(result)
|
||||
except Exception:
|
||||
logger.debug("durable savings recording failed", exc_info=True)
|
||||
|
||||
return [TextContent(type="text", text=json.dumps(result, indent=2))]
|
||||
|
||||
def _record_savings(self, result: dict[str, Any]) -> None:
|
||||
"""Append a durable savings event for a completed compression."""
|
||||
try:
|
||||
before = int(result.get("original_tokens", 0) or 0)
|
||||
after = int(result.get("compressed_tokens", 0) or 0)
|
||||
except (TypeError, ValueError):
|
||||
return
|
||||
if before <= after:
|
||||
return
|
||||
savings_ledger.record_savings_event(
|
||||
tokens_before=before,
|
||||
tokens_after=after,
|
||||
# The MCP tool doesn't know the agent's upstream model; an optional
|
||||
# hint lets a host attribute it, otherwise it records as "unknown".
|
||||
model=os.environ.get("HEADROOM_MCP_MODEL"),
|
||||
client=self._current_client(),
|
||||
source="mcp",
|
||||
)
|
||||
|
||||
def _current_client(self) -> str:
|
||||
"""Name of the MCP client driving this session (best-effort)."""
|
||||
override = os.environ.get("HEADROOM_MCP_CLIENT")
|
||||
if override:
|
||||
return override
|
||||
try:
|
||||
params = self.server.request_context.session.client_params
|
||||
info = getattr(params, "clientInfo", None) if params else None
|
||||
name = getattr(info, "name", None)
|
||||
if name:
|
||||
return str(name)
|
||||
except Exception:
|
||||
pass
|
||||
return "unknown"
|
||||
|
||||
async def _handle_retrieve(self, arguments: dict[str, Any]) -> list[TextContent]:
|
||||
"""Handle headroom_retrieve tool call."""
|
||||
hash_key = arguments.get("hash")
|
||||
|
|
|
|||
|
|
@ -61,6 +61,7 @@ def _register_commands() -> None:
|
|||
output_savings, # noqa: F401
|
||||
perf, # noqa: F401
|
||||
proxy, # noqa: F401
|
||||
savings, # noqa: F401
|
||||
tools, # noqa: F401
|
||||
update, # noqa: F401
|
||||
wrap, # noqa: F401
|
||||
|
|
|
|||
105
headroom/cli/savings.py
Normal file
105
headroom/cli/savings.py
Normal file
|
|
@ -0,0 +1,105 @@
|
|||
"""CLI: show durable compression savings over time.
|
||||
|
||||
Reads the append-only savings ledger (``~/.headroom/savings_events.jsonl``,
|
||||
written by both the MCP tool path and the proxy) and renders a cost-avoided
|
||||
summary with Today / Last 7 days / All time bars plus per-model and
|
||||
per-client breakdowns. Durable across restarts; aggregated on read.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
import click
|
||||
|
||||
from headroom import savings_ledger
|
||||
|
||||
from .main import main
|
||||
|
||||
_BAR_WIDTH = 16
|
||||
|
||||
|
||||
def _bar(percent: float, width: int = _BAR_WIDTH) -> str:
|
||||
filled = int(round(percent / 100 * width))
|
||||
filled = max(0, min(width, filled))
|
||||
return "█" * filled + "░" * (width - filled)
|
||||
|
||||
|
||||
def _money(value: float, places: int = 4) -> str:
|
||||
return f"${value:,.{places}f}"
|
||||
|
||||
|
||||
def _tokens(value: int) -> str:
|
||||
return f"{value:,}"
|
||||
|
||||
|
||||
def _window_line(label: str, window: dict[str, Any]) -> str:
|
||||
pct = float(window.get("savings_percent", 0.0) or 0.0)
|
||||
saved = int(window.get("tokens_saved", 0) or 0)
|
||||
before = int(window.get("tokens_before", 0) or 0)
|
||||
cost = float(window.get("cost_usd", 0.0) or 0.0)
|
||||
return (
|
||||
f"{label:<11} {_bar(pct)} {pct:5.1f}% "
|
||||
f"saved {_tokens(saved)} / {_tokens(before)} tokens {_money(cost)}"
|
||||
)
|
||||
|
||||
|
||||
@main.command(name="savings")
|
||||
@click.option("--json", "as_json", is_flag=True, help="Emit the raw report as JSON.")
|
||||
@click.option(
|
||||
"--days",
|
||||
type=int,
|
||||
default=savings_ledger.DEFAULT_RETENTION_DAYS,
|
||||
show_default=True,
|
||||
help="Retention/lookback window for the ledger, in days.",
|
||||
)
|
||||
@click.option("--reset", is_flag=True, help="Delete the savings ledger and start fresh.")
|
||||
def savings(as_json: bool, days: int, reset: bool) -> None:
|
||||
"""Show durable compression savings over time."""
|
||||
|
||||
if reset:
|
||||
path = savings_ledger._resolve_path(None)
|
||||
if path.exists():
|
||||
path.unlink()
|
||||
click.echo(f"Ledger reset: {path}")
|
||||
else:
|
||||
click.echo("Nothing to reset — ledger does not exist.")
|
||||
return
|
||||
|
||||
report = savings_ledger.aggregate_savings(retention_days=days)
|
||||
|
||||
if as_json:
|
||||
click.echo(json.dumps(report.to_dict(), indent=2))
|
||||
return
|
||||
|
||||
lifetime = report.lifetime
|
||||
calls = int(lifetime.get("calls", 0) or 0)
|
||||
if calls == 0:
|
||||
click.echo("No savings recorded yet.")
|
||||
click.echo(
|
||||
"Compress via the Headroom MCP tool or route traffic through the "
|
||||
"proxy, then re-run `headroom savings`."
|
||||
)
|
||||
click.echo(f"Ledger: {report.path}")
|
||||
return
|
||||
|
||||
click.echo("")
|
||||
click.echo(_window_line("Today", report.windows["today"]))
|
||||
click.echo(_window_line("Last 7 days", report.windows["last_7_days"]))
|
||||
click.echo(_window_line("All time", report.windows["all_time"]))
|
||||
|
||||
if report.by_model:
|
||||
click.echo("")
|
||||
click.echo("Cost avoided per model:")
|
||||
for row in report.by_model:
|
||||
click.echo(f" {str(row['model']):<24} {_money(float(row['cost_usd']))}")
|
||||
|
||||
if report.by_client:
|
||||
click.echo("")
|
||||
click.echo("Savings by client:")
|
||||
for row in report.by_client:
|
||||
click.echo(
|
||||
f" {str(row['client']):<24} {int(row['calls']):,} calls · "
|
||||
f"{_tokens(int(row['tokens_saved']))} tokens saved"
|
||||
)
|
||||
|
|
@ -47,6 +47,7 @@ HEADROOM_WORKSPACE_DIR_ENV = "HEADROOM_WORKSPACE_DIR"
|
|||
# ---------------------------------------------------------------------------
|
||||
|
||||
HEADROOM_SAVINGS_PATH_ENV = "HEADROOM_SAVINGS_PATH"
|
||||
HEADROOM_SAVINGS_EVENTS_PATH_ENV = "HEADROOM_SAVINGS_EVENTS_PATH"
|
||||
HEADROOM_TOIN_PATH_ENV = "HEADROOM_TOIN_PATH"
|
||||
HEADROOM_SUBSCRIPTION_STATE_PATH_ENV = "HEADROOM_SUBSCRIPTION_STATE_PATH"
|
||||
|
||||
|
|
@ -66,6 +67,7 @@ _MEMORY_DB_FILE = "memory.db"
|
|||
_MEMORIES_DIR = "memories"
|
||||
_LICENSE_CACHE_FILE = "license_cache.json"
|
||||
_SESSION_STATS_FILE = "session_stats.jsonl"
|
||||
_SAVINGS_EVENTS_FILE = "savings_events.jsonl"
|
||||
_SYNC_STATE_FILE = "sync_state.json"
|
||||
_BRIDGE_STATE_FILE = "bridge_state.json"
|
||||
_LOGS_DIR = "logs"
|
||||
|
|
@ -228,6 +230,21 @@ def session_stats_path() -> Path:
|
|||
return workspace_dir() / _SESSION_STATS_FILE
|
||||
|
||||
|
||||
def savings_events_path(explicit: str | os.PathLike[str] | None = None) -> Path:
|
||||
"""Return the path for the durable append-only savings event ledger.
|
||||
|
||||
Unlike :func:`session_stats_path` (pruned to a short rolling window), this
|
||||
file accrues one line per compression across proxy restarts and concurrent
|
||||
MCP processes, and is the source of truth for ``headroom savings``.
|
||||
"""
|
||||
|
||||
return _resolve(
|
||||
explicit,
|
||||
HEADROOM_SAVINGS_EVENTS_PATH_ENV,
|
||||
workspace_dir() / _SAVINGS_EVENTS_FILE,
|
||||
)
|
||||
|
||||
|
||||
def sync_state_path() -> Path:
|
||||
"""Return the path for memory sync state."""
|
||||
|
||||
|
|
@ -344,6 +361,7 @@ __all__ = [
|
|||
"HEADROOM_CONFIG_DIR_ENV",
|
||||
"HEADROOM_WORKSPACE_DIR_ENV",
|
||||
"HEADROOM_SAVINGS_PATH_ENV",
|
||||
"HEADROOM_SAVINGS_EVENTS_PATH_ENV",
|
||||
"HEADROOM_TOIN_PATH_ENV",
|
||||
"HEADROOM_SUBSCRIPTION_STATE_PATH_ENV",
|
||||
"config_dir",
|
||||
|
|
@ -357,6 +375,7 @@ __all__ = [
|
|||
"native_memory_dir",
|
||||
"license_cache_path",
|
||||
"session_stats_path",
|
||||
"savings_events_path",
|
||||
"sync_state_path",
|
||||
"bridge_state_path",
|
||||
"log_dir",
|
||||
|
|
|
|||
|
|
@ -369,6 +369,7 @@ async def emit_request_outcome(handler: Any, outcome: RequestOutcome) -> None:
|
|||
uncached_input_tokens=outcome.uncached_input_tokens,
|
||||
attempted_input_tokens=outcome.attempted_input_tokens,
|
||||
project=project,
|
||||
client=outcome.client,
|
||||
)
|
||||
|
||||
# 2. Cost tracker (optional).
|
||||
|
|
|
|||
|
|
@ -20,6 +20,7 @@ if TYPE_CHECKING:
|
|||
from headroom.observability import HeadroomOtelMetrics
|
||||
from headroom.proxy.cost import CostTracker
|
||||
|
||||
from headroom import savings_ledger
|
||||
from headroom.observability import get_otel_metrics
|
||||
from headroom.proxy.savings_tracker import SavingsTracker
|
||||
|
||||
|
|
@ -563,6 +564,7 @@ class PrometheusMetrics:
|
|||
uncached_input_tokens: int = 0,
|
||||
attempted_input_tokens: int = 0,
|
||||
project: str | None = None,
|
||||
client: str | None = None,
|
||||
):
|
||||
"""Record metrics for a request."""
|
||||
async with self._lock:
|
||||
|
|
@ -658,6 +660,21 @@ class PrometheusMetrics:
|
|||
total_input_cost_usd=total_input_cost_usd,
|
||||
)
|
||||
|
||||
# Also append to the durable, multi-process savings ledger so
|
||||
# `headroom savings` reflects proxy traffic alongside MCP-tool usage.
|
||||
# The real upstream model means litellm prices it accurately. The
|
||||
# client is the harness classified from the User-Agent / X-Client
|
||||
# (claude-code, codex, cursor, ...); it falls back to "proxy" only
|
||||
# when the harness is unidentified.
|
||||
if tokens_saved > 0:
|
||||
savings_ledger.record_savings_event(
|
||||
tokens_before=input_tokens,
|
||||
tokens_after=max(input_tokens - tokens_saved, 0),
|
||||
model=model,
|
||||
client=client or "proxy",
|
||||
source="proxy",
|
||||
)
|
||||
|
||||
self._get_otel_metrics().record_proxy_request(
|
||||
provider=provider,
|
||||
model=model,
|
||||
|
|
|
|||
392
headroom/savings_ledger.py
Normal file
392
headroom/savings_ledger.py
Normal file
|
|
@ -0,0 +1,392 @@
|
|||
"""Durable append-only savings event ledger.
|
||||
|
||||
Every compression — interactive ``headroom_compress`` MCP calls *and* proxy
|
||||
requests — appends a single JSON line to a file-locked JSONL ledger. Unlike the
|
||||
in-memory ``SessionStats`` and the 2-hour ``session_stats.jsonl`` window, this
|
||||
ledger survives proxy/agent restarts and is safe across concurrent writers
|
||||
(the main MCP server, each subagent's MCP server, and the proxy all append to
|
||||
the same file under an advisory lock). ``headroom savings`` aggregates it on
|
||||
read, so there is no shared mutable state to clobber and totals stay accurate.
|
||||
|
||||
Cost is computed and stored at write time so historical numbers do not drift if
|
||||
model pricing changes later. litellm list pricing is used where the model is
|
||||
known; a blended input-token rate is used as a fallback (MCP-tool compressions
|
||||
do not know the agent's upstream model, so they record ``model="unknown"`` and
|
||||
fall back to the blended rate rather than recording ``$0``).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from headroom import paths as _paths
|
||||
|
||||
# Reuse the proxy tracker's pricing + normalization so MCP and proxy events
|
||||
# bucket models identically and price them through one implementation.
|
||||
from headroom.proxy.savings_tracker import (
|
||||
_estimate_compression_savings_usd,
|
||||
_normalize_model,
|
||||
_parse_timestamp,
|
||||
sanitize_project_name,
|
||||
)
|
||||
|
||||
# fcntl is Unix-only; on Windows we skip locking (append is still best-effort).
|
||||
fcntl: Any | None = None
|
||||
try:
|
||||
import fcntl as _fcntl
|
||||
|
||||
fcntl = _fcntl
|
||||
_HAS_FCNTL = True
|
||||
except ImportError:
|
||||
_HAS_FCNTL = False
|
||||
|
||||
SCHEMA_VERSION = 1
|
||||
UNKNOWN = "unknown"
|
||||
|
||||
DEFAULT_RETENTION_DAYS = 365
|
||||
# Blended input price ($/token) used only when litellm cannot price the model.
|
||||
# Mirrors the ~$3 / 1M input-token assumption the MCP stats path already uses.
|
||||
DEFAULT_FALLBACK_INPUT_COST_PER_TOKEN = 3.0 / 1_000_000
|
||||
|
||||
# Disk hygiene: compact the ledger once it grows past this size. Retention is
|
||||
# also enforced on read, so accuracy never depends on compaction having run.
|
||||
_COMPACT_SIZE_BYTES = 8 * 1024 * 1024
|
||||
|
||||
|
||||
def _utc_now() -> datetime:
|
||||
return datetime.now(timezone.utc)
|
||||
|
||||
|
||||
def _coerce_timestamp(value: Any) -> datetime:
|
||||
if isinstance(value, datetime):
|
||||
aware = value if value.tzinfo else value.replace(tzinfo=timezone.utc)
|
||||
return aware.astimezone(timezone.utc)
|
||||
parsed = _parse_timestamp(value)
|
||||
if parsed is not None:
|
||||
return parsed
|
||||
return _utc_now()
|
||||
|
||||
|
||||
def _label(value: Any) -> str:
|
||||
"""Sanitize a free-form client label, defaulting to ``unknown``."""
|
||||
|
||||
cleaned = sanitize_project_name(value)
|
||||
return cleaned or UNKNOWN
|
||||
|
||||
|
||||
def _resolve_path(path: str | os.PathLike[str] | None) -> Path:
|
||||
return _paths.savings_events_path(path)
|
||||
|
||||
|
||||
def estimate_cost_usd(
|
||||
model: str,
|
||||
tokens_saved: int,
|
||||
*,
|
||||
fallback_rate: float = DEFAULT_FALLBACK_INPUT_COST_PER_TOKEN,
|
||||
) -> float:
|
||||
"""Dollar value of saved input tokens.
|
||||
|
||||
Uses litellm list pricing when the model resolves; otherwise falls back to a
|
||||
blended per-token rate so unknown-model traffic still accrues a non-zero
|
||||
cost-avoided figure.
|
||||
"""
|
||||
|
||||
if tokens_saved <= 0:
|
||||
return 0.0
|
||||
# Skip the litellm lookup for unknown models: it can't price them and emits
|
||||
# noisy "Provider List" warnings, and the MCP path is unknown-model by far
|
||||
# the most often. Go straight to the blended fallback.
|
||||
if model and model != UNKNOWN:
|
||||
priced = _estimate_compression_savings_usd(model, tokens_saved)
|
||||
if priced > 0:
|
||||
return round(priced, 6)
|
||||
return round(float(tokens_saved) * float(fallback_rate), 6)
|
||||
|
||||
|
||||
def record_savings_event(
|
||||
*,
|
||||
tokens_before: int,
|
||||
tokens_after: int,
|
||||
model: Any = None,
|
||||
client: Any = None,
|
||||
source: str = "mcp",
|
||||
timestamp: Any = None,
|
||||
cost_usd: float | None = None,
|
||||
fallback_rate: float = DEFAULT_FALLBACK_INPUT_COST_PER_TOKEN,
|
||||
path: str | os.PathLike[str] | None = None,
|
||||
) -> bool:
|
||||
"""Append one savings event to the durable ledger. Never raises.
|
||||
|
||||
Returns ``True`` when a line was written. ``cost_usd`` is computed from the
|
||||
model + tokens saved when not supplied by the caller.
|
||||
"""
|
||||
|
||||
try:
|
||||
before = max(int(tokens_before), 0)
|
||||
after = max(int(tokens_after), 0)
|
||||
except (TypeError, ValueError):
|
||||
return False
|
||||
|
||||
saved = max(before - after, 0)
|
||||
if saved <= 0:
|
||||
return False
|
||||
|
||||
model_label = _normalize_model(model)
|
||||
if cost_usd is None:
|
||||
cost = estimate_cost_usd(model_label, saved, fallback_rate=fallback_rate)
|
||||
else:
|
||||
try:
|
||||
cost = max(float(cost_usd), 0.0)
|
||||
except (TypeError, ValueError):
|
||||
cost = 0.0
|
||||
|
||||
event = {
|
||||
"v": SCHEMA_VERSION,
|
||||
"ts": _coerce_timestamp(timestamp).isoformat(),
|
||||
"before": before,
|
||||
"after": after,
|
||||
"saved": saved,
|
||||
"cost_usd": round(cost, 6),
|
||||
"model": model_label,
|
||||
"client": _label(client),
|
||||
"source": str(source or UNKNOWN),
|
||||
"pid": os.getpid(),
|
||||
}
|
||||
|
||||
target = _resolve_path(path)
|
||||
try:
|
||||
target.parent.mkdir(parents=True, exist_ok=True)
|
||||
line = json.dumps(event, separators=(",", ":")) + "\n"
|
||||
with open(target, "a", encoding="utf-8") as handle:
|
||||
if _HAS_FCNTL and fcntl is not None:
|
||||
fcntl.flock(handle, fcntl.LOCK_EX)
|
||||
try:
|
||||
handle.write(line)
|
||||
finally:
|
||||
if _HAS_FCNTL and fcntl is not None:
|
||||
fcntl.flock(handle, fcntl.LOCK_UN)
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
_maybe_compact(target)
|
||||
return True
|
||||
|
||||
|
||||
def _read_events(
|
||||
path: str | os.PathLike[str] | None,
|
||||
*,
|
||||
retention_days: int,
|
||||
now: datetime,
|
||||
) -> list[dict[str, Any]]:
|
||||
target = _resolve_path(path)
|
||||
if not target.exists():
|
||||
return []
|
||||
|
||||
cutoff = now - timedelta(days=retention_days) if retention_days else None
|
||||
events: list[dict[str, Any]] = []
|
||||
try:
|
||||
with open(target, encoding="utf-8") as handle:
|
||||
if _HAS_FCNTL and fcntl is not None:
|
||||
fcntl.flock(handle, fcntl.LOCK_SH)
|
||||
try:
|
||||
for raw in handle:
|
||||
raw = raw.strip()
|
||||
if not raw:
|
||||
continue
|
||||
try:
|
||||
event = json.loads(raw)
|
||||
except (ValueError, TypeError):
|
||||
continue
|
||||
parsed = _parse_timestamp(event.get("ts"))
|
||||
if parsed is None:
|
||||
continue
|
||||
if cutoff is not None and parsed < cutoff:
|
||||
continue
|
||||
event["_ts"] = parsed
|
||||
events.append(event)
|
||||
finally:
|
||||
if _HAS_FCNTL and fcntl is not None:
|
||||
fcntl.flock(handle, fcntl.LOCK_UN)
|
||||
except Exception:
|
||||
return []
|
||||
return events
|
||||
|
||||
|
||||
@dataclass
|
||||
class _Bucket:
|
||||
tokens_saved: int = 0
|
||||
tokens_before: int = 0
|
||||
cost_usd: float = 0.0
|
||||
calls: int = 0
|
||||
|
||||
def add(self, *, saved: int, before: int, cost: float) -> None:
|
||||
self.tokens_saved += saved
|
||||
self.tokens_before += before
|
||||
self.cost_usd += cost
|
||||
self.calls += 1
|
||||
|
||||
@property
|
||||
def savings_percent(self) -> float:
|
||||
if self.tokens_before <= 0:
|
||||
return 0.0
|
||||
return round(self.tokens_saved / self.tokens_before * 100, 1)
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
return {
|
||||
"tokens_saved": self.tokens_saved,
|
||||
"tokens_before": self.tokens_before,
|
||||
"cost_usd": round(self.cost_usd, 6),
|
||||
"calls": self.calls,
|
||||
"savings_percent": self.savings_percent,
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class SavingsReport:
|
||||
path: str
|
||||
schema_version: int
|
||||
lifetime: dict[str, Any]
|
||||
windows: dict[str, dict[str, Any]]
|
||||
by_model: list[dict[str, Any]]
|
||||
by_client: list[dict[str, Any]]
|
||||
top_model: str = UNKNOWN
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
return {
|
||||
"schema_version": self.schema_version,
|
||||
"path": self.path,
|
||||
"top_model": self.top_model,
|
||||
"lifetime": self.lifetime,
|
||||
"windows": self.windows,
|
||||
"by_model": self.by_model,
|
||||
"by_client": self.by_client,
|
||||
}
|
||||
|
||||
|
||||
def _ranked(buckets: dict[str, _Bucket], key_name: str) -> list[dict[str, Any]]:
|
||||
rows = []
|
||||
for name, bucket in buckets.items():
|
||||
row = {key_name: name, **bucket.to_dict()}
|
||||
rows.append(row)
|
||||
rows.sort(key=lambda r: (r["cost_usd"], r["tokens_saved"]), reverse=True)
|
||||
return rows
|
||||
|
||||
|
||||
def aggregate_savings(
|
||||
path: str | os.PathLike[str] | None = None,
|
||||
*,
|
||||
now: datetime | None = None,
|
||||
retention_days: int = DEFAULT_RETENTION_DAYS,
|
||||
) -> SavingsReport:
|
||||
"""Aggregate the durable ledger into lifetime / windowed / per-dimension views."""
|
||||
|
||||
now = now or _utc_now()
|
||||
events = _read_events(path, retention_days=retention_days, now=now)
|
||||
|
||||
# "Today" is local-calendar-day; the 7-day window is a rolling 168h.
|
||||
today_cutoff = (
|
||||
now.astimezone().replace(hour=0, minute=0, second=0, microsecond=0).astimezone(timezone.utc)
|
||||
)
|
||||
week_cutoff = now - timedelta(days=7)
|
||||
|
||||
all_time = _Bucket()
|
||||
today = _Bucket()
|
||||
last_7 = _Bucket()
|
||||
by_model: dict[str, _Bucket] = {}
|
||||
by_client: dict[str, _Bucket] = {}
|
||||
|
||||
for event in events:
|
||||
ts: datetime = event["_ts"]
|
||||
saved = max(int(event.get("saved", 0) or 0), 0)
|
||||
before = max(int(event.get("before", 0) or 0), 0)
|
||||
try:
|
||||
cost = max(float(event.get("cost_usd", 0.0) or 0.0), 0.0)
|
||||
except (TypeError, ValueError):
|
||||
cost = 0.0
|
||||
|
||||
all_time.add(saved=saved, before=before, cost=cost)
|
||||
if ts >= today_cutoff:
|
||||
today.add(saved=saved, before=before, cost=cost)
|
||||
if ts >= week_cutoff:
|
||||
last_7.add(saved=saved, before=before, cost=cost)
|
||||
|
||||
by_model.setdefault(str(event.get("model") or UNKNOWN), _Bucket()).add(
|
||||
saved=saved, before=before, cost=cost
|
||||
)
|
||||
by_client.setdefault(str(event.get("client") or UNKNOWN), _Bucket()).add(
|
||||
saved=saved, before=before, cost=cost
|
||||
)
|
||||
|
||||
model_rows = _ranked(by_model, "model")
|
||||
top_model = model_rows[0]["model"] if model_rows else UNKNOWN
|
||||
|
||||
return SavingsReport(
|
||||
path=str(_resolve_path(path)),
|
||||
schema_version=SCHEMA_VERSION,
|
||||
lifetime=all_time.to_dict(),
|
||||
windows={
|
||||
"today": today.to_dict(),
|
||||
"last_7_days": last_7.to_dict(),
|
||||
"all_time": all_time.to_dict(),
|
||||
},
|
||||
by_model=model_rows,
|
||||
by_client=_ranked(by_client, "client"),
|
||||
top_model=top_model,
|
||||
)
|
||||
|
||||
|
||||
def _maybe_compact(target: Path) -> None:
|
||||
"""Rewrite the ledger dropping out-of-retention events once it grows large."""
|
||||
|
||||
try:
|
||||
if target.stat().st_size <= _COMPACT_SIZE_BYTES:
|
||||
return
|
||||
except OSError:
|
||||
return
|
||||
|
||||
now = _utc_now()
|
||||
cutoff = now - timedelta(days=DEFAULT_RETENTION_DAYS)
|
||||
try:
|
||||
with open(target, "r+", encoding="utf-8") as handle:
|
||||
if _HAS_FCNTL and fcntl is not None:
|
||||
fcntl.flock(handle, fcntl.LOCK_EX)
|
||||
try:
|
||||
kept: list[str] = []
|
||||
handle.seek(0)
|
||||
for raw in handle:
|
||||
stripped = raw.strip()
|
||||
if not stripped:
|
||||
continue
|
||||
try:
|
||||
event = json.loads(stripped)
|
||||
except (ValueError, TypeError):
|
||||
continue
|
||||
parsed = _parse_timestamp(event.get("ts"))
|
||||
if parsed is None or parsed < cutoff:
|
||||
continue
|
||||
kept.append(stripped)
|
||||
handle.seek(0)
|
||||
handle.truncate()
|
||||
if kept:
|
||||
handle.write("\n".join(kept) + "\n")
|
||||
finally:
|
||||
if _HAS_FCNTL and fcntl is not None:
|
||||
fcntl.flock(handle, fcntl.LOCK_UN)
|
||||
except Exception:
|
||||
return
|
||||
|
||||
|
||||
__all__ = [
|
||||
"SCHEMA_VERSION",
|
||||
"DEFAULT_RETENTION_DAYS",
|
||||
"DEFAULT_FALLBACK_INPUT_COST_PER_TOKEN",
|
||||
"SavingsReport",
|
||||
"estimate_cost_usd",
|
||||
"record_savings_event",
|
||||
"aggregate_savings",
|
||||
]
|
||||
281
tests/test_savings_ledger.py
Normal file
281
tests/test_savings_ledger.py
Normal file
|
|
@ -0,0 +1,281 @@
|
|||
"""Tests for the durable savings event ledger and the `headroom savings` CLI."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from datetime import datetime, timedelta, timezone
|
||||
|
||||
import pytest
|
||||
|
||||
from headroom import savings_ledger as L
|
||||
|
||||
UTC = timezone.utc
|
||||
|
||||
|
||||
def _events_env(monkeypatch, tmp_path):
|
||||
path = tmp_path / "savings_events.jsonl"
|
||||
monkeypatch.setenv("HEADROOM_SAVINGS_EVENTS_PATH", str(path))
|
||||
return path
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# core ledger
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
def test_unknown_model_uses_blended_fallback(monkeypatch, tmp_path):
|
||||
_events_env(monkeypatch, tmp_path)
|
||||
assert L.record_savings_event(tokens_before=1000, tokens_after=400, model=None, client="c")
|
||||
report = L.aggregate_savings()
|
||||
assert report.lifetime["tokens_saved"] == 600
|
||||
expected = round(600 * L.DEFAULT_FALLBACK_INPUT_COST_PER_TOKEN, 6)
|
||||
assert report.lifetime["cost_usd"] == pytest.approx(expected)
|
||||
assert any(row["model"] == "unknown" for row in report.by_model)
|
||||
|
||||
|
||||
def test_estimate_cost_unknown_short_circuits_to_fallback():
|
||||
assert L.estimate_cost_usd("unknown", 1000, fallback_rate=1e-6) == pytest.approx(0.001)
|
||||
assert L.estimate_cost_usd(L.UNKNOWN, 0) == 0.0
|
||||
|
||||
|
||||
def test_explicit_cost_is_honored(monkeypatch, tmp_path):
|
||||
_events_env(monkeypatch, tmp_path)
|
||||
L.record_savings_event(tokens_before=100, tokens_after=10, model="x", client="c", cost_usd=1.25)
|
||||
assert L.aggregate_savings().lifetime["cost_usd"] == pytest.approx(1.25)
|
||||
|
||||
|
||||
def test_zero_or_negative_savings_not_recorded(monkeypatch, tmp_path):
|
||||
path = _events_env(monkeypatch, tmp_path)
|
||||
assert L.record_savings_event(tokens_before=100, tokens_after=100) is False
|
||||
assert L.record_savings_event(tokens_before=50, tokens_after=80) is False
|
||||
assert not path.exists() or path.read_text().strip() == ""
|
||||
assert L.aggregate_savings().lifetime["calls"] == 0
|
||||
|
||||
|
||||
def test_breakdowns_aggregate_by_dimension(monkeypatch, tmp_path):
|
||||
_events_env(monkeypatch, tmp_path)
|
||||
L.record_savings_event(tokens_before=1000, tokens_after=300, model=None, client="claude-code")
|
||||
L.record_savings_event(tokens_before=500, tokens_after=200, model=None, client="claude-code")
|
||||
L.record_savings_event(
|
||||
tokens_before=2000, tokens_after=600, model="gpt", client="proxy", cost_usd=0.5
|
||||
)
|
||||
report = L.aggregate_savings()
|
||||
|
||||
clients = {row["client"]: row for row in report.by_client}
|
||||
assert clients["claude-code"]["calls"] == 2
|
||||
assert clients["proxy"]["tokens_saved"] == 1400
|
||||
|
||||
|
||||
def test_windows_today_week_alltime(monkeypatch, tmp_path):
|
||||
_events_env(monkeypatch, tmp_path)
|
||||
now = datetime(2026, 6, 17, 12, 0, tzinfo=UTC)
|
||||
L.record_savings_event(
|
||||
tokens_before=1000, tokens_after=500, model=None, client="c", timestamp=now
|
||||
)
|
||||
L.record_savings_event(
|
||||
tokens_before=1000,
|
||||
tokens_after=600,
|
||||
model=None,
|
||||
client="c",
|
||||
timestamp=now - timedelta(days=3),
|
||||
)
|
||||
L.record_savings_event(
|
||||
tokens_before=1000,
|
||||
tokens_after=700,
|
||||
model=None,
|
||||
client="c",
|
||||
timestamp=now - timedelta(days=30),
|
||||
)
|
||||
report = L.aggregate_savings(now=now)
|
||||
assert report.windows["today"]["tokens_saved"] == 500
|
||||
assert report.windows["last_7_days"]["tokens_saved"] == 500 + 400
|
||||
assert report.windows["all_time"]["tokens_saved"] == 500 + 400 + 300
|
||||
assert report.windows["all_time"]["calls"] == 3
|
||||
# 500 saved out of 1000 before today
|
||||
assert report.windows["today"]["savings_percent"] == pytest.approx(50.0)
|
||||
|
||||
|
||||
def test_retention_excludes_old_events(monkeypatch, tmp_path):
|
||||
_events_env(monkeypatch, tmp_path)
|
||||
now = datetime(2026, 6, 17, 12, 0, tzinfo=UTC)
|
||||
L.record_savings_event(
|
||||
tokens_before=1000, tokens_after=500, model=None, client="c", timestamp=now
|
||||
)
|
||||
L.record_savings_event(
|
||||
tokens_before=1000,
|
||||
tokens_after=500,
|
||||
model=None,
|
||||
client="c",
|
||||
timestamp=now - timedelta(days=400),
|
||||
)
|
||||
report = L.aggregate_savings(now=now, retention_days=365)
|
||||
assert report.lifetime["calls"] == 1
|
||||
|
||||
|
||||
def test_appends_do_not_clobber_and_survive_restart(monkeypatch, tmp_path):
|
||||
path = _events_env(monkeypatch, tmp_path)
|
||||
for _ in range(5):
|
||||
L.record_savings_event(tokens_before=100, tokens_after=10, model=None, client="c")
|
||||
lines = [ln for ln in path.read_text().splitlines() if ln.strip()]
|
||||
assert len(lines) == 5
|
||||
# aggregate_savings holds no in-memory state — it reads purely from disk,
|
||||
# so this also proves durability across a process restart.
|
||||
assert L.aggregate_savings().lifetime["calls"] == 5
|
||||
assert L.aggregate_savings().lifetime["tokens_saved"] == 5 * 90
|
||||
|
||||
|
||||
def test_corrupt_lines_are_skipped(monkeypatch, tmp_path):
|
||||
path = _events_env(monkeypatch, tmp_path)
|
||||
L.record_savings_event(tokens_before=1000, tokens_after=400, model=None, client="c")
|
||||
with open(path, "a", encoding="utf-8") as handle:
|
||||
handle.write("not json\n\n")
|
||||
assert L.aggregate_savings().lifetime["calls"] == 1
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# CLI
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
def test_cli_reset_deletes_ledger(monkeypatch, tmp_path):
|
||||
pytest.importorskip("click")
|
||||
from click.testing import CliRunner
|
||||
|
||||
from headroom.cli.savings import savings
|
||||
|
||||
path = _events_env(monkeypatch, tmp_path)
|
||||
L.record_savings_event(tokens_before=1000, tokens_after=300, model=None, client="claude-code")
|
||||
assert path.exists()
|
||||
|
||||
result = CliRunner().invoke(savings, ["--reset"])
|
||||
assert result.exit_code == 0
|
||||
assert "reset" in result.output.lower()
|
||||
assert not path.exists()
|
||||
|
||||
# second reset on missing file is a no-op
|
||||
result2 = CliRunner().invoke(savings, ["--reset"])
|
||||
assert result2.exit_code == 0
|
||||
assert "Nothing to reset" in result2.output
|
||||
|
||||
|
||||
def test_cli_empty_state(monkeypatch, tmp_path):
|
||||
pytest.importorskip("click")
|
||||
from click.testing import CliRunner
|
||||
|
||||
from headroom.cli.savings import savings
|
||||
|
||||
_events_env(monkeypatch, tmp_path)
|
||||
result = CliRunner().invoke(savings, [])
|
||||
assert result.exit_code == 0
|
||||
assert "No savings recorded yet." in result.output
|
||||
|
||||
|
||||
def test_cli_renders_sections_and_json(monkeypatch, tmp_path):
|
||||
pytest.importorskip("click")
|
||||
from click.testing import CliRunner
|
||||
|
||||
from headroom.cli.savings import savings
|
||||
|
||||
_events_env(monkeypatch, tmp_path)
|
||||
L.record_savings_event(tokens_before=1000, tokens_after=300, model=None, client="claude-code")
|
||||
runner = CliRunner()
|
||||
|
||||
result = runner.invoke(savings, [])
|
||||
assert result.exit_code == 0
|
||||
# No redundant top-line headline; the windows lead the output.
|
||||
assert "cost avoided" not in result.output
|
||||
assert "Today" in result.output and "All time" in result.output
|
||||
assert "Savings by client" in result.output and "claude-code" in result.output
|
||||
assert "Per-repo totals" not in result.output
|
||||
|
||||
result_json = runner.invoke(savings, ["--json"])
|
||||
assert result_json.exit_code == 0
|
||||
payload = json.loads(result_json.output)
|
||||
assert payload["lifetime"]["tokens_saved"] == 700
|
||||
assert payload["windows"]["all_time"]["calls"] == 1
|
||||
assert "by_repo" not in payload
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# MCP tool path
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
def test_mcp_compress_records_durable_event(monkeypatch, tmp_path):
|
||||
pytest.importorskip("mcp", reason="MCP SDK required")
|
||||
from headroom.ccr import mcp_server
|
||||
|
||||
_events_env(monkeypatch, tmp_path)
|
||||
monkeypatch.setenv("HEADROOM_MCP_CLIENT", "claude-code")
|
||||
|
||||
server = mcp_server.HeadroomMCPServer(check_proxy=False)
|
||||
server._record_savings({"original_tokens": 1000, "compressed_tokens": 250})
|
||||
|
||||
report = L.aggregate_savings()
|
||||
assert report.lifetime["tokens_saved"] == 750
|
||||
assert {row["client"] for row in report.by_client} == {"claude-code"}
|
||||
|
||||
|
||||
def test_mcp_record_savings_ignores_noop(monkeypatch, tmp_path):
|
||||
pytest.importorskip("mcp", reason="MCP SDK required")
|
||||
from headroom.ccr import mcp_server
|
||||
|
||||
_events_env(monkeypatch, tmp_path)
|
||||
server = mcp_server.HeadroomMCPServer(check_proxy=False)
|
||||
server._record_savings({"original_tokens": 500, "compressed_tokens": 500})
|
||||
assert L.aggregate_savings().lifetime["calls"] == 0
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# proxy path
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
def test_proxy_record_request_appends_ledger_event(tmp_path, monkeypatch):
|
||||
pytest.importorskip("fastapi")
|
||||
import asyncio
|
||||
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from headroom.proxy.server import ProxyConfig, create_app
|
||||
|
||||
monkeypatch.setenv("HEADROOM_SAVINGS_PATH", str(tmp_path / "proxy_savings.json"))
|
||||
monkeypatch.setenv("HEADROOM_SAVINGS_EVENTS_PATH", str(tmp_path / "savings_events.jsonl"))
|
||||
monkeypatch.setattr(
|
||||
"headroom.proxy.server.CostTracker._get_cache_prices",
|
||||
lambda self, model: (0.001, 0.0015, 0.002),
|
||||
)
|
||||
|
||||
config = ProxyConfig(cache_enabled=False, rate_limit_enabled=False, log_requests=False)
|
||||
with TestClient(create_app(config)) as client:
|
||||
proxy = client.app.state.proxy
|
||||
# identified harness -> recorded as that client
|
||||
asyncio.run(
|
||||
proxy.metrics.record_request(
|
||||
provider="openai",
|
||||
model="gpt-4o",
|
||||
input_tokens=120,
|
||||
output_tokens=24,
|
||||
tokens_saved=40,
|
||||
latency_ms=15.0,
|
||||
client="claude-code",
|
||||
)
|
||||
)
|
||||
# unidentified harness -> falls back to "proxy"
|
||||
asyncio.run(
|
||||
proxy.metrics.record_request(
|
||||
provider="openai",
|
||||
model="gpt-4o",
|
||||
input_tokens=80,
|
||||
output_tokens=10,
|
||||
tokens_saved=20,
|
||||
latency_ms=15.0,
|
||||
)
|
||||
)
|
||||
|
||||
report = L.aggregate_savings()
|
||||
assert report.lifetime["tokens_saved"] == 60
|
||||
clients = {row["client"] for row in report.by_client}
|
||||
assert "claude-code" in clients
|
||||
assert "proxy" in clients
|
||||
assert any(row["model"] == "gpt-4o" for row in report.by_model)
|
||||
Loading…
Add table
Add a link
Reference in a new issue