mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
feat(dashboard): persist lifetime proxy metrics (#2198)
## Description Persist bounded, aggregate-only Lifetime dashboard metrics across proxy restarts and expose them through a new `/stats-lifetime` endpoint. The change keeps session/runtime stats separate from durable lifetime stats, gates sensitive dashboard metadata for loopback or explicitly trusted dashboard clients, and updates the dashboard Lifetime view to consume the new endpoint. Closes #2137 ## 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 - [x] Code refactoring (no functional changes) ## Changes Made - Added bounded persistent lifetime metrics state and wired proxy metric events into it. - Added `/stats-lifetime` with sensitive project/persistence details gated behind dashboard metadata access checks. - Extended loopback/dashboard metadata access policy for trusted dashboard client CIDRs without widening admin/debug endpoints. - Reorganized dashboard session/lifetime presentation around runtime counters versus durable aggregates. - Added focused tests for persistent aggregation, persistence, endpoint registration, loopback gating, trusted dashboard CIDRs, and recent request ordering. - Fixed current Ruff/mypy issues in the lifetime metrics normalization code. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check`) - [x] Type checking passes (`mypy`) - [x] New tests added for new functionality - [x] Manual testing performed ### Test Output ```text uv run --extra proxy --with pytest --with pytest-asyncio pytest tests/test_persistent_metrics.py tests/test_persistent_metrics_integration.py tests/test_persistent_metrics_persistence.py tests/test_proxy_loopback_gating.py tests/test_proxy_stats_recent_requests.py -q 53 passed, 1 warning uvx ruff==0.15.17 check headroom/proxy/forwarded_headers.py headroom/proxy/loopback_guard.py headroom/proxy/persistent_metrics.py headroom/proxy/savings_tracker.py headroom/proxy/server.py tests/test_forwarded_headers.py tests/test_persistent_metrics.py tests/test_persistent_metrics_integration.py tests/test_persistent_metrics_persistence.py tests/test_proxy_loopback_gating.py tests/test_proxy_project_savings.py tests/test_proxy_stats_recent_requests.py All checks passed! uv run --extra proxy --with mypy mypy headroom/proxy/persistent_metrics.py Success: no issues found in 1 source file ``` ## Real Behavior Proof - Environment: Windows review worktree, Python 3.13.3 via uv. - Exact command / steps: Ran the focused persistent metrics, persistence, loopback gating, and recent request tests; ran CI-matching Ruff on touched files; ran mypy on the new persistent metrics module. - Observed result: `/stats-lifetime` is registered, non-loopback callers receive only non-sensitive aggregate data, loopback/trusted dashboard clients receive the full lifetime payload, admin/debug endpoints remain loopback-only, and persistent metrics normalize malformed stored state without type/lint errors. - Not tested: Full repository pytest suite, full dashboard browser screenshot pass, or live long-running proxy traffic. ## 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 ## Screenshots (if applicable) N/A ## Additional Notes No changelog entry is required for this dashboard/internal metrics iteration. The endpoint intentionally exposes only aggregate lifetime data to ordinary network callers and strips project/persistence details unless the caller passes the dashboard metadata access policy. --------- Co-authored-by: JerrettDavis <mxjerrett@gmail.com>
This commit is contained in:
parent
551f473e04
commit
0537cbfde4
15 changed files with 2656 additions and 1248 deletions
2
.gitignore
vendored
2
.gitignore
vendored
|
|
@ -261,3 +261,5 @@ uv.lock
|
|||
/headroom/_core.*.so
|
||||
/headroom/_core.so
|
||||
.tokensave
|
||||
|
||||
.codebase-memory/
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -72,7 +72,9 @@ if TYPE_CHECKING:
|
|||
logger = logging.getLogger(__name__)
|
||||
|
||||
__all__ = [
|
||||
"TRUSTED_DASHBOARD_CLIENT_CIDRS_ENV",
|
||||
"TRUSTED_GATEWAY_CIDRS_ENV",
|
||||
"load_trusted_dashboard_client_cidrs",
|
||||
"load_trusted_gateway_cidrs",
|
||||
"peer_is_trusted_gateway",
|
||||
"resolve_client_ip",
|
||||
|
|
@ -82,6 +84,7 @@ __all__ = [
|
|||
|
||||
#: Environment variable that holds the comma-separated CIDR allow-list.
|
||||
TRUSTED_GATEWAY_CIDRS_ENV = "HEADROOM_PROXY_TRUSTED_GATEWAY_CIDRS"
|
||||
TRUSTED_DASHBOARD_CLIENT_CIDRS_ENV = "HEADROOM_PROXY_TRUSTED_DASHBOARD_CLIENT_CIDRS"
|
||||
|
||||
|
||||
def _parse_cidr_list(
|
||||
|
|
@ -115,6 +118,20 @@ def load_trusted_gateway_cidrs(
|
|||
return _parse_cidr_list(raw)
|
||||
|
||||
|
||||
def load_trusted_dashboard_client_cidrs(
|
||||
raw: str | None = None,
|
||||
) -> tuple[ipaddress.IPv4Network | ipaddress.IPv6Network, ...]:
|
||||
"""Parse the Dashboard client CIDR allow-list from its environment variable."""
|
||||
if raw is None:
|
||||
raw = os.environ.get(TRUSTED_DASHBOARD_CLIENT_CIDRS_ENV, "")
|
||||
try:
|
||||
return _parse_cidr_list(raw)
|
||||
except ValueError as exc:
|
||||
raise ValueError(
|
||||
f"Invalid {TRUSTED_DASHBOARD_CLIENT_CIDRS_ENV} entry: {exc}"
|
||||
) from exc
|
||||
|
||||
|
||||
def _normalize_ip(
|
||||
host: str,
|
||||
) -> ipaddress.IPv4Address | ipaddress.IPv6Address | None:
|
||||
|
|
|
|||
|
|
@ -50,6 +50,7 @@ except ImportError: # pragma: no cover - fastapi is a hard dep in practice
|
|||
|
||||
__all__ = [
|
||||
"LOOPBACK_HOSTS",
|
||||
"is_ip_literal_host_header",
|
||||
"is_loopback_host",
|
||||
"is_loopback_host_header",
|
||||
"require_loopback",
|
||||
|
|
@ -128,6 +129,46 @@ def is_loopback_host_header(header_value: str | None) -> bool:
|
|||
return is_loopback_host(host_part)
|
||||
|
||||
|
||||
def is_ip_literal_host_header(header_value: str | None) -> bool:
|
||||
"""Return whether ``Host:`` contains an IPv4 or bracketed IPv6 literal.
|
||||
|
||||
Dashboard clients may use a non-loopback server address, but retaining an
|
||||
IP-literal Host requirement prevents DNS-rebinding requests from using an
|
||||
attacker-controlled hostname. Ports are accepted in normal HTTP forms.
|
||||
"""
|
||||
if not header_value:
|
||||
return False
|
||||
|
||||
candidate = header_value.strip()
|
||||
if not candidate or "/" in candidate or "@" in candidate:
|
||||
return False
|
||||
|
||||
if candidate.startswith("["):
|
||||
closing = candidate.find("]")
|
||||
if closing == -1 or candidate.count("[") != 1 or candidate.count("]") != 1:
|
||||
return False
|
||||
host_part = candidate[1:closing]
|
||||
suffix = candidate[closing + 1 :]
|
||||
if suffix and (not suffix.startswith(":") or not suffix[1:].isdigit()):
|
||||
return False
|
||||
try:
|
||||
return isinstance(ipaddress.ip_address(host_part), ipaddress.IPv6Address)
|
||||
except ValueError:
|
||||
return False
|
||||
|
||||
if candidate.count(":") == 1:
|
||||
host_part, port = candidate.rsplit(":", 1)
|
||||
if not port.isdigit():
|
||||
return False
|
||||
else:
|
||||
host_part = candidate
|
||||
|
||||
try:
|
||||
return isinstance(ipaddress.ip_address(host_part), ipaddress.IPv4Address)
|
||||
except ValueError:
|
||||
return False
|
||||
|
||||
|
||||
def require_loopback(request: Request) -> None: # type: ignore[valid-type]
|
||||
"""FastAPI dependency: 404 any non-loopback caller.
|
||||
|
||||
|
|
|
|||
470
headroom/proxy/persistent_metrics.py
Normal file
470
headroom/proxy/persistent_metrics.py
Normal file
|
|
@ -0,0 +1,470 @@
|
|||
"""Pure, bounded aggregate state for the durable Dashboard Lifetime view."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from collections.abc import Callable
|
||||
from copy import deepcopy
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
|
||||
SCHEMA_VERSION = 5
|
||||
MAX_PROVIDER_VALUES = 32
|
||||
MAX_STACK_VALUES = 64
|
||||
MAX_TRACKED_MODELS = 200
|
||||
MAX_EXPOSED_MODELS = 100
|
||||
MAX_LABEL_LENGTH = 128
|
||||
|
||||
KNOWN_MISS_REASONS = frozenset({"ttl_expiry", "prefix_change", "unknown"})
|
||||
KNOWN_WASTE_SIGNALS = frozenset(
|
||||
{
|
||||
"json_noise",
|
||||
"html_noise",
|
||||
"base64",
|
||||
"whitespace",
|
||||
"dynamic_date",
|
||||
"repetition",
|
||||
"reread",
|
||||
"reread_compressed",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def utc_now() -> datetime:
|
||||
"""Return the current UTC time without sub-second noise in persisted state."""
|
||||
|
||||
return datetime.now(timezone.utc).replace(microsecond=0)
|
||||
|
||||
|
||||
def _to_iso(value: datetime | None) -> str | None:
|
||||
if value is None:
|
||||
return None
|
||||
return value.astimezone(timezone.utc).replace(microsecond=0).isoformat().replace("+00:00", "Z")
|
||||
|
||||
|
||||
def _coerce_int(value: Any) -> int:
|
||||
try:
|
||||
return max(int(value), 0)
|
||||
except (TypeError, ValueError, OverflowError):
|
||||
return 0
|
||||
|
||||
|
||||
def _coerce_float(value: Any) -> float:
|
||||
try:
|
||||
result = float(value)
|
||||
except (TypeError, ValueError, OverflowError):
|
||||
return 0.0
|
||||
return result if math.isfinite(result) and result >= 0 else 0.0
|
||||
|
||||
|
||||
def _label(value: Any) -> str:
|
||||
if not isinstance(value, str):
|
||||
return "other"
|
||||
value = value.strip()
|
||||
return value[:MAX_LABEL_LENGTH] if value else "other"
|
||||
|
||||
|
||||
def _model_entry(raw: Any = None) -> dict[str, Any]:
|
||||
raw = raw if isinstance(raw, dict) else {}
|
||||
return {
|
||||
"requests": _coerce_int(raw.get("requests")),
|
||||
"input_tokens": _coerce_int(raw.get("input_tokens")),
|
||||
"output_tokens": _coerce_int(raw.get("output_tokens")),
|
||||
"attempted_input_tokens": _coerce_int(raw.get("attempted_input_tokens")),
|
||||
"tokens_saved": _coerce_int(raw.get("tokens_saved")),
|
||||
"last_activity_at": raw.get("last_activity_at")
|
||||
if isinstance(raw.get("last_activity_at"), str)
|
||||
else None,
|
||||
}
|
||||
|
||||
|
||||
def _empty_state() -> dict[str, Any]:
|
||||
return {
|
||||
"started_at": None,
|
||||
"last_activity_at": None,
|
||||
"full_fidelity_started_at": None,
|
||||
"requests": {
|
||||
"total": 0,
|
||||
"cached": 0,
|
||||
"failed": 0,
|
||||
"rate_limited": 0,
|
||||
"by_provider": {},
|
||||
"by_stack": {},
|
||||
},
|
||||
"tokens": {"input": 0, "output": 0, "attempted_input": 0, "saved": 0},
|
||||
"prefix_cache": {
|
||||
"requests": 0,
|
||||
"hit_requests": 0,
|
||||
"cache_read_tokens": 0,
|
||||
"cache_write_tokens": 0,
|
||||
"cache_write_5m_tokens": 0,
|
||||
"cache_write_1h_tokens": 0,
|
||||
"uncached_input_tokens": 0,
|
||||
"bust_count": 0,
|
||||
"bust_tokens": 0,
|
||||
"misses_by_reason": {},
|
||||
"by_provider": {},
|
||||
},
|
||||
"cost": {"input_usd": 0.0, "compression_savings_usd": 0.0, "cache_savings_usd": 0.0},
|
||||
"waste_signals": {},
|
||||
"models": {"tracked": {}, "other": _model_entry()},
|
||||
"persistence": {"last_saved_at": None},
|
||||
}
|
||||
|
||||
|
||||
def _dict_or_empty(value: Any) -> dict[Any, Any]:
|
||||
return value if isinstance(value, dict) else {}
|
||||
|
||||
|
||||
class PersistentMetricsState:
|
||||
"""In-memory Lifetime aggregate with deterministic, bounded dimensions."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
raw: dict[str, Any] | None = None,
|
||||
*,
|
||||
now: Callable[[], datetime] = utc_now,
|
||||
) -> None:
|
||||
self._now = now
|
||||
self._state = self._normalize(raw)
|
||||
self._compact_models()
|
||||
|
||||
def _normalize(self, raw: dict[str, Any] | None) -> dict[str, Any]:
|
||||
source = raw if isinstance(raw, dict) else {}
|
||||
result = _empty_state()
|
||||
for timestamp_key in ("started_at", "last_activity_at", "full_fidelity_started_at"):
|
||||
value = source.get(timestamp_key)
|
||||
if isinstance(value, str):
|
||||
result[timestamp_key] = value
|
||||
|
||||
raw_requests = _dict_or_empty(source.get("requests"))
|
||||
for key in ("total", "cached", "failed", "rate_limited"):
|
||||
result["requests"][key] = _coerce_int(raw_requests.get(key))
|
||||
result["requests"]["by_provider"] = self._normalize_count_map(
|
||||
raw_requests.get("by_provider"), MAX_PROVIDER_VALUES
|
||||
)
|
||||
result["requests"]["by_stack"] = self._normalize_count_map(
|
||||
raw_requests.get("by_stack"), MAX_STACK_VALUES
|
||||
)
|
||||
|
||||
raw_tokens = _dict_or_empty(source.get("tokens"))
|
||||
for key in ("input", "output", "attempted_input", "saved"):
|
||||
result["tokens"][key] = _coerce_int(raw_tokens.get(key))
|
||||
|
||||
raw_cache = _dict_or_empty(source.get("prefix_cache"))
|
||||
for key in (
|
||||
"requests",
|
||||
"hit_requests",
|
||||
"cache_read_tokens",
|
||||
"cache_write_tokens",
|
||||
"cache_write_5m_tokens",
|
||||
"cache_write_1h_tokens",
|
||||
"uncached_input_tokens",
|
||||
"bust_count",
|
||||
"bust_tokens",
|
||||
):
|
||||
result["prefix_cache"][key] = _coerce_int(raw_cache.get(key))
|
||||
result["prefix_cache"]["by_provider"] = self._normalize_count_map(
|
||||
raw_cache.get("by_provider"), MAX_PROVIDER_VALUES
|
||||
)
|
||||
result["prefix_cache"]["misses_by_reason"] = self._normalize_enum_map(
|
||||
raw_cache.get("misses_by_reason"), KNOWN_MISS_REASONS
|
||||
)
|
||||
|
||||
raw_cost = _dict_or_empty(source.get("cost"))
|
||||
for key in ("input_usd", "compression_savings_usd", "cache_savings_usd"):
|
||||
result["cost"][key] = round(_coerce_float(raw_cost.get(key)), 6)
|
||||
result["waste_signals"] = self._normalize_enum_map(
|
||||
source.get("waste_signals"), KNOWN_WASTE_SIGNALS
|
||||
)
|
||||
|
||||
raw_models = _dict_or_empty(source.get("models"))
|
||||
raw_tracked = _dict_or_empty(raw_models.get("tracked"))
|
||||
for name, entry in raw_tracked.items():
|
||||
normalized_name = self._model_name(name)
|
||||
if normalized_name == "other":
|
||||
self._merge_model_entry(result["models"]["other"], _model_entry(entry))
|
||||
continue
|
||||
result["models"]["tracked"][normalized_name] = _model_entry(entry)
|
||||
self._merge_model_entry(result["models"]["other"], _model_entry(raw_models.get("other")))
|
||||
raw_persistence = _dict_or_empty(source.get("persistence"))
|
||||
if isinstance(raw_persistence.get("last_saved_at"), str):
|
||||
result["persistence"]["last_saved_at"] = raw_persistence["last_saved_at"]
|
||||
return result
|
||||
|
||||
@staticmethod
|
||||
def _normalize_count_map(raw: Any, limit: int) -> dict[str, int]:
|
||||
result: dict[str, int] = {}
|
||||
if not isinstance(raw, dict):
|
||||
return result
|
||||
for key, value in raw.items():
|
||||
label = _label(key)
|
||||
result[label] = result.get(label, 0) + _coerce_int(value)
|
||||
PersistentMetricsState._compact_count_map(result, limit)
|
||||
return result
|
||||
|
||||
@staticmethod
|
||||
def _normalize_enum_map(raw: Any, allowed: frozenset[str]) -> dict[str, int]:
|
||||
result: dict[str, int] = {}
|
||||
if not isinstance(raw, dict):
|
||||
return result
|
||||
for key, value in raw.items():
|
||||
label = key if isinstance(key, str) and key in allowed else "unknown"
|
||||
result[label] = result.get(label, 0) + _coerce_int(value)
|
||||
return result
|
||||
|
||||
@staticmethod
|
||||
def _compact_count_map(values: dict[str, int], limit: int) -> None:
|
||||
named = [key for key in values if key != "other"]
|
||||
while len(named) > limit:
|
||||
evicted = min(named, key=lambda key: (values[key], key))
|
||||
values["other"] = values.get("other", 0) + values.pop(evicted)
|
||||
named.remove(evicted)
|
||||
|
||||
def _record_activity(self) -> str:
|
||||
timestamp = _to_iso(self._now())
|
||||
if self._state["started_at"] is None:
|
||||
self._state["started_at"] = timestamp
|
||||
if self._state["full_fidelity_started_at"] is None:
|
||||
self._state["full_fidelity_started_at"] = timestamp
|
||||
self._state["last_activity_at"] = timestamp
|
||||
return timestamp or ""
|
||||
|
||||
@staticmethod
|
||||
def _increment_count(values: dict[str, int], label: str, limit: int) -> None:
|
||||
values[label] = values.get(label, 0) + 1
|
||||
PersistentMetricsState._compact_count_map(values, limit)
|
||||
|
||||
@staticmethod
|
||||
def _model_name(value: Any) -> str:
|
||||
"""Never retain the legacy unknown model bucket as a named model."""
|
||||
|
||||
label = _label(value)
|
||||
return "other" if label.lower() == "unknown" else label
|
||||
|
||||
@staticmethod
|
||||
def _merge_model_entry(destination: dict[str, Any], source: dict[str, Any]) -> None:
|
||||
for key in (
|
||||
"requests",
|
||||
"input_tokens",
|
||||
"output_tokens",
|
||||
"attempted_input_tokens",
|
||||
"tokens_saved",
|
||||
):
|
||||
destination[key] += _coerce_int(source.get(key))
|
||||
if destination["last_activity_at"] is None or (
|
||||
source["last_activity_at"] is not None
|
||||
and source["last_activity_at"] > destination["last_activity_at"]
|
||||
):
|
||||
destination["last_activity_at"] = source["last_activity_at"]
|
||||
|
||||
@staticmethod
|
||||
def _model_rank(item: tuple[str, dict[str, Any]]) -> tuple[int, str, str]:
|
||||
name, entry = item
|
||||
observed_tokens = entry["input_tokens"] + entry["output_tokens"]
|
||||
return (-observed_tokens, entry["last_activity_at"] or "", name)
|
||||
|
||||
def _compact_models(self) -> None:
|
||||
tracked = self._state["models"]["tracked"]
|
||||
if len(tracked) <= MAX_TRACKED_MODELS:
|
||||
return
|
||||
ranked = sorted(tracked.items(), key=self._model_rank)
|
||||
kept = dict(ranked[:MAX_EXPOSED_MODELS])
|
||||
other = self._state["models"]["other"]
|
||||
for _, entry in ranked[MAX_EXPOSED_MODELS:]:
|
||||
self._merge_model_entry(other, entry)
|
||||
self._state["models"]["tracked"] = kept
|
||||
|
||||
def _record_model(
|
||||
self,
|
||||
*,
|
||||
model: str | None,
|
||||
timestamp: str,
|
||||
input_tokens: int,
|
||||
output_tokens: int,
|
||||
attempted_input_tokens: int,
|
||||
tokens_saved: int,
|
||||
) -> None:
|
||||
name = self._model_name(model)
|
||||
models = self._state["models"]
|
||||
entry = (
|
||||
models["other"]
|
||||
if name == "other"
|
||||
else models["tracked"].setdefault(name, _model_entry())
|
||||
)
|
||||
entry["requests"] += 1
|
||||
entry["input_tokens"] += input_tokens
|
||||
entry["output_tokens"] += output_tokens
|
||||
entry["attempted_input_tokens"] += attempted_input_tokens
|
||||
entry["tokens_saved"] += tokens_saved
|
||||
entry["last_activity_at"] = timestamp
|
||||
self._compact_models()
|
||||
|
||||
def record_request(
|
||||
self,
|
||||
*,
|
||||
provider: str | None,
|
||||
stack: str | None,
|
||||
model: str | None,
|
||||
input_tokens: Any = 0,
|
||||
output_tokens: Any = 0,
|
||||
attempted_input_tokens: Any = 0,
|
||||
tokens_saved: Any = 0,
|
||||
cached: bool = False,
|
||||
record_stack: bool = True,
|
||||
cache_read_tokens: Any = 0,
|
||||
cache_write_tokens: Any = 0,
|
||||
cache_write_5m_tokens: Any = 0,
|
||||
cache_write_1h_tokens: Any = 0,
|
||||
uncached_input_tokens: Any = 0,
|
||||
input_usd: Any = 0.0,
|
||||
compression_savings_usd: Any = 0.0,
|
||||
cache_savings_usd: Any = 0.0,
|
||||
waste_signals: dict[str, Any] | None = None,
|
||||
) -> None:
|
||||
"""Accumulate one completed top-level request after coercing all deltas."""
|
||||
|
||||
timestamp = self._record_activity()
|
||||
input_delta = _coerce_int(input_tokens)
|
||||
output_delta = _coerce_int(output_tokens)
|
||||
attempted_delta = _coerce_int(attempted_input_tokens)
|
||||
saved_delta = _coerce_int(tokens_saved)
|
||||
provider_label = _label(provider)
|
||||
stack_label = _label(stack)
|
||||
|
||||
requests = self._state["requests"]
|
||||
requests["total"] += 1
|
||||
requests["cached"] += int(bool(cached))
|
||||
self._increment_count(requests["by_provider"], provider_label, MAX_PROVIDER_VALUES)
|
||||
if record_stack:
|
||||
self._increment_count(requests["by_stack"], stack_label, MAX_STACK_VALUES)
|
||||
|
||||
tokens = self._state["tokens"]
|
||||
tokens["input"] += input_delta
|
||||
tokens["output"] += output_delta
|
||||
tokens["attempted_input"] += attempted_delta
|
||||
tokens["saved"] += saved_delta
|
||||
|
||||
cache = self._state["prefix_cache"]
|
||||
cache["requests"] += 1
|
||||
cache["hit_requests"] += int(bool(cached))
|
||||
cache["cache_read_tokens"] += _coerce_int(cache_read_tokens)
|
||||
cache["cache_write_tokens"] += _coerce_int(cache_write_tokens)
|
||||
cache["cache_write_5m_tokens"] += _coerce_int(cache_write_5m_tokens)
|
||||
cache["cache_write_1h_tokens"] += _coerce_int(cache_write_1h_tokens)
|
||||
cache["uncached_input_tokens"] += _coerce_int(uncached_input_tokens)
|
||||
self._increment_count(cache["by_provider"], provider_label, MAX_PROVIDER_VALUES)
|
||||
|
||||
cost = self._state["cost"]
|
||||
cost["input_usd"] = round(cost["input_usd"] + _coerce_float(input_usd), 6)
|
||||
cost["compression_savings_usd"] = round(
|
||||
cost["compression_savings_usd"] + _coerce_float(compression_savings_usd), 6
|
||||
)
|
||||
cost["cache_savings_usd"] = round(
|
||||
cost["cache_savings_usd"] + _coerce_float(cache_savings_usd), 6
|
||||
)
|
||||
|
||||
if isinstance(waste_signals, dict):
|
||||
for name, token_count in waste_signals.items():
|
||||
bucket = name if isinstance(name, str) and name in KNOWN_WASTE_SIGNALS else "other"
|
||||
self._state["waste_signals"][bucket] = self._state["waste_signals"].get(
|
||||
bucket, 0
|
||||
) + _coerce_int(token_count)
|
||||
|
||||
self._record_model(
|
||||
model=model,
|
||||
timestamp=timestamp,
|
||||
input_tokens=input_delta,
|
||||
output_tokens=output_delta,
|
||||
attempted_input_tokens=attempted_delta,
|
||||
tokens_saved=saved_delta,
|
||||
)
|
||||
|
||||
def record_stack(self, stack: str | None) -> None:
|
||||
"""Accumulate the existing inbound stack label without adding a request."""
|
||||
|
||||
if stack is None:
|
||||
return
|
||||
self._record_activity()
|
||||
self._increment_count(self._state["requests"]["by_stack"], _label(stack), MAX_STACK_VALUES)
|
||||
|
||||
def record_failed(self, *, provider: str | None = None, model: str | None = None) -> None:
|
||||
"""Record a failed request without changing the completed-request denominator."""
|
||||
|
||||
self._record_activity()
|
||||
self._state["requests"]["failed"] += 1
|
||||
|
||||
def record_rate_limited(self, *, provider: str | None = None, model: str | None = None) -> None:
|
||||
"""Record a rate-limited request without redefining total request semantics."""
|
||||
|
||||
self._record_activity()
|
||||
self._state["requests"]["rate_limited"] += 1
|
||||
|
||||
def record_cache_bust(self, *, tokens_lost: Any = 0) -> None:
|
||||
self._record_activity()
|
||||
cache = self._state["prefix_cache"]
|
||||
cache["bust_count"] += 1
|
||||
cache["bust_tokens"] += _coerce_int(tokens_lost)
|
||||
|
||||
def record_cache_miss(self, *, provider: str | None, reason: str | None) -> None:
|
||||
self._record_activity()
|
||||
bucket = reason if isinstance(reason, str) and reason in KNOWN_MISS_REASONS else "unknown"
|
||||
misses = self._state["prefix_cache"]["misses_by_reason"]
|
||||
misses[bucket] = misses.get(bucket, 0) + 1
|
||||
|
||||
def set_last_saved_at(self, value: str | None) -> None:
|
||||
self._state["persistence"]["last_saved_at"] = value
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
"""Return the persisted form without derived values or I/O metadata."""
|
||||
|
||||
return deepcopy(self._state)
|
||||
|
||||
@staticmethod
|
||||
def _percent(numerator: int, denominator: int) -> float | None:
|
||||
if denominator <= 0:
|
||||
return None
|
||||
return round(numerator / denominator * 100, 6)
|
||||
|
||||
def _by_model_snapshot(self) -> dict[str, dict[str, Any]]:
|
||||
ranked = sorted(self._state["models"]["tracked"].items(), key=self._model_rank)
|
||||
visible = ranked[:MAX_EXPOSED_MODELS]
|
||||
other = _model_entry(self._state["models"]["other"])
|
||||
for _, entry in ranked[MAX_EXPOSED_MODELS:]:
|
||||
self._merge_model_entry(other, entry)
|
||||
result = {name: deepcopy(entry) for name, entry in visible}
|
||||
result["other"] = other
|
||||
return result
|
||||
|
||||
def snapshot(self, *, persistence: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Return an API-safe aggregate and derive all percentages at read time."""
|
||||
|
||||
cache = self._state["prefix_cache"]
|
||||
cache_write_total = cache["cache_write_1h_tokens"] + cache["cache_write_5m_tokens"]
|
||||
tokens = self._state["tokens"]
|
||||
return {
|
||||
"scope": "lifetime",
|
||||
"schema_version": SCHEMA_VERSION,
|
||||
"generated_at": _to_iso(self._now()),
|
||||
"started_at": self._state["started_at"],
|
||||
"last_activity_at": self._state["last_activity_at"],
|
||||
"full_fidelity_started_at": self._state["full_fidelity_started_at"],
|
||||
"requests": deepcopy(self._state["requests"]),
|
||||
"tokens": {
|
||||
**deepcopy(tokens),
|
||||
"token_savings_percent": self._percent(tokens["saved"], tokens["attempted_input"]),
|
||||
},
|
||||
"prefix_cache": {
|
||||
**deepcopy(cache),
|
||||
"cache_hit_rate": self._percent(cache["hit_requests"], cache["requests"]),
|
||||
"ttl_1h_percent": self._percent(cache["cache_write_1h_tokens"], cache_write_total),
|
||||
"ttl_5m_percent": self._percent(cache["cache_write_5m_tokens"], cache_write_total),
|
||||
},
|
||||
"cost": deepcopy(self._state["cost"]),
|
||||
"waste_signals": deepcopy(self._state["waste_signals"]),
|
||||
"by_model": self._by_model_snapshot(),
|
||||
"persistence": {
|
||||
**deepcopy(persistence),
|
||||
"last_saved_at": self._state["persistence"]["last_saved_at"],
|
||||
},
|
||||
}
|
||||
|
|
@ -446,6 +446,7 @@ class PrometheusMetrics:
|
|||
):
|
||||
return
|
||||
self.requests_by_stack[slug] += 1
|
||||
self.savings_tracker.record_lifetime_stack(slug)
|
||||
|
||||
def record_compression(
|
||||
self,
|
||||
|
|
@ -732,6 +733,24 @@ class PrometheusMetrics:
|
|||
if len(self.savings_history) > 500:
|
||||
self.savings_history = self.savings_history[-500:]
|
||||
|
||||
self.savings_tracker.record_lifetime_request(
|
||||
persist=False,
|
||||
provider=provider,
|
||||
stack=None,
|
||||
record_stack=False,
|
||||
model=model,
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
attempted_input_tokens=attempted_input_tokens,
|
||||
tokens_saved=tokens_saved,
|
||||
cached=cached,
|
||||
cache_read_tokens=cache_read_tokens,
|
||||
cache_write_tokens=cache_write_tokens,
|
||||
cache_write_5m_tokens=cache_write_5m_tokens,
|
||||
cache_write_1h_tokens=cache_write_1h_tokens,
|
||||
uncached_input_tokens=uncached_input_tokens,
|
||||
waste_signals=waste_signals,
|
||||
)
|
||||
total_input_tokens, total_input_cost_usd = self._current_savings_tracker_totals()
|
||||
self.savings_tracker.record_request(
|
||||
model=model,
|
||||
|
|
@ -828,6 +847,7 @@ class PrometheusMetrics:
|
|||
async with self._lock:
|
||||
self.cache_bust_tokens_lost += tokens_lost
|
||||
self.cache_bust_count += 1
|
||||
self.savings_tracker.record_lifetime_cache_bust(tokens_lost=tokens_lost)
|
||||
self._get_otel_metrics().record_proxy_cache_bust(tokens_lost=tokens_lost)
|
||||
|
||||
async def record_cache_miss_attribution(self, provider: str, reason: str) -> None:
|
||||
|
|
@ -841,6 +861,7 @@ class PrometheusMetrics:
|
|||
"""
|
||||
async with self._lock:
|
||||
self.cache_miss_attribution_by_provider[provider][reason] += 1
|
||||
self.savings_tracker.record_lifetime_cache_miss(provider=provider, reason=reason)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Unit 3: WS session lifecycle gauges / histogram
|
||||
|
|
@ -886,11 +907,13 @@ class PrometheusMetrics:
|
|||
async def record_rate_limited(self, *, provider: str | None = None, model: str | None = None):
|
||||
async with self._lock:
|
||||
self.requests_rate_limited += 1
|
||||
self.savings_tracker.record_lifetime_rate_limited(provider=provider, model=model)
|
||||
self._get_otel_metrics().record_proxy_rate_limited(provider=provider, model=model)
|
||||
|
||||
async def record_failed(self, *, provider: str | None = None, model: str | None = None):
|
||||
async with self._lock:
|
||||
self.requests_failed += 1
|
||||
self.savings_tracker.record_lifetime_failed(provider=provider, model=model)
|
||||
self._get_otel_metrics().record_proxy_failed(provider=provider, model=model)
|
||||
|
||||
async def export(self) -> str:
|
||||
|
|
|
|||
|
|
@ -22,6 +22,7 @@ from typing import Any
|
|||
|
||||
from headroom import paths as _paths
|
||||
from headroom.proxy import project_name_policy
|
||||
from headroom.proxy.persistent_metrics import PersistentMetricsState
|
||||
|
||||
PROJECT_NAME_MAX_LENGTH = project_name_policy.PROJECT_NAME_MAX_LENGTH
|
||||
sanitize_project_name = project_name_policy.sanitize_project_name
|
||||
|
|
@ -31,7 +32,7 @@ logger = logging.getLogger(__name__)
|
|||
HEADROOM_SAVINGS_PATH_ENV_VAR = _paths.HEADROOM_SAVINGS_PATH_ENV
|
||||
DEFAULT_SAVINGS_DIR = ".headroom"
|
||||
DEFAULT_SAVINGS_FILE = "proxy_savings.json"
|
||||
SCHEMA_VERSION = 4
|
||||
SCHEMA_VERSION = 5
|
||||
DEFAULT_MAX_HISTORY_POINTS = 5000
|
||||
DEFAULT_MAX_PROJECTS = 50
|
||||
DEFAULT_MAX_HISTORY_AGE_DAYS = 365
|
||||
|
|
@ -546,7 +547,13 @@ class SavingsTracker:
|
|||
self._save_flush_every = max(_coerce_int(save_flush_every, 1), 1)
|
||||
self._since_save = 0
|
||||
self._lock = threading.Lock()
|
||||
self._persistence_healthy = True
|
||||
self._persistence_error: str | None = None
|
||||
self._needs_schema_save = False
|
||||
self._state = self._load_state()
|
||||
self._persistent_metrics = PersistentMetricsState(
|
||||
self._state.pop("lifetime_metrics", None)
|
||||
)
|
||||
|
||||
@property
|
||||
def storage_path(self) -> str:
|
||||
|
|
@ -784,6 +791,66 @@ class SavingsTracker:
|
|||
self._maybe_save_locked()
|
||||
return True
|
||||
|
||||
def record_lifetime_request(self, *, persist: bool = True, **metrics: Any) -> None:
|
||||
"""Record one completed request in the durable Lifetime aggregate."""
|
||||
|
||||
model = _normalize_model(metrics.get("model"))
|
||||
input_tokens = _coerce_int(metrics.get("input_tokens"))
|
||||
cache_read_tokens = _coerce_int(metrics.get("cache_read_tokens"))
|
||||
cache_write_tokens = _coerce_int(metrics.get("cache_write_tokens"))
|
||||
uncached_input_tokens = _coerce_int(metrics.get("uncached_input_tokens"))
|
||||
metrics.setdefault(
|
||||
"input_usd",
|
||||
_estimate_input_cost_usd(
|
||||
model,
|
||||
input_tokens,
|
||||
cache_read_tokens=cache_read_tokens,
|
||||
cache_write_tokens=cache_write_tokens,
|
||||
uncached_input_tokens=uncached_input_tokens,
|
||||
),
|
||||
)
|
||||
metrics.setdefault(
|
||||
"compression_savings_usd",
|
||||
_estimate_compression_savings_usd(model, _coerce_int(metrics.get("tokens_saved"))),
|
||||
)
|
||||
metrics.setdefault("cache_savings_usd", _estimate_cache_savings_usd(model, cache_read_tokens))
|
||||
with self._lock:
|
||||
self._persistent_metrics.record_request(**metrics)
|
||||
if persist:
|
||||
self._maybe_save_locked()
|
||||
|
||||
def record_lifetime_stack(self, stack: str | None) -> None:
|
||||
"""Mirror the existing inbound stack dimension without another save."""
|
||||
|
||||
with self._lock:
|
||||
self._persistent_metrics.record_stack(stack)
|
||||
|
||||
def record_lifetime_failed(self, *, provider: str | None = None, model: str | None = None) -> None:
|
||||
"""Record a failed proxy request without changing legacy history."""
|
||||
|
||||
with self._lock:
|
||||
self._persistent_metrics.record_failed(provider=provider, model=model)
|
||||
self._maybe_save_locked()
|
||||
|
||||
def record_lifetime_rate_limited(
|
||||
self, *, provider: str | None = None, model: str | None = None
|
||||
) -> None:
|
||||
"""Record a rate-limited proxy request without changing legacy history."""
|
||||
|
||||
with self._lock:
|
||||
self._persistent_metrics.record_rate_limited(provider=provider, model=model)
|
||||
self._maybe_save_locked()
|
||||
|
||||
def record_lifetime_cache_bust(self, *, tokens_lost: int) -> None:
|
||||
with self._lock:
|
||||
self._persistent_metrics.record_cache_bust(tokens_lost=tokens_lost)
|
||||
self._maybe_save_locked()
|
||||
|
||||
def record_lifetime_cache_miss(self, *, provider: str | None, reason: str | None) -> None:
|
||||
with self._lock:
|
||||
self._persistent_metrics.record_cache_miss(provider=provider, reason=reason)
|
||||
self._maybe_save_locked()
|
||||
|
||||
def _record_project_locked(
|
||||
self,
|
||||
project: str | None,
|
||||
|
|
@ -891,6 +958,24 @@ class SavingsTracker:
|
|||
)
|
||||
result[model] = view
|
||||
return result
|
||||
def lifetime_response(self) -> dict[str, Any]:
|
||||
"""Return the durable aggregate used only by ``/stats-lifetime``."""
|
||||
|
||||
with self._lock:
|
||||
response = self._persistent_metrics.snapshot(
|
||||
persistence={
|
||||
"enabled": not self._stateless,
|
||||
"healthy": self._persistence_healthy,
|
||||
"error": (
|
||||
"Lifetime metrics unavailable in stateless mode"
|
||||
if self._stateless
|
||||
else self._persistence_error
|
||||
),
|
||||
"pending_records": 0 if self._stateless else self._since_save,
|
||||
}
|
||||
)
|
||||
response["projects"] = self._projects_snapshot_locked()
|
||||
return response
|
||||
|
||||
def stats_preview(self, recent_points: int = 20) -> dict[str, Any]:
|
||||
"""Return a compact preview for `/stats`."""
|
||||
|
|
@ -1031,10 +1116,30 @@ class SavingsTracker:
|
|||
raw = json.load(f)
|
||||
except (json.JSONDecodeError, OSError) as e:
|
||||
logger.warning("Failed to load savings history from %s: %s", self._path, e)
|
||||
self._persistence_healthy = False
|
||||
self._persistence_error = str(e)
|
||||
self._preserve_corrupt_file()
|
||||
return self._default_state()
|
||||
|
||||
if not isinstance(raw, dict):
|
||||
self._persistence_healthy = False
|
||||
self._persistence_error = "Savings state root must be a JSON object"
|
||||
self._preserve_corrupt_file()
|
||||
return self._default_state()
|
||||
if _coerce_int(raw.get("schema_version")) != SCHEMA_VERSION:
|
||||
self._needs_schema_save = True
|
||||
return self._sanitize_state(raw)
|
||||
|
||||
def _preserve_corrupt_file(self) -> None:
|
||||
"""Best-effort retention of unreadable state before starting fresh."""
|
||||
|
||||
timestamp = _to_utc_iso(_utc_now()).replace(":", "-")
|
||||
corrupt_path = self._path.with_name(f"{self._path.name}.corrupt-{timestamp}")
|
||||
try:
|
||||
self._path.replace(corrupt_path)
|
||||
except OSError:
|
||||
logger.warning("Could not preserve corrupt savings history at %s", self._path)
|
||||
|
||||
def _sanitize_state(self, raw: Any) -> dict[str, Any]:
|
||||
if not isinstance(raw, dict):
|
||||
return self._default_state()
|
||||
|
|
@ -1101,6 +1206,12 @@ class SavingsTracker:
|
|||
"projects": _normalize_projects(raw.get("projects")),
|
||||
"by_model": _normalize_by_model(raw.get("by_model")),
|
||||
}
|
||||
raw_lifetime_metrics = raw.get("lifetime_metrics")
|
||||
if isinstance(raw_lifetime_metrics, dict):
|
||||
state["lifetime_metrics"] = raw_lifetime_metrics
|
||||
else:
|
||||
state["lifetime_metrics"] = self._migrate_v4_lifetime_metrics(state)
|
||||
self._needs_schema_save = True
|
||||
|
||||
if normalized_history:
|
||||
reference_time = _parse_timestamp(normalized_history[-1]["timestamp"]) or _utc_now()
|
||||
|
|
@ -1115,6 +1226,49 @@ class SavingsTracker:
|
|||
|
||||
return state
|
||||
|
||||
def _migrate_v4_lifetime_metrics(self, state: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Seed the v5 aggregate from the best information available in v4."""
|
||||
|
||||
history = state["history"]
|
||||
display_session = state["display_session"]
|
||||
started_at = (
|
||||
history[0]["timestamp"]
|
||||
if history
|
||||
else display_session.get("started_at") or _to_utc_iso(_utc_now())
|
||||
)
|
||||
last_activity_at = (
|
||||
history[-1]["timestamp"] if history else display_session.get("last_activity_at")
|
||||
)
|
||||
legacy = state["lifetime"]
|
||||
return {
|
||||
"started_at": started_at,
|
||||
"last_activity_at": last_activity_at,
|
||||
"full_fidelity_started_at": _to_utc_iso(_utc_now()),
|
||||
"requests": {"total": legacy["requests"]},
|
||||
"tokens": {
|
||||
"input": legacy["total_input_tokens"],
|
||||
"attempted_input": legacy["total_input_tokens"] + legacy["tokens_saved"],
|
||||
"saved": legacy["tokens_saved"],
|
||||
},
|
||||
"prefix_cache": {"cache_read_tokens": legacy["cache_read_tokens"]},
|
||||
"cost": {
|
||||
"input_usd": legacy["total_input_cost_usd"],
|
||||
"compression_savings_usd": legacy["compression_savings_usd"],
|
||||
"cache_savings_usd": legacy["cache_savings_usd"],
|
||||
},
|
||||
"models": {
|
||||
"tracked": {},
|
||||
"other": {
|
||||
"requests": legacy["requests"],
|
||||
"input_tokens": legacy["total_input_tokens"],
|
||||
"attempted_input_tokens": legacy["total_input_tokens"]
|
||||
+ legacy["tokens_saved"],
|
||||
"tokens_saved": legacy["tokens_saved"],
|
||||
"last_activity_at": last_activity_at,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
def _trim_history_locked(self, reference_time: datetime | None = None) -> None:
|
||||
history = self._state["history"]
|
||||
if not history:
|
||||
|
|
@ -1190,7 +1344,7 @@ class SavingsTracker:
|
|||
recent requests. No-op when nothing is buffered.
|
||||
"""
|
||||
with self._lock:
|
||||
if self._since_save > 0:
|
||||
if self._since_save > 0 or self._needs_schema_save:
|
||||
self._save_locked()
|
||||
|
||||
def _maybe_save_locked(self) -> None:
|
||||
|
|
@ -1206,9 +1360,13 @@ class SavingsTracker:
|
|||
if self._stateless:
|
||||
# Stateless mode: live counters stay in memory; nothing is persisted.
|
||||
self._since_save = 0
|
||||
self._needs_schema_save = False
|
||||
return
|
||||
try:
|
||||
self._path.parent.mkdir(parents=True, exist_ok=True)
|
||||
saved_at = _to_utc_iso(_utc_now())
|
||||
lifetime_metrics = self._persistent_metrics.to_dict()
|
||||
lifetime_metrics["persistence"]["last_saved_at"] = saved_at
|
||||
payload = {
|
||||
"schema_version": SCHEMA_VERSION,
|
||||
"lifetime": self._state["lifetime"],
|
||||
|
|
@ -1216,6 +1374,7 @@ class SavingsTracker:
|
|||
"history": self._state["history"],
|
||||
"projects": self._state.get("projects", {}),
|
||||
"by_model": self._state.get("by_model", {}),
|
||||
"lifetime_metrics": lifetime_metrics,
|
||||
}
|
||||
json_data = json.dumps(payload, indent=2)
|
||||
|
||||
|
|
@ -1255,9 +1414,15 @@ class SavingsTracker:
|
|||
|
||||
# Reset only after a durable write. A failed save leaves the counter
|
||||
# untouched so the next record retries instead of waiting a full window.
|
||||
self._persistent_metrics.set_last_saved_at(saved_at)
|
||||
self._persistence_healthy = True
|
||||
self._persistence_error = None
|
||||
self._needs_schema_save = False
|
||||
self._since_save = 0
|
||||
except OSError as e:
|
||||
logger.warning("Failed to save savings history to %s: %s", self._path, e)
|
||||
self._persistence_healthy = False
|
||||
self._persistence_error = str(e)
|
||||
|
||||
def _display_session_snapshot_locked(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -28,6 +28,7 @@ import asyncio
|
|||
import concurrent.futures
|
||||
import contextlib
|
||||
import hmac
|
||||
import ipaddress
|
||||
import json
|
||||
import logging
|
||||
import math
|
||||
|
|
@ -40,6 +41,7 @@ from dataclasses import fields, is_dataclass, replace
|
|||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any, Literal, TypedDict, cast
|
||||
from urllib.parse import urlsplit
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ..backends.base import Backend
|
||||
|
|
@ -2036,6 +2038,77 @@ def _request_is_loopback(request: Request) -> bool:
|
|||
return peer_is_trusted_gateway(client_host, load_trusted_gateway_cidrs())
|
||||
|
||||
|
||||
def _request_can_view_dashboard_metadata(
|
||||
request: Request,
|
||||
trusted_dashboard_client_cidrs: tuple[
|
||||
ipaddress.IPv4Network | ipaddress.IPv6Network, ...
|
||||
],
|
||||
) -> bool:
|
||||
"""Authorize sensitive ``/stats`` metadata without widening admin access."""
|
||||
if _request_is_loopback(request):
|
||||
return True
|
||||
|
||||
from headroom.proxy.forwarded_headers import peer_is_trusted_gateway, resolve_client_ip
|
||||
from headroom.proxy.loopback_guard import is_ip_literal_host_header
|
||||
|
||||
try:
|
||||
host_header = request.headers.get("host")
|
||||
except AttributeError:
|
||||
return False
|
||||
if not is_ip_literal_host_header(host_header):
|
||||
return False
|
||||
|
||||
# CIDR authorization makes this endpoint usable by a remote dashboard, but
|
||||
# it must not let an unrelated site read sensitive metadata through a
|
||||
# victim's browser. Native CLI clients usually send neither header, so
|
||||
# absence remains valid. If either browser provenance header is present,
|
||||
# require it to identify this exact scheme/host/port.
|
||||
if not _request_has_same_origin_or_no_provenance(request, host_header):
|
||||
return False
|
||||
|
||||
return peer_is_trusted_gateway(
|
||||
resolve_client_ip(request),
|
||||
trusted_dashboard_client_cidrs,
|
||||
)
|
||||
|
||||
|
||||
def _request_has_same_origin_or_no_provenance(
|
||||
request: Request, host_header: str
|
||||
) -> bool:
|
||||
"""Accept no browser provenance, otherwise require same-origin headers."""
|
||||
|
||||
from headroom.proxy.forwarded_headers import trusted_forwarded_headers
|
||||
|
||||
forwarded_proto = trusted_forwarded_headers(request)["proto"]
|
||||
request_scheme = forwarded_proto or request.url.scheme
|
||||
expected_origin = _normalized_http_origin(f"{request_scheme}://{host_header}")
|
||||
if expected_origin is None:
|
||||
return False
|
||||
|
||||
for header_name in ("origin", "referer"):
|
||||
header_value = request.headers.get(header_name)
|
||||
if header_value and _normalized_http_origin(header_value) != expected_origin:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def _normalized_http_origin(value: str) -> tuple[str, str, int] | None:
|
||||
"""Return a normalized HTTP(S) origin tuple."""
|
||||
|
||||
try:
|
||||
parsed = urlsplit(value.strip())
|
||||
port = parsed.port
|
||||
except ValueError:
|
||||
return None
|
||||
|
||||
scheme = parsed.scheme.lower()
|
||||
if scheme not in {"http", "https"} or not parsed.hostname:
|
||||
return None
|
||||
if port is None:
|
||||
port = 80 if scheme == "http" else 443
|
||||
return scheme, parsed.hostname.lower(), port
|
||||
|
||||
|
||||
_is_known_websocket_callback_failure = is_known_websocket_callback_failure
|
||||
|
||||
|
||||
|
|
@ -2047,6 +2120,11 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
|
|||
if not FASTAPI_AVAILABLE:
|
||||
raise ImportError("FastAPI required. Install: pip install fastapi uvicorn httpx")
|
||||
|
||||
from headroom.proxy.forwarded_headers import load_trusted_dashboard_client_cidrs
|
||||
|
||||
# Parse once at startup so invalid operator configuration fails loudly.
|
||||
trusted_dashboard_client_cidrs = load_trusted_dashboard_client_cidrs()
|
||||
|
||||
from contextlib import asynccontextmanager
|
||||
|
||||
# Always-on file logging to ~/.headroom/logs/ for `headroom perf` analysis.
|
||||
|
|
@ -2993,7 +3071,7 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
|
|||
def _build_recent_request_payload(limit: int = RECENT_REQUEST_LOG_WINDOW) -> dict[str, Any]:
|
||||
recent_request_logs = proxy.logger.get_recent(limit) if proxy.logger else []
|
||||
dashboard_recent_requests = []
|
||||
for log in recent_request_logs:
|
||||
for log in reversed(recent_request_logs):
|
||||
token_accounting_status = _recent_request_token_accounting_status(log)
|
||||
dashboard_recent_requests.append(
|
||||
{
|
||||
|
|
@ -3021,7 +3099,7 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
|
|||
"tool_schema_saved_tokens": _tool_schema_saved_from_tags(log.get("tags")),
|
||||
}
|
||||
)
|
||||
dashboard_recent_requests = dashboard_recent_requests[-10:]
|
||||
dashboard_recent_requests = dashboard_recent_requests[:25]
|
||||
return {
|
||||
"request_logs": recent_request_logs[-10:],
|
||||
"recent_requests": dashboard_recent_requests,
|
||||
|
|
@ -3659,7 +3737,10 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
|
|||
only for loopback callers — the local dashboard. Network callers still
|
||||
get the aggregate counters but never the per-request metadata.
|
||||
"""
|
||||
include_sensitive = _request_is_loopback(request)
|
||||
include_sensitive = _request_can_view_dashboard_metadata(
|
||||
request,
|
||||
trusted_dashboard_client_cidrs,
|
||||
)
|
||||
if cached:
|
||||
payload = dict(await _get_cached_stats_payload())
|
||||
if include_sensitive:
|
||||
|
|
@ -3676,6 +3757,21 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
|
|||
payload.pop("request_logs", None)
|
||||
return payload
|
||||
|
||||
@app.get("/stats-lifetime")
|
||||
async def stats_lifetime(request: Request):
|
||||
"""Return persisted lifetime aggregates with sensitive fields gated."""
|
||||
payload = dict(proxy.metrics.savings_tracker.lifetime_response())
|
||||
include_sensitive = _request_can_view_dashboard_metadata(
|
||||
request,
|
||||
trusted_dashboard_client_cidrs,
|
||||
)
|
||||
if not include_sensitive:
|
||||
payload.pop("projects", None)
|
||||
persistence = payload.get("persistence")
|
||||
if isinstance(persistence, dict):
|
||||
payload["persistence"] = {**persistence, "error": None}
|
||||
return payload
|
||||
|
||||
@app.post("/stats/reset", dependencies=[Depends(_require_loopback)])
|
||||
async def stats_reset():
|
||||
"""Reset in-memory proxy stats for local test/debug isolation."""
|
||||
|
|
|
|||
|
|
@ -19,7 +19,9 @@ from fastapi.testclient import TestClient
|
|||
from starlette.datastructures import Headers, State
|
||||
|
||||
from headroom.proxy.forwarded_headers import (
|
||||
TRUSTED_DASHBOARD_CLIENT_CIDRS_ENV,
|
||||
TRUSTED_GATEWAY_CIDRS_ENV,
|
||||
load_trusted_dashboard_client_cidrs,
|
||||
load_trusted_gateway_cidrs,
|
||||
peer_is_trusted_gateway,
|
||||
resolve_client_ip,
|
||||
|
|
@ -115,6 +117,36 @@ def test_load_cidrs_reads_env_by_default(monkeypatch: pytest.MonkeyPatch) -> Non
|
|||
assert [str(c) for c in cidrs] == ["10.0.0.0/8"]
|
||||
|
||||
|
||||
def test_load_dashboard_client_cidrs_uses_their_own_env(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setenv(
|
||||
TRUSTED_DASHBOARD_CLIENT_CIDRS_ENV,
|
||||
"100.90.0.5/32, fd7a:115c:a1e0::/48",
|
||||
)
|
||||
|
||||
cidrs = load_trusted_dashboard_client_cidrs()
|
||||
|
||||
assert [str(cidr) for cidr in cidrs] == [
|
||||
"100.90.0.5/32",
|
||||
"fd7a:115c:a1e0::/48",
|
||||
]
|
||||
|
||||
|
||||
def test_load_dashboard_client_cidrs_unset_is_empty(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.delenv(TRUSTED_DASHBOARD_CLIENT_CIDRS_ENV, raising=False)
|
||||
assert load_trusted_dashboard_client_cidrs() == ()
|
||||
|
||||
|
||||
def test_load_dashboard_client_cidrs_empty_and_malformed_values() -> None:
|
||||
assert load_trusted_dashboard_client_cidrs("") == ()
|
||||
assert load_trusted_dashboard_client_cidrs(" ") == ()
|
||||
with pytest.raises(ValueError, match=TRUSTED_DASHBOARD_CLIENT_CIDRS_ENV):
|
||||
load_trusted_dashboard_client_cidrs("100.90.0.5/32,not-a-cidr")
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────
|
||||
# Membership check (peer_is_trusted_gateway)
|
||||
# ──────────────────────────────────────────────────────────────────
|
||||
|
|
|
|||
155
tests/test_persistent_metrics.py
Normal file
155
tests/test_persistent_metrics.py
Normal file
|
|
@ -0,0 +1,155 @@
|
|||
"""Tests for durable, aggregate-only proxy Lifetime metrics."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime, timezone
|
||||
|
||||
import pytest
|
||||
|
||||
from headroom.proxy.persistent_metrics import PersistentMetricsState
|
||||
|
||||
FIXED_NOW = datetime(2026, 7, 14, 8, 30, tzinfo=timezone.utc)
|
||||
|
||||
|
||||
def _new_state() -> PersistentMetricsState:
|
||||
return PersistentMetricsState(now=lambda: FIXED_NOW)
|
||||
|
||||
|
||||
def test_snapshot_accumulates_request_token_cache_cost_and_waste_metrics() -> None:
|
||||
state = _new_state()
|
||||
|
||||
state.record_request(
|
||||
provider="anthropic",
|
||||
stack="codex",
|
||||
model="claude-test",
|
||||
input_tokens=100,
|
||||
output_tokens=20,
|
||||
attempted_input_tokens=150,
|
||||
tokens_saved=50,
|
||||
cached=True,
|
||||
cache_read_tokens=80,
|
||||
cache_write_tokens=40,
|
||||
cache_write_5m_tokens=10,
|
||||
cache_write_1h_tokens=30,
|
||||
uncached_input_tokens=20,
|
||||
input_usd=0.4,
|
||||
compression_savings_usd=0.2,
|
||||
cache_savings_usd=0.1,
|
||||
waste_signals={"repetition": 7},
|
||||
)
|
||||
state.record_failed(provider="anthropic", model="claude-test")
|
||||
state.record_rate_limited(provider="anthropic", model="claude-test")
|
||||
state.record_cache_bust(tokens_lost=9)
|
||||
state.record_cache_miss(provider="anthropic", reason="prefix_change")
|
||||
|
||||
snapshot = state.snapshot(persistence={"enabled": True, "healthy": True})
|
||||
|
||||
assert snapshot["scope"] == "lifetime"
|
||||
assert snapshot["requests"] == {
|
||||
"total": 1,
|
||||
"cached": 1,
|
||||
"failed": 1,
|
||||
"rate_limited": 1,
|
||||
"by_provider": {"anthropic": 1},
|
||||
"by_stack": {"codex": 1},
|
||||
}
|
||||
assert snapshot["tokens"] == {
|
||||
"input": 100,
|
||||
"output": 20,
|
||||
"attempted_input": 150,
|
||||
"saved": 50,
|
||||
"token_savings_percent": pytest.approx(50 / 150 * 100),
|
||||
}
|
||||
assert snapshot["prefix_cache"]["requests"] == 1
|
||||
assert snapshot["prefix_cache"]["hit_requests"] == 1
|
||||
assert snapshot["prefix_cache"]["cache_read_tokens"] == 80
|
||||
assert snapshot["prefix_cache"]["cache_write_tokens"] == 40
|
||||
assert snapshot["prefix_cache"]["cache_hit_rate"] == 100.0
|
||||
assert snapshot["prefix_cache"]["ttl_1h_percent"] == 75.0
|
||||
assert snapshot["prefix_cache"]["ttl_5m_percent"] == 25.0
|
||||
assert snapshot["prefix_cache"]["bust_count"] == 1
|
||||
assert snapshot["prefix_cache"]["bust_tokens"] == 9
|
||||
assert snapshot["prefix_cache"]["misses_by_reason"] == {"prefix_change": 1}
|
||||
assert snapshot["cost"] == {
|
||||
"input_usd": 0.4,
|
||||
"compression_savings_usd": 0.2,
|
||||
"cache_savings_usd": 0.1,
|
||||
}
|
||||
assert snapshot["waste_signals"] == {"repetition": 7}
|
||||
assert snapshot["by_model"]["claude-test"]["input_tokens"] == 100
|
||||
|
||||
|
||||
def test_snapshot_uses_null_for_ratios_without_a_denominator() -> None:
|
||||
snapshot = _new_state().snapshot(persistence={"enabled": True, "healthy": True})
|
||||
|
||||
assert snapshot["tokens"]["token_savings_percent"] is None
|
||||
assert snapshot["prefix_cache"]["cache_hit_rate"] is None
|
||||
assert snapshot["prefix_cache"]["ttl_1h_percent"] is None
|
||||
assert snapshot["prefix_cache"]["ttl_5m_percent"] is None
|
||||
|
||||
|
||||
def test_candidate_models_remain_available_until_the_two_hundred_and_first_model() -> None:
|
||||
state = _new_state()
|
||||
for index in range(200):
|
||||
state.record_request(
|
||||
provider="provider",
|
||||
stack="stack",
|
||||
model=f"model-{index:03}",
|
||||
input_tokens=index + 1,
|
||||
)
|
||||
|
||||
persisted = state.to_dict()
|
||||
snapshot = state.snapshot(persistence={"enabled": True, "healthy": True})
|
||||
|
||||
assert len(persisted["models"]["tracked"]) == 200
|
||||
assert "model-000" not in snapshot["by_model"]
|
||||
assert snapshot["by_model"]["other"]["input_tokens"] == sum(range(1, 101))
|
||||
|
||||
|
||||
def test_two_hundred_and_first_model_permanently_compacts_non_top_candidates() -> None:
|
||||
state = _new_state()
|
||||
for index in range(201):
|
||||
state.record_request(
|
||||
provider="provider",
|
||||
stack="stack",
|
||||
model=f"model-{index:03}",
|
||||
input_tokens=index + 1,
|
||||
)
|
||||
|
||||
persisted = state.to_dict()
|
||||
snapshot = state.snapshot(persistence={"enabled": True, "healthy": True})
|
||||
|
||||
assert len(persisted["models"]["tracked"]) == 100
|
||||
assert set(snapshot["by_model"]) == {
|
||||
*(f"model-{index:03}" for index in range(101, 201)),
|
||||
"other",
|
||||
}
|
||||
assert snapshot["by_model"]["other"]["input_tokens"] == sum(range(1, 102))
|
||||
|
||||
|
||||
def test_state_normalizes_invalid_values_and_unknown_dimension_labels() -> None:
|
||||
state = PersistentMetricsState(
|
||||
{
|
||||
"requests": {"total": "not-a-number"},
|
||||
"tokens": {"input": float("nan"), "output": -3},
|
||||
"models": {"tracked": {"unknown": {"input_tokens": "7"}}},
|
||||
},
|
||||
now=lambda: FIXED_NOW,
|
||||
)
|
||||
state.record_request(
|
||||
provider=" ",
|
||||
stack=None,
|
||||
model=" ",
|
||||
input_tokens=-1,
|
||||
output_tokens=float("inf"),
|
||||
waste_signals={"unrecognized": 9},
|
||||
)
|
||||
|
||||
snapshot = state.snapshot(persistence={"enabled": True, "healthy": True})
|
||||
|
||||
assert snapshot["tokens"]["input"] == 0
|
||||
assert snapshot["tokens"]["output"] == 0
|
||||
assert snapshot["requests"]["by_provider"] == {"other": 1}
|
||||
assert snapshot["requests"]["by_stack"] == {"other": 1}
|
||||
assert snapshot["by_model"]["other"]["input_tokens"] == 7
|
||||
assert snapshot["waste_signals"] == {"other": 9}
|
||||
49
tests/test_persistent_metrics_integration.py
Normal file
49
tests/test_persistent_metrics_integration.py
Normal file
|
|
@ -0,0 +1,49 @@
|
|||
"""Tests forwarding existing proxy metric events to Lifetime storage."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
|
||||
from headroom.proxy.prometheus_metrics import PrometheusMetrics
|
||||
from headroom.proxy.savings_tracker import SavingsTracker
|
||||
|
||||
|
||||
def test_runtime_metric_events_feed_lifetime_without_resetting_runtime_counters(tmp_path) -> None:
|
||||
tracker = SavingsTracker(path=str(tmp_path / "proxy_savings.json"), save_flush_every=25)
|
||||
metrics = PrometheusMetrics(savings_tracker=tracker)
|
||||
|
||||
metrics.record_stack("codex")
|
||||
asyncio.run(
|
||||
metrics.record_request(
|
||||
provider="anthropic",
|
||||
model="claude-test",
|
||||
input_tokens=10,
|
||||
output_tokens=3,
|
||||
tokens_saved=2,
|
||||
latency_ms=1,
|
||||
cached=True,
|
||||
attempted_input_tokens=12,
|
||||
cache_read_tokens=5,
|
||||
cache_write_1h_tokens=2,
|
||||
waste_signals={"repetition": 4},
|
||||
)
|
||||
)
|
||||
asyncio.run(metrics.record_failed(provider="anthropic", model="claude-test"))
|
||||
asyncio.run(metrics.record_rate_limited(provider="anthropic", model="claude-test"))
|
||||
asyncio.run(metrics.record_cache_bust(tokens_lost=7))
|
||||
asyncio.run(metrics.record_cache_miss_attribution("anthropic", "prefix_change"))
|
||||
|
||||
lifetime = tracker.lifetime_response()
|
||||
|
||||
assert lifetime["requests"]["total"] == 1
|
||||
assert lifetime["requests"]["cached"] == 1
|
||||
assert lifetime["requests"]["failed"] == 1
|
||||
assert lifetime["requests"]["rate_limited"] == 1
|
||||
assert lifetime["requests"]["by_provider"] == {"anthropic": 1}
|
||||
assert lifetime["requests"]["by_stack"] == {"codex": 1}
|
||||
assert lifetime["tokens"]["output"] == 3
|
||||
assert lifetime["tokens"]["attempted_input"] == 12
|
||||
assert lifetime["prefix_cache"]["bust_tokens"] == 7
|
||||
assert lifetime["prefix_cache"]["misses_by_reason"] == {"prefix_change": 1}
|
||||
assert lifetime["waste_signals"] == {"repetition": 4}
|
||||
assert metrics.requests_total == 1
|
||||
87
tests/test_persistent_metrics_persistence.py
Normal file
87
tests/test_persistent_metrics_persistence.py
Normal file
|
|
@ -0,0 +1,87 @@
|
|||
"""Tests schema v5 persistence around the pure Lifetime aggregate."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
|
||||
from headroom.proxy.savings_tracker import SavingsTracker
|
||||
|
||||
|
||||
def test_savings_tracker_migrates_v4_lifetime_to_v5_metrics_and_preserves_legacy_state(tmp_path):
|
||||
path = tmp_path / "proxy_savings.json"
|
||||
legacy_state = {
|
||||
"schema_version": 4,
|
||||
"lifetime": {
|
||||
"requests": 7,
|
||||
"tokens_saved": 20,
|
||||
"compression_savings_usd": 0.5,
|
||||
"cache_read_tokens": 5,
|
||||
"cache_savings_usd": 0.2,
|
||||
"total_input_tokens": 80,
|
||||
"total_input_cost_usd": 1.5,
|
||||
},
|
||||
"display_session": {
|
||||
"requests": 2,
|
||||
"tokens_saved": 4,
|
||||
"compression_savings_usd": 0.1,
|
||||
"cache_read_tokens": 1,
|
||||
"cache_savings_usd": 0.01,
|
||||
"total_input_tokens": 10,
|
||||
"total_input_cost_usd": 0.2,
|
||||
"started_at": "2026-07-01T00:00:00Z",
|
||||
"last_activity_at": "2026-07-02T00:00:00Z",
|
||||
},
|
||||
"history": [
|
||||
{
|
||||
"timestamp": "2026-07-02T00:00:00Z",
|
||||
"total_tokens_saved": 20,
|
||||
"compression_savings_usd": 0.5,
|
||||
"total_input_tokens": 80,
|
||||
"total_input_cost_usd": 1.5,
|
||||
}
|
||||
],
|
||||
"projects": {"keep-me": {"requests": 1}},
|
||||
}
|
||||
path.write_text(json.dumps(legacy_state), encoding="utf-8")
|
||||
|
||||
tracker = SavingsTracker(path=str(path), save_flush_every=25)
|
||||
lifetime = tracker.lifetime_response()
|
||||
|
||||
assert lifetime["schema_version"] == 5
|
||||
assert lifetime["requests"]["total"] == 7
|
||||
assert lifetime["tokens"]["input"] == 80
|
||||
assert lifetime["tokens"]["attempted_input"] == 100
|
||||
assert lifetime["tokens"]["saved"] == 20
|
||||
assert lifetime["prefix_cache"]["cache_read_tokens"] == 5
|
||||
assert lifetime["cost"] == {
|
||||
"input_usd": 1.5,
|
||||
"compression_savings_usd": 0.5,
|
||||
"cache_savings_usd": 0.2,
|
||||
}
|
||||
assert lifetime["by_model"]["other"]["input_tokens"] == 80
|
||||
|
||||
tracker.flush()
|
||||
saved = json.loads(path.read_text(encoding="utf-8"))
|
||||
assert saved["schema_version"] == 5
|
||||
assert saved["lifetime"] == legacy_state["lifetime"]
|
||||
assert saved["display_session"]["requests"] == 2
|
||||
assert saved["projects"]["keep-me"]["requests"] == 1
|
||||
assert saved["lifetime_metrics"]["models"]["other"]["input_tokens"] == 80
|
||||
assert isinstance(saved["lifetime_metrics"]["persistence"]["last_saved_at"], str)
|
||||
|
||||
|
||||
def test_lifetime_response_reports_stateless_mode_without_writing(tmp_path):
|
||||
path = tmp_path / "proxy_savings.json"
|
||||
tracker = SavingsTracker(path=str(path), stateless=True, save_flush_every=1)
|
||||
|
||||
tracker.record_lifetime_request(provider="openai", stack="codex", model="gpt-test", input_tokens=3)
|
||||
|
||||
response = tracker.lifetime_response()
|
||||
assert response["persistence"] == {
|
||||
"enabled": False,
|
||||
"healthy": True,
|
||||
"error": "Lifetime metrics unavailable in stateless mode",
|
||||
"pending_records": 0,
|
||||
"last_saved_at": None,
|
||||
}
|
||||
assert path.exists() is False
|
||||
|
|
@ -16,6 +16,7 @@ from fastapi.testclient import TestClient
|
|||
|
||||
from headroom.cache.backends import InMemoryBackend
|
||||
from headroom.cache.compression_store import get_compression_store, reset_compression_store
|
||||
from headroom.proxy.loopback_guard import is_ip_literal_host_header
|
||||
from headroom.proxy.server import ProxyConfig, create_app
|
||||
|
||||
GATED = [
|
||||
|
|
@ -75,6 +76,57 @@ def test_loopback_caller_allowed(method: str, path: str) -> None:
|
|||
|
||||
|
||||
# CCR data endpoints — cached session content, gated to 404 off-loopback (#1227).
|
||||
def test_stats_lifetime_route_uses_dashboard_metadata_access_policy(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setenv(
|
||||
"HEADROOM_PROXY_TRUSTED_DASHBOARD_CLIENT_CIDRS",
|
||||
"100.90.0.5/32",
|
||||
)
|
||||
app = _make_app()
|
||||
expected = {
|
||||
"requests": {"total": 7},
|
||||
"projects": {"headroom": {"requests": 3}},
|
||||
"persistence": {
|
||||
"enabled": True,
|
||||
"healthy": False,
|
||||
"error": "D:/private/proxy_savings.json: access denied",
|
||||
},
|
||||
}
|
||||
monkeypatch.setattr(
|
||||
app.state.proxy.metrics.savings_tracker,
|
||||
"lifetime_response",
|
||||
lambda: expected,
|
||||
)
|
||||
|
||||
network = TestClient(app).get("/stats-lifetime")
|
||||
assert network.status_code == 200, network.text
|
||||
assert network.json() == {
|
||||
"requests": {"total": 7},
|
||||
"persistence": {
|
||||
"enabled": True,
|
||||
"healthy": False,
|
||||
"error": None,
|
||||
},
|
||||
}
|
||||
|
||||
loopback = TestClient(
|
||||
app,
|
||||
base_url="http://127.0.0.1",
|
||||
client=("127.0.0.1", 12345),
|
||||
).get("/stats-lifetime")
|
||||
assert loopback.status_code == 200, loopback.text
|
||||
assert loopback.json() == expected
|
||||
|
||||
trusted_dashboard = TestClient(
|
||||
app,
|
||||
base_url="http://100.82.0.2:8787",
|
||||
client=("100.90.0.5", 12345),
|
||||
).get("/stats-lifetime")
|
||||
assert trusted_dashboard.status_code == 200, trusted_dashboard.text
|
||||
assert trusted_dashboard.json() == expected
|
||||
|
||||
|
||||
CCR_GATED = [
|
||||
("post", "/v1/retrieve"),
|
||||
("get", "/v1/retrieve/stats"),
|
||||
|
|
@ -112,6 +164,22 @@ def test_dns_rebinding_host_header_rejected() -> None:
|
|||
assert resp.status_code == 404, resp.text
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"host_header",
|
||||
["100.82.0.2", "100.82.0.2:8787", "[fd7a:115c:a1e0::2]", "[fd7a:115c:a1e0::2]:8787"],
|
||||
)
|
||||
def test_ip_literal_host_header_accepts_ip_addresses(host_header: str) -> None:
|
||||
assert is_ip_literal_host_header(host_header) is True
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"host_header",
|
||||
[None, "", "attacker.example", "localhost", "user@100.82.0.2", "100.82.0.2/path", "[fd7a::1"],
|
||||
)
|
||||
def test_ip_literal_host_header_rejects_non_addresses(host_header: str | None) -> None:
|
||||
assert is_ip_literal_host_header(host_header) is False
|
||||
|
||||
|
||||
def _client(*, loopback: bool) -> TestClient:
|
||||
app = _make_app()
|
||||
if loopback:
|
||||
|
|
@ -185,3 +253,198 @@ def test_stats_metadata_served_to_trusted_gateway_peer(
|
|||
rebind = TestClient(app, base_url="http://attacker.example", client=(gateway_ip, 54321))
|
||||
payload = rebind.get("/stats").json()
|
||||
assert "recent_requests" not in payload
|
||||
|
||||
|
||||
@pytest.mark.parametrize("cached", [False, True])
|
||||
def test_dashboard_client_cidr_grants_stats_metadata_for_ip_literal_host(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
cached: bool,
|
||||
) -> None:
|
||||
monkeypatch.setenv("HEADROOM_PROXY_TRUSTED_DASHBOARD_CLIENT_CIDRS", "100.90.0.5/32")
|
||||
app = _make_app()
|
||||
client = TestClient(
|
||||
app,
|
||||
base_url="http://100.82.0.2:8787",
|
||||
client=("100.90.0.5", 12345),
|
||||
)
|
||||
|
||||
payload = client.get("/stats", params={"cached": int(cached)}).json()
|
||||
|
||||
assert "recent_requests" in payload
|
||||
assert "request_logs" in payload
|
||||
assert "config" in payload
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"headers",
|
||||
[
|
||||
{"origin": "http://100.82.0.2:8787"},
|
||||
{"referer": "http://100.82.0.2:8787/dashboard"},
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("cached", [False, True])
|
||||
def test_dashboard_client_cidr_grants_stats_metadata_to_same_origin_browser(
|
||||
monkeypatch: pytest.MonkeyPatch, headers: dict[str, str], cached: bool
|
||||
) -> None:
|
||||
monkeypatch.setenv("HEADROOM_PROXY_TRUSTED_DASHBOARD_CLIENT_CIDRS", "100.90.0.5/32")
|
||||
client = TestClient(
|
||||
_make_app(),
|
||||
base_url="http://100.82.0.2:8787",
|
||||
client=("100.90.0.5", 12345),
|
||||
)
|
||||
|
||||
payload = client.get(
|
||||
"/stats", params={"cached": int(cached)}, headers=headers
|
||||
).json()
|
||||
|
||||
assert "recent_requests" in payload
|
||||
assert "request_logs" in payload
|
||||
assert "config" in payload
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"headers",
|
||||
[
|
||||
{"origin": "http://attacker.example"},
|
||||
{"referer": "http://attacker.example/dashboard"},
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("cached", [False, True])
|
||||
def test_dashboard_client_cidr_hides_stats_metadata_from_cross_origin_browser(
|
||||
monkeypatch: pytest.MonkeyPatch, headers: dict[str, str], cached: bool
|
||||
) -> None:
|
||||
monkeypatch.setenv("HEADROOM_PROXY_TRUSTED_DASHBOARD_CLIENT_CIDRS", "100.90.0.5/32")
|
||||
client = TestClient(
|
||||
_make_app(),
|
||||
base_url="http://100.82.0.2:8787",
|
||||
client=("100.90.0.5", 12345),
|
||||
)
|
||||
|
||||
response = client.get(
|
||||
"/stats", params={"cached": int(cached)}, headers=headers
|
||||
)
|
||||
payload = response.json()
|
||||
|
||||
assert response.status_code == 200
|
||||
assert "tokens" in payload
|
||||
assert "recent_requests" not in payload
|
||||
assert "request_logs" not in payload
|
||||
assert "config" not in payload
|
||||
|
||||
|
||||
def test_dashboard_client_cidr_only_uses_forwarded_proto_from_trusted_gateway(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setenv("HEADROOM_PROXY_TRUSTED_DASHBOARD_CLIENT_CIDRS", "100.90.0.5/32")
|
||||
monkeypatch.setenv("HEADROOM_PROXY_TRUSTED_GATEWAY_CIDRS", "172.18.0.0/16")
|
||||
client = TestClient(
|
||||
_make_app(),
|
||||
base_url="http://100.82.0.2:8787",
|
||||
client=("172.18.0.1", 12345),
|
||||
)
|
||||
|
||||
payload = client.get(
|
||||
"/stats",
|
||||
headers={
|
||||
"origin": "https://100.82.0.2:8787",
|
||||
"x-forwarded-for": "100.90.0.5",
|
||||
"x-forwarded-proto": "https",
|
||||
},
|
||||
).json()
|
||||
|
||||
assert "recent_requests" in payload
|
||||
assert "request_logs" in payload
|
||||
assert "config" in payload
|
||||
|
||||
spoofed = TestClient(
|
||||
_make_app(),
|
||||
base_url="http://100.82.0.2:8787",
|
||||
client=("100.90.0.5", 12345),
|
||||
).get(
|
||||
"/stats",
|
||||
headers={
|
||||
"origin": "https://100.82.0.2:8787",
|
||||
"x-forwarded-proto": "https",
|
||||
},
|
||||
).json()
|
||||
|
||||
assert "recent_requests" not in spoofed
|
||||
assert "request_logs" not in spoofed
|
||||
assert "config" not in spoofed
|
||||
|
||||
|
||||
def test_dashboard_client_cidr_rejects_unlisted_clients_and_hostname_hosts(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setenv("HEADROOM_PROXY_TRUSTED_DASHBOARD_CLIENT_CIDRS", "100.90.0.5/32")
|
||||
app = _make_app()
|
||||
|
||||
unlisted = TestClient(
|
||||
app,
|
||||
base_url="http://100.82.0.2:8787",
|
||||
client=("100.90.0.6", 12345),
|
||||
).get("/stats").json()
|
||||
hostname = TestClient(
|
||||
app,
|
||||
base_url="http://100.82.0.2:8787",
|
||||
client=("100.90.0.5", 12345),
|
||||
).get("/stats", headers={"host": "attacker.example"}).json()
|
||||
|
||||
for payload in (unlisted, hostname):
|
||||
assert "recent_requests" not in payload
|
||||
assert "request_logs" not in payload
|
||||
assert "config" not in payload
|
||||
|
||||
|
||||
def test_dashboard_client_cidr_only_accepts_forwarded_client_from_trusted_gateway(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setenv("HEADROOM_PROXY_TRUSTED_DASHBOARD_CLIENT_CIDRS", "100.90.0.5/32")
|
||||
monkeypatch.setenv("HEADROOM_PROXY_TRUSTED_GATEWAY_CIDRS", "172.18.0.0/16")
|
||||
app = _make_app()
|
||||
|
||||
trusted = TestClient(
|
||||
app,
|
||||
base_url="http://100.82.0.2:8787",
|
||||
client=("172.18.0.1", 12345),
|
||||
).get("/stats", headers={"x-forwarded-for": "100.90.0.5"}).json()
|
||||
forged = TestClient(
|
||||
app,
|
||||
base_url="http://100.82.0.2:8787",
|
||||
client=("198.51.100.10", 12345),
|
||||
).get("/stats", headers={"x-forwarded-for": "100.90.0.5"}).json()
|
||||
|
||||
assert "recent_requests" in trusted
|
||||
assert "recent_requests" not in forged
|
||||
|
||||
|
||||
def test_dashboard_client_cidr_normalizes_ipv4_mapped_ipv6(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setenv("HEADROOM_PROXY_TRUSTED_DASHBOARD_CLIENT_CIDRS", "100.90.0.0/24")
|
||||
app = _make_app()
|
||||
payload = TestClient(
|
||||
app,
|
||||
base_url="http://100.82.0.2:8787",
|
||||
client=("::ffff:100.90.0.5", 12345),
|
||||
).get("/stats").json()
|
||||
|
||||
assert "recent_requests" in payload
|
||||
|
||||
|
||||
def test_dashboard_client_cidr_does_not_expand_other_management_endpoints(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setenv("HEADROOM_PROXY_TRUSTED_DASHBOARD_CLIENT_CIDRS", "100.90.0.5/32")
|
||||
client = TestClient(
|
||||
_make_app(),
|
||||
base_url="http://100.82.0.2:8787",
|
||||
client=("100.90.0.5", 12345),
|
||||
)
|
||||
|
||||
health = client.get("/health")
|
||||
assert health.status_code == 200
|
||||
assert "config" not in health.json()
|
||||
assert client.get("/admin/upstream").status_code == 404
|
||||
assert client.get("/debug/tasks").status_code == 404
|
||||
assert client.post("/stats/reset").status_code == 404
|
||||
|
|
|
|||
|
|
@ -122,6 +122,7 @@ def test_tracker_accumulates_per_project_and_persists(tmp_path):
|
|||
# Survives a restart via the persisted JSON state.
|
||||
reloaded = SavingsTracker(path=str(path))
|
||||
assert reloaded.stats_preview()["projects"]["api"]["tokens_saved"] == 500
|
||||
assert reloaded.lifetime_response()["projects"]["api"]["tokens_saved"] == 500
|
||||
|
||||
|
||||
def test_tracker_migrates_v2_state_without_projects(tmp_path):
|
||||
|
|
|
|||
|
|
@ -81,7 +81,7 @@ def test_stats_refreshes_recent_requests_when_cached() -> None:
|
|||
assert second_response.status_code == 200
|
||||
second_payload = second_response.json()
|
||||
|
||||
assert second_payload["recent_requests"][-1]["model"] == "claude-sonnet"
|
||||
assert second_payload["recent_requests"][0]["model"] == "claude-sonnet"
|
||||
assert second_payload["request_logs"][-1]["model"] == "claude-sonnet"
|
||||
|
||||
|
||||
|
|
@ -148,18 +148,18 @@ def test_stats_recent_requests_includes_token_incomplete_requests() -> None:
|
|||
assert response.status_code == 200
|
||||
payload = response.json()
|
||||
assert [req["model"] for req in payload["recent_requests"]] == [
|
||||
"claude-haiku",
|
||||
"claude-haiku",
|
||||
"claude-sonnet",
|
||||
"claude-haiku",
|
||||
"claude-haiku",
|
||||
]
|
||||
assert payload["recent_requests"][0]["input_tokens_optimized"] is None
|
||||
assert payload["recent_requests"][0]["token_accounting_status"] == "missing"
|
||||
assert payload["recent_requests"][0]["has_exact_tokens"] is False
|
||||
assert payload["recent_requests"][0]["token_accounting_status"] == "complete"
|
||||
assert payload["recent_requests"][0]["has_exact_tokens"] is True
|
||||
assert payload["recent_requests"][1]["output_tokens"] is None
|
||||
assert payload["recent_requests"][1]["token_accounting_status"] == "partial"
|
||||
assert payload["recent_requests"][1]["tokens_saved"] == 0
|
||||
assert payload["recent_requests"][2]["token_accounting_status"] == "complete"
|
||||
assert payload["recent_requests"][2]["has_exact_tokens"] is True
|
||||
assert payload["recent_requests"][2]["input_tokens_optimized"] is None
|
||||
assert payload["recent_requests"][2]["token_accounting_status"] == "missing"
|
||||
assert payload["recent_requests"][2]["has_exact_tokens"] is False
|
||||
assert payload["summary"]["uncompressed_requests"]["unknown_token_accounting"] == 2
|
||||
assert payload["request_logs"][-1]["model"] == "claude-sonnet"
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue