diff --git a/docs/observability.md b/docs/observability.md index b6de50dc6..b43f5a970 100644 --- a/docs/observability.md +++ b/docs/observability.md @@ -245,6 +245,21 @@ Every label vocabulary is bounded by code, not customer input: `"other"` and a `tracing::warn!` is emitted so wire-format drift surfaces loudly in logs. - `status`: 5-variant enum. +- `tool` (Python-side `wrap_rtk_invocations_total`): bounded by the + set of tools the wrap CLI rewrites, captured by + `headroom.cli.wrap_rtk_metrics`. +- `model` (Python-side `requests_by_model` / + `_cache_requests_by_model`): unlike the Rust path above, the Python + proxy reads `model` from the request body, so it is client-supplied. + It is bounded at record time by `MAX_DISTINCT_MODELS` + (`headroom.telemetry.context`): once the cap is reached, further + distinct models bucket into the `"other"` sentinel and a one-time + warning is logged, mirroring the `tier` discipline above. The + in-memory dicts and the exported `headroom_requests_by_model` series + can never exceed the cap plus `"other"`. + +Every label vocabulary listed above is bounded by code, so no +client-supplied value can drive label cardinality unbounded. There is no code path where a malicious client can drive label cardinality unbounded. diff --git a/headroom/proxy/prometheus_metrics.py b/headroom/proxy/prometheus_metrics.py index 9cea10638..39bef4754 100644 --- a/headroom/proxy/prometheus_metrics.py +++ b/headroom/proxy/prometheus_metrics.py @@ -26,6 +26,10 @@ from headroom.proxy.savings_tracker import SavingsTracker logger = logging.getLogger("headroom.proxy") +# Sentinel label value that models past MAX_DISTINCT_MODELS collapse into, so +# client-supplied model cardinality stays bounded (see record_request). +_OTHER_MODEL = "other" + def _escape_label_value(value: str) -> str: # The /metrics body is emitted whole with .encode("utf-8") (server.py). A @@ -89,6 +93,9 @@ class PrometheusMetrics: self.requests_total = 0 self.requests_by_provider: dict[str, int] = defaultdict(int) self.requests_by_model: dict[str, int] = defaultdict(int) + # Set once when requests_by_model first reaches MAX_DISTINCT_MODELS, so the + # cardinality-cap warning fires exactly once instead of per request. + self._model_cardinality_warned = False # Populated via X-Headroom-Stack header (TS SDK adapters, etc.) self.requests_by_stack: dict[str, int] = defaultdict(int) self.requests_cached = 0 @@ -322,6 +329,7 @@ class PrometheusMetrics: self.requests_total = 0 self.requests_by_provider.clear() self.requests_by_model.clear() + self._model_cardinality_warned = False self.requests_by_stack.clear() self.requests_cached = 0 self.requests_rate_limited = 0 @@ -731,6 +739,10 @@ class PrometheusMetrics: reduction/yield/ledger math never straddles two rulers. Defaults to ``input_tokens`` when omitted, preserving pre-split behaviour. """ + # Local import mirrors record_stack: defers to call-time (the telemetry + # package is fully loaded by then), avoiding an import cycle at module load. + from headroom.telemetry.context import MAX_DISTINCT_MODELS + ledger_input_tokens = input_tokens if local_input_tokens is None else local_input_tokens # Post-guard invariant (all providers): Headroom never forwards a request # larger than the original — handlers revert any inflation before sending @@ -748,7 +760,25 @@ class PrometheusMetrics: async with self._lock: self.requests_total += 1 self.requests_by_provider[provider] += 1 - self.requests_by_model[model] += 1 + # Cap client-supplied model cardinality. `model` is client-controlled + # (body.get("model") in the openai/gemini/bedrock handlers), so an + # arbitrary-model client would otherwise grow requests_by_model and the + # exported series without bound. Bucket over-cap models into "other" + # (the sentinel docs/observability.md documents for `tier`), mirroring + # the requests_by_stack cap. Membership test, never a defaultdict index: + # indexing would materialize the key and defeat the cap. + if model in self.requests_by_model or len(self.requests_by_model) < MAX_DISTINCT_MODELS: + bounded_model = model + else: + bounded_model = _OTHER_MODEL + if not self._model_cardinality_warned: + self._model_cardinality_warned = True + logger.warning( + "metrics.record: model cardinality cap (%d) reached; " + 'bucketing further models into "other"', + MAX_DISTINCT_MODELS, + ) + self.requests_by_model[bounded_model] += 1 if cached: self.requests_cached += 1 @@ -780,8 +810,13 @@ class PrometheusMetrics: # is always a cold start (100% write, 0% read) — not a bust. # Only flag as bust when a previously-warm model suddenly has # high write ratio, indicating prefix invalidation. - model_req_num = self._cache_requests_by_model[model] - self._cache_requests_by_model[model] += 1 + # bounded_model can be "other" once the cardinality cap trips, which + # mixes distinct models in this bust heuristic. That is acceptable: + # it only happens past MAX_DISTINCT_MODELS distinct models on cached + # anthropic traffic, the worst case is a mis-attributed bust stat, + # and it keeps _cache_requests_by_model bounded. + model_req_num = self._cache_requests_by_model[bounded_model] + self._cache_requests_by_model[bounded_model] += 1 if provider == "anthropic" and model_req_num > 0: total_cached = cache_read_tokens + cache_write_tokens if total_cached > 0 and cache_write_tokens > total_cached * 0.5: diff --git a/headroom/telemetry/context.py b/headroom/telemetry/context.py index e845b1ed7..93d5d3956 100644 --- a/headroom/telemetry/context.py +++ b/headroom/telemetry/context.py @@ -46,6 +46,14 @@ _STACK_SLUG_RE = re.compile(r"^[a-z][a-z0-9_]{0,63}$") # header values. MAX_DISTINCT_STACKS = 32 +# Cardinality cap on the per-process requests_by_model / _cache_requests_by_model +# dicts. Protects the Prometheus scrape, the in-memory counters, and telemetry +# from unbounded label explosion when clients send arbitrary `model` values. +# 32x MAX_DISTINCT_STACKS: models are a larger legitimate vocabulary (provider +# and snapshot variants across a multi-tenant deployment) than stacks, while the +# cap stays a hard ceiling. Over-cap models bucket into the "other" sentinel. +MAX_DISTINCT_MODELS = 1024 + def normalize_stack(raw: str | None) -> str | None: """Validate and normalize a stack slug. diff --git a/tests/test_observability_metrics.py b/tests/test_observability_metrics.py index 5a7297476..ae539bddb 100644 --- a/tests/test_observability_metrics.py +++ b/tests/test_observability_metrics.py @@ -2,6 +2,7 @@ from __future__ import annotations +import logging from dataclasses import dataclass, field from typing import Any @@ -16,6 +17,7 @@ from headroom.observability import ( set_otel_metrics, ) from headroom.proxy.prometheus_metrics import PrometheusMetrics +from headroom.telemetry.context import MAX_DISTINCT_MODELS from headroom.transforms.pipeline import TransformPipeline @@ -236,3 +238,108 @@ async def test_prometheus_metrics_clamps_negative_token_savings() -> None: assert metrics.tokens_saved_total == 0 assert metrics.savings_history[-1][1] == 0 + + +@pytest.mark.asyncio +async def test_prometheus_metrics_caps_model_cardinality() -> None: + """A client sending unbounded distinct models cannot grow the per-model dicts + past MAX_DISTINCT_MODELS + the "other" sentinel, while accounting stays exact.""" + metrics = PrometheusMetrics(stateless=True) + + async def record(model: str) -> None: + await metrics.record_request( + provider="anthropic", + model=model, + input_tokens=10, + output_tokens=1, + tokens_saved=1, + latency_ms=1.0, + cache_read_tokens=1, # enter the prefix-cache block -> _cache_requests_by_model + ) + + # Fill exactly to the cap with distinct models: no bucketing yet. + for i in range(MAX_DISTINCT_MODELS): + await record(f"model_{i}") + assert len(metrics.requests_by_model) == MAX_DISTINCT_MODELS + assert len(metrics._cache_requests_by_model) == MAX_DISTINCT_MODELS + assert "other" not in metrics.requests_by_model + + # New distinct models past the cap bucket into "other", never their own key. + for i in range(5): + await record(f"overflow_{i}") + assert "overflow_0" not in metrics.requests_by_model + assert metrics.requests_by_model["other"] == 5 + assert metrics._cache_requests_by_model["other"] == 5 + assert len(metrics.requests_by_model) == MAX_DISTINCT_MODELS + 1 + assert len(metrics._cache_requests_by_model) == MAX_DISTINCT_MODELS + 1 + + # An already-tracked model keeps incrementing after the cap is reached. + await record("model_0") + assert metrics.requests_by_model["model_0"] == 2 + + # Accounting is preserved: every request is counted somewhere. + total_calls = MAX_DISTINCT_MODELS + 5 + 1 + assert metrics.requests_total == total_calls + assert sum(metrics.requests_by_model.values()) == total_calls + + +@pytest.mark.asyncio +async def test_prometheus_metrics_model_cardinality_warns_once( + caplog: pytest.LogCaptureFixture, +) -> None: + """Bucketing into "other" logs exactly one warning, not one per request.""" + metrics = PrometheusMetrics(stateless=True) + with caplog.at_level(logging.WARNING, logger="headroom.proxy"): + for i in range(MAX_DISTINCT_MODELS + 10): + await metrics.record_request( + provider="openai", + model=f"model_{i}", + input_tokens=10, + output_tokens=1, + tokens_saved=1, + latency_ms=1.0, + ) + cap_warnings = [r for r in caplog.records if "cardinality cap" in r.getMessage()] + assert len(cap_warnings) == 1 + + +@pytest.mark.asyncio +async def test_prometheus_metrics_reset_rearms_cardinality_warning() -> None: + """reset_runtime clears the model dicts and re-arms the one-shot cap warning.""" + metrics = PrometheusMetrics(stateless=True) + for i in range(MAX_DISTINCT_MODELS + 5): + await metrics.record_request( + provider="openai", + model=f"model_{i}", + input_tokens=1, + output_tokens=1, + tokens_saved=1, + latency_ms=1.0, + cache_read_tokens=1, + ) + assert metrics._model_cardinality_warned is True + + await metrics.reset_runtime() + + assert metrics._model_cardinality_warned is False + assert len(metrics.requests_by_model) == 0 + assert len(metrics._cache_requests_by_model) == 0 + + +@pytest.mark.asyncio +async def test_prometheus_metrics_export_bounds_model_series() -> None: + """export() emits at most MAX_DISTINCT_MODELS model series plus the 'other' bucket.""" + metrics = PrometheusMetrics(stateless=True) + for i in range(MAX_DISTINCT_MODELS + 20): + await metrics.record_request( + provider="openai", + model=f"model_{i}", + input_tokens=1, + output_tokens=1, + tokens_saved=1, + latency_ms=1.0, + ) + text = await metrics.export() + series = text.count("headroom_requests_by_model{") + assert series <= MAX_DISTINCT_MODELS + 1 + assert 'headroom_requests_by_model{model="other"}' in text