diff --git a/headroom/proxy/prometheus_metrics.py b/headroom/proxy/prometheus_metrics.py index 8cd83846b..9cea10638 100644 --- a/headroom/proxy/prometheus_metrics.py +++ b/headroom/proxy/prometheus_metrics.py @@ -28,6 +28,13 @@ logger = logging.getLogger("headroom.proxy") def _escape_label_value(value: str) -> str: + # The /metrics body is emitted whole with .encode("utf-8") (server.py). A + # client-supplied value can be a valid str that is not UTF-8-encodable — a + # lone surrogate decoded from a JSON model id — which raises in the response + # encoder and 500s every scrape, not just its own line. Drop un-encodable + # code points before escaping so one malformed request can't down the + # endpoint. Byte-identical for encodable values, including non-ASCII. + value = value.encode("utf-8", "replace").decode("utf-8") return value.replace("\\", "\\\\").replace("\n", "\\n").replace('"', '\\"') @@ -1267,10 +1274,11 @@ class PrometheusMetrics: ] ) for _provider, _reasons in self.cache_miss_attribution_by_provider.items(): + _safe_provider = _escape_label_value(str(_provider)) for _reason, _count in _reasons.items(): lines.append( - f'headroom_cache_miss_attribution_total{{provider="{_provider}",' - f'reason="{_reason}"}} {_count}' + f'headroom_cache_miss_attribution_total{{provider="{_safe_provider}",' + f'reason="{_escape_label_value(str(_reason))}"}} {_count}' ) lines.append("") @@ -1327,7 +1335,9 @@ class PrometheusMetrics: ] ) for provider, count in self.requests_by_provider.items(): - lines.append(f'headroom_requests_by_provider{{provider="{provider}"}} {count}') + lines.append( + f'headroom_requests_by_provider{{provider="{_escape_label_value(str(provider))}"}} {count}' + ) lines.append("") lines.extend( @@ -1337,7 +1347,9 @@ class PrometheusMetrics: ] ) for model, count in self.requests_by_model.items(): - lines.append(f'headroom_requests_by_model{{model="{model}"}} {count}') + lines.append( + f'headroom_requests_by_model{{model="{_escape_label_value(str(model))}"}} {count}' + ) lines.append("") if self.transform_timing_sum: @@ -1472,13 +1484,20 @@ class PrometheusMetrics: lines.append("") if self.cache_by_provider: + # The exposition format wants each family's samples grouped, so the + # blocks below re-walk this dict once per family. Escape the provider + # keys once here instead of at all eleven emission sites. + cache_by_provider = { + _escape_label_value(str(name)): stats + for name, stats in self.cache_by_provider.items() + } lines.extend( [ "# HELP headroom_cache_read_tokens_total Provider cache read tokens", "# TYPE headroom_cache_read_tokens_total counter", ] ) - for provider, stats in self.cache_by_provider.items(): + for provider, stats in cache_by_provider.items(): lines.append( f'headroom_cache_read_tokens_total{{provider="{provider}"}} {stats["cache_read_tokens"]}' ) @@ -1489,7 +1508,7 @@ class PrometheusMetrics: "# TYPE headroom_cache_write_tokens_total counter", ] ) - for provider, stats in self.cache_by_provider.items(): + for provider, stats in cache_by_provider.items(): lines.append( f'headroom_cache_write_tokens_total{{provider="{provider}"}} {stats["cache_write_tokens"]}' ) @@ -1500,7 +1519,7 @@ class PrometheusMetrics: "# TYPE headroom_cache_write_ttl_tokens_total counter", ] ) - for provider, stats in self.cache_by_provider.items(): + for provider, stats in cache_by_provider.items(): lines.append( f'headroom_cache_write_ttl_tokens_total{{provider="{provider}",ttl="5m"}} {stats["cache_write_5m_tokens"]}' ) @@ -1514,7 +1533,7 @@ class PrometheusMetrics: "# TYPE headroom_cache_write_ttl_requests_total counter", ] ) - for provider, stats in self.cache_by_provider.items(): + for provider, stats in cache_by_provider.items(): lines.append( f'headroom_cache_write_ttl_requests_total{{provider="{provider}",ttl="5m"}} {stats["cache_write_5m_requests"]}' ) @@ -1528,7 +1547,7 @@ class PrometheusMetrics: "# TYPE headroom_uncached_input_tokens_total counter", ] ) - for provider, stats in self.cache_by_provider.items(): + for provider, stats in cache_by_provider.items(): lines.append( f'headroom_uncached_input_tokens_total{{provider="{provider}"}} {stats["uncached_input_tokens"]}' ) @@ -1539,7 +1558,7 @@ class PrometheusMetrics: "# TYPE headroom_provider_cache_requests_total counter", ] ) - for provider, stats in self.cache_by_provider.items(): + for provider, stats in cache_by_provider.items(): lines.append( f'headroom_provider_cache_requests_total{{provider="{provider}"}} {stats["requests"]}' ) @@ -1550,7 +1569,7 @@ class PrometheusMetrics: "# TYPE headroom_provider_cache_hit_requests_total counter", ] ) - for provider, stats in self.cache_by_provider.items(): + for provider, stats in cache_by_provider.items(): lines.append( f'headroom_provider_cache_hit_requests_total{{provider="{provider}"}} {stats["hit_requests"]}' ) @@ -1561,7 +1580,7 @@ class PrometheusMetrics: "# TYPE headroom_provider_cache_bust_total counter", ] ) - for provider, stats in self.cache_by_provider.items(): + for provider, stats in cache_by_provider.items(): lines.append( f'headroom_provider_cache_bust_total{{provider="{provider}"}} {stats["bust_count"]}' ) @@ -1572,7 +1591,7 @@ class PrometheusMetrics: "# TYPE headroom_provider_cache_bust_write_tokens_total counter", ] ) - for provider, stats in self.cache_by_provider.items(): + for provider, stats in cache_by_provider.items(): lines.append( f'headroom_provider_cache_bust_write_tokens_total{{provider="{provider}"}} {stats["bust_write_tokens"]}' ) diff --git a/tests/test_prometheus_label_escaping.py b/tests/test_prometheus_label_escaping.py new file mode 100644 index 000000000..80dca7a9d --- /dev/null +++ b/tests/test_prometheus_label_escaping.py @@ -0,0 +1,273 @@ +"""Label-value escaping in the Prometheus text exposition output. + +``PrometheusMetrics.export()`` builds the exposition text by hand, so every label +value has to pass through ``_escape_label_value`` before it is interpolated. The +format reserves ``"``, ``\\`` and the line feed, and a standard scraper does not +degrade gracefully on a malformed line — it aborts the parse, losing every +family emitted at or after the bad sample. + +``model`` reaches ``requests_by_model`` straight from the parsed client request +body (``handlers/openai.py`` reads ``body.get("model", "unknown")`` with no +sanitisation, and the Anthropic path's ``sanitize_anthropic_model_id`` only +strips ANSI sequences and whitespace), so an unescaped value is remotely +reachable. + +Imports only the metrics module so the test stays free of heavy ML deps. +""" + +from __future__ import annotations + +import re + +import pytest + +from headroom.proxy.prometheus_metrics import PrometheusMetrics + +# A label whose value contains only unreserved characters or well-formed escape +# pairs. An unescaped quote inside a value stops this matching, which is exactly +# the failure a scraper hits. +_LABEL_RE = re.compile(r'([a-zA-Z_][a-zA-Z0-9_]*)="((?:[^"\\]|\\.)*)"') +_SAMPLE_RE = re.compile(r"^(?P[a-zA-Z_:][a-zA-Z0-9_:]*)\{(?P.*)\} \S+$") +_ESCAPE_RE = re.compile(r"\\(.)") +_UNESCAPE = {"n": "\n", '"': '"', "\\": "\\"} + + +def _unescape(value: str) -> str: + def replace(match: re.Match[str]) -> str: + char = match.group(1) + if char not in _UNESCAPE: + raise ValueError(f"undefined escape sequence '\\{char}' in {value!r}") + return _UNESCAPE[char] + + return _ESCAPE_RE.sub(replace, value) + + +def _parse_label_block(block: str) -> dict[str, str]: + """Parse ``key="value",key="value"`` the way a scraper would. + + Raises ``ValueError`` on anything the exposition grammar rejects, so a line + carrying an unescaped quote fails loudly instead of yielding a + plausible-looking dict. + """ + labels: dict[str, str] = {} + pos = 0 + while pos < len(block): + match = _LABEL_RE.match(block, pos) + if match is None: + raise ValueError(f"malformed label block at offset {pos}: {block!r}") + labels[match.group(1)] = _unescape(match.group(2)) + pos = match.end() + if pos < len(block): + if block[pos] != ",": + raise ValueError(f"expected ',' at offset {pos}: {block!r}") + pos += 1 + return labels + + +def _labelled_samples(text: str) -> list[tuple[str, dict[str, str]]]: + """Every labelled sample in a scrape, as (metric name, decoded labels). + + Raises on any line a scraper would reject — including the fragments an + unescaped line feed splits a sample into. + """ + samples: list[tuple[str, dict[str, str]]] = [] + for line in text.splitlines(): + if not line or line.startswith("#") or "{" not in line: + continue + match = _SAMPLE_RE.match(line) + if match is None: + raise ValueError(f"malformed sample line: {line!r}") + samples.append((match.group("name"), _parse_label_block(match.group("labels")))) + return samples + + +async def _record(metrics: PrometheusMetrics, **overrides: object) -> None: + kwargs: dict[str, object] = { + "provider": "anthropic", + "model": "claude-sonnet-4-5", + "input_tokens": 100, + "output_tokens": 20, + # tokens_saved=0 keeps the durable savings-ledger write out of the test. + "tokens_saved": 0, + "latency_ms": 10.0, + } + kwargs.update(overrides) + await metrics.record_request(**kwargs) # type: ignore[arg-type] + + +@pytest.mark.asyncio +async def test_quote_in_model_is_escaped() -> None: + metrics = PrometheusMetrics() + + await _record(metrics, model='claude-sonnet-4-5"evil') + + text = await metrics.export() + + assert 'headroom_requests_by_model{model="claude-sonnet-4-5\\"evil"} 1' in text + assert 'headroom_requests_by_model{model="claude-sonnet-4-5"evil"}' not in text + + +@pytest.mark.asyncio +async def test_quote_in_provider_is_escaped() -> None: + metrics = PrometheusMetrics() + + await _record(metrics, provider='anth"ropic') + + text = await metrics.export() + + assert 'headroom_requests_by_provider{provider="anth\\"ropic"} 1' in text + assert 'headroom_requests_by_provider{provider="anth"ropic"}' not in text + + +@pytest.mark.asyncio +async def test_backslash_and_newline_in_model_are_escaped() -> None: + metrics = PrometheusMetrics() + + await _record(metrics, model="back\\slash") + await _record(metrics, model="line\nfeed") + + text = await metrics.export() + + # Backslash first, so the escapes this inserts are not re-escaped. + assert 'headroom_requests_by_model{model="back\\\\slash"} 1' in text + assert 'headroom_requests_by_model{model="line\\nfeed"} 1' in text + # The line feed must not survive as a real newline splitting the sample. + assert "line\nfeed" not in text + + +@pytest.mark.asyncio +async def test_provider_cache_families_escape_provider() -> None: + # The families PR #2450 added inherit `provider` from the same parameter, + # so they need naming explicitly rather than assuming coverage. + metrics = PrometheusMetrics() + + await _record( + metrics, + provider='anth"ropic', + cache_read_tokens=40, + cache_write_tokens=60, + cache_write_5m_tokens=10, + cache_write_1h_tokens=50, + uncached_input_tokens=20, + ) + + text = await metrics.export() + + families = [ + "headroom_cache_read_tokens_total", + "headroom_cache_write_tokens_total", + "headroom_cache_write_ttl_tokens_total", + "headroom_cache_write_ttl_requests_total", + "headroom_uncached_input_tokens_total", + "headroom_provider_cache_requests_total", + "headroom_provider_cache_hit_requests_total", + "headroom_provider_cache_bust_total", + "headroom_provider_cache_bust_write_tokens_total", + ] + for family in families: + assert f'{family}{{provider="anth\\"ropic"' in text, f"{family} left provider raw" + + +@pytest.mark.asyncio +async def test_cache_miss_attribution_escapes_both_labels() -> None: + metrics = PrometheusMetrics() + + await metrics.record_cache_miss_attribution('anth"ropic', 'ttl"expiry') + + text = await metrics.export() + + assert ( + 'headroom_cache_miss_attribution_total{provider="anth\\"ropic",reason="ttl\\"expiry"} 1' + in text + ) + + +@pytest.mark.asyncio +async def test_no_emitted_label_value_is_malformed() -> None: + # The regression guard: poison every reachable label input, then read the + # whole scrape the way a scraper does. A future emission that forgets to + # escape fails here even when no assertion above names it. + metrics = PrometheusMetrics() + + # The model poison carries a comma and an inner quote. The parse alone + # can't catch comma-injection (this value raises on the quote first), so the + # round-trip assertion below is the real guard: after escaping, the value + # must decode back to the exact raw string, comma and all, rather than + # splitting into extra labels. + await _record( + metrics, + provider='pro"vider\\one', + model='mo"del,evil="1', + cache_read_tokens=40, + cache_write_tokens=60, + cache_write_5m_tokens=10, + cache_write_1h_tokens=50, + uncached_input_tokens=20, + ) + await metrics.record_cache_miss_attribution('pro"vider\\one', 'rea"son') + + samples = _labelled_samples(await metrics.export()) + + values = {value for _, labels in samples for value in labels.values()} + assert 'pro"vider\\one' in values, "provider did not round-trip through the escape" + assert 'mo"del,evil="1' in values, "model did not round-trip through the escape" + + +@pytest.mark.asyncio +async def test_non_string_label_values_are_coerced() -> None: + # A JSON body can carry `"model": 123`, and the handlers pass the decoded + # value through untouched (handlers/openai.py reads body.get("model")). The + # hand-rolled f-strings used to call str() implicitly, so escaping has to + # keep tolerating a non-str. /metrics has no error handling around export(), + # and the key survives in the dict, so a raise here would take out every + # later scrape too. + metrics = PrometheusMetrics() + + await _record(metrics, provider=456, model=123, cache_read_tokens=5, cache_write_tokens=5) + await metrics.record_cache_miss_attribution(456, 789) + + text = await metrics.export() + + assert 'headroom_requests_by_model{model="123"} 1' in text + assert 'headroom_requests_by_provider{provider="456"} 1' in text + assert 'headroom_cache_read_tokens_total{provider="456"}' in text + assert 'headroom_cache_miss_attribution_total{provider="456",reason="789"} 1' in text + + +@pytest.mark.asyncio +async def test_well_formed_values_are_emitted_unchanged() -> None: + metrics = PrometheusMetrics() + + await _record(metrics) + + text = await metrics.export() + + assert 'headroom_requests_by_provider{provider="anthropic"} 1' in text + assert 'headroom_requests_by_model{model="claude-sonnet-4-5"} 1' in text + + +@pytest.mark.asyncio +async def test_export_is_utf8_encodable_with_surrogate_model() -> None: + # `/metrics` renders the whole body with `.encode("utf-8")` (server.py). A + # client can decode a lone surrogate from JSON (`{"model": "x-\ud83d-y"}`) — + # a valid str that is NOT UTF-8-encodable and passes escaping untouched. It + # would raise in the response encoder and, because the poisoned key persists + # in requests_by_model, 500 every later scrape until restart. Escaping must + # leave the whole export encodable. + metrics = PrometheusMetrics() + + await _record(metrics, model="x-\ud83d-y") + await _record(metrics, model="clean-model") # a healthy series alongside + + text = await metrics.export() + + # The load-bearing assertion: the body a scraper receives must encode. + text.encode("utf-8") + # And the healthy series is still readable, i.e. the poison did not corrupt + # the surrounding output. + assert 'headroom_requests_by_model{model="clean-model"} 1' in text + # Legitimate astral characters (a real emoji is one code point, encodable) + # are preserved, not scrubbed — only un-encodable lone surrogates change. + metrics2 = PrometheusMetrics() + await _record(metrics2, model="gpt-\U0001f600") + assert 'headroom_requests_by_model{model="gpt-\U0001f600"} 1' in await metrics2.export()