From 41a2dfefdb30de2301ddbb369328d463f4388825 Mon Sep 17 00:00:00 2001 From: Peter Wilson Date: Tue, 15 Sep 2026 20:18:55 +0100 Subject: [PATCH 1/8] test(metrics): pin the metric families the scrape exposes Moving each counter beside the code that increments it (#1177) must not rename a metric or change its labels, because dashboards and alerts outside this repository key on them. Pin the families before any move so each move is checked against the same set. --- tests/unit/test_gateway_metrics.py | 40 ++++++++++++++++++++++++++++++ 1 file changed, 40 insertions(+) diff --git a/tests/unit/test_gateway_metrics.py b/tests/unit/test_gateway_metrics.py index ec3760ce55..efc739cad3 100644 --- a/tests/unit/test_gateway_metrics.py +++ b/tests/unit/test_gateway_metrics.py @@ -121,6 +121,46 @@ def test_record_abandoned_attempt_labels_by_reason_and_position() -> None: assert _sample("gateway_abandoned_attempts_total", upstream_labels) - before_upstream == 1.0 +# Every family the gateway registers, as (name, type, label names). A metric +# may move to the module that increments it, but the scrape is an external +# contract: dashboards, recording rules and alerts outside this repository +# key on these names and labels. +_EXPOSED_FAMILIES: set[tuple[str, str, tuple[str, ...]]] = { + ("gateway_abandoned_attempts", "counter", ("provider", "model", "reason", "position")), + ("gateway_active_requests", "gauge", ()), + ("gateway_auth_failures", "counter", ("reason",)), + ("gateway_budget_exceeded", "counter", ()), + ("gateway_inline_cost_settlements", "counter", ("outcome",)), + ("gateway_rate_limit_hits", "counter", ()), + ("gateway_request_cost_dollars", "histogram", ("provider", "model")), + ("gateway_request_duration_seconds", "histogram", ("method", "endpoint", "api_version")), + ("gateway_requests", "counter", ("method", "endpoint", "api_version", "status")), + ("gateway_tokens", "counter", ("provider", "model", "type")), + ("gateway_usage_log_batch_size", "histogram", ("writer",)), + ("gateway_usage_log_flush_duration_seconds", "histogram", ("writer", "result")), + ("gateway_usage_log_queue_depth", "gauge", ()), + ("gateway_usage_log_rows", "counter", ("writer", "result")), +} + + +def test_scrape_exposes_the_pinned_families() -> None: + """The set of gateway metric families, with their types and label names, is fixed. + + A labeled family with no series yet shows only its HELP and TYPE lines in a + scrape, so the label names are read off the collectors rather than the text. + """ + import gateway.main # noqa: F401 # imports every module that registers a metric + + families: set[tuple[str, str, tuple[str, ...]]] = set() + for collector in REGISTRY._collector_to_names: + describe = getattr(collector, "describe", collector.collect) + for metric in describe(): + if metric.name.startswith("gateway_"): + families.add((metric.name, metric.type, tuple(getattr(collector, "_labelnames", ())))) + + assert families == _EXPOSED_FAMILIES + + def test_config_enable_metrics_defaults_to_false() -> None: config = GatewayConfig() assert config.enable_metrics is False From e1042c3549eae831607ceecb4217e643ffeaa448 Mon Sep 17 00:00:00 2001 From: Peter Wilson Date: Tue, 15 Sep 2026 20:20:28 +0100 Subject: [PATCH 2/8] refactor(metrics): move the rate limit counter beside the limiter The rate limiter is the only code that increments gateway_rate_limit_hits, so the counter now lives in rate_limit.py, registered on the shared registry. The wrapper function went with it: a one-line helper with a single caller in the same module said nothing the counter does not. Part of #1177. --- src/gateway/metrics.py | 11 ----------- src/gateway/rate_limit.py | 11 +++++++++-- tests/unit/test_gateway_metrics.py | 25 ------------------------- tests/unit/test_rate_limiter_core.py | 15 +++++++++++++++ 4 files changed, 24 insertions(+), 38 deletions(-) diff --git a/src/gateway/metrics.py b/src/gateway/metrics.py index 8181e78b66..b78cc429e8 100644 --- a/src/gateway/metrics.py +++ b/src/gateway/metrics.py @@ -77,12 +77,6 @@ registry=REGISTRY, ) -RATE_LIMIT_HITS = Counter( - "gateway_rate_limit_hits", - "Total number of rate limit hits", - registry=REGISTRY, -) - BUDGET_EXCEEDED = Counter( "gateway_budget_exceeded", "Total number of budget exceeded events", @@ -234,11 +228,6 @@ def record_inline_cost_settlement(outcome: str) -> None: INLINE_COST_SETTLEMENTS.labels(outcome=outcome).inc() -def record_rate_limit_hit() -> None: - """Record a rate limit hit.""" - RATE_LIMIT_HITS.inc() - - def record_budget_exceeded() -> None: """Record a budget exceeded event.""" BUDGET_EXCEEDED.inc() diff --git a/src/gateway/rate_limit.py b/src/gateway/rate_limit.py index 775849cc04..1cc3d3749e 100644 --- a/src/gateway/rate_limit.py +++ b/src/gateway/rate_limit.py @@ -6,8 +6,15 @@ from dataclasses import dataclass from fastapi import HTTPException, Request, status +from prometheus_client import Counter -from gateway.metrics import record_rate_limit_hit +from gateway.metrics import REGISTRY + +RATE_LIMIT_HITS = Counter( + "gateway_rate_limit_hits", + "Total number of rate limit hits", + registry=REGISTRY, +) @dataclass @@ -59,7 +66,7 @@ def check(self, user_id: str) -> RateLimitInfo: if len(timestamps) >= self._rpm: oldest = timestamps[0] retry_after = math.ceil(oldest - cutoff) - record_rate_limit_hit() + RATE_LIMIT_HITS.inc() raise HTTPException( status_code=status.HTTP_429_TOO_MANY_REQUESTS, detail="Rate limit exceeded", diff --git a/tests/unit/test_gateway_metrics.py b/tests/unit/test_gateway_metrics.py index efc739cad3..c3250fef38 100644 --- a/tests/unit/test_gateway_metrics.py +++ b/tests/unit/test_gateway_metrics.py @@ -18,7 +18,6 @@ record_budget_exceeded, record_cost, record_inline_cost_settlement, - record_rate_limit_hit, record_tokens, ) @@ -72,14 +71,6 @@ def test_record_inline_cost_settlement_increments_counter(outcome: str) -> None: assert _sample("gateway_inline_cost_settlements_total", labels) - before == 1.0 -def test_record_rate_limit_hit_increments_counter() -> None: - before = _sample("gateway_rate_limit_hits_total") - - record_rate_limit_hit() - - assert _sample("gateway_rate_limit_hits_total") - before == 1.0 - - def test_record_budget_exceeded_increments_counter() -> None: before = _sample("gateway_budget_exceeded_total") @@ -331,22 +322,6 @@ def test_active_requests_returns_to_zero() -> None: assert _sample("gateway_active_requests") == before -def test_rate_limiter_records_metric_on_429() -> None: - """RateLimiter.check() records a metric before raising 429.""" - from gateway.rate_limit import RateLimiter - - limiter = RateLimiter(rpm=1) - limiter.check("metric-rl-user") - - before = _sample("gateway_rate_limit_hits_total") - - with pytest.raises(HTTPException) as exc_info: - limiter.check("metric-rl-user") - - assert exc_info.value.status_code == 429 - assert _sample("gateway_rate_limit_hits_total") - before == 1.0 - - @pytest.mark.skipif(not os.path.exists("/proc/stat"), reason="ProcessCollector needs /proc") def test_metrics_expose_process_memory() -> None: """Without this the only way to read the resident set is a shell in the container.""" diff --git a/tests/unit/test_rate_limiter_core.py b/tests/unit/test_rate_limiter_core.py index ffb5a02fb1..7570a76e0c 100644 --- a/tests/unit/test_rate_limiter_core.py +++ b/tests/unit/test_rate_limiter_core.py @@ -8,6 +8,7 @@ from pydantic import ValidationError from gateway.core.config import GatewayConfig +from gateway.metrics import REGISTRY from gateway.rate_limit import RateLimiter, RateLimitInfo @@ -140,3 +141,17 @@ def test_config_accepts_positive_rate_limit() -> None: def test_config_accepts_none_rate_limit() -> None: config = GatewayConfig(rate_limit_rpm=None) assert config.rate_limit_rpm is None + + +def test_check_records_metric_on_429() -> None: + """RateLimiter.check() records a metric before raising 429.""" + limiter = RateLimiter(rpm=1) + limiter.check("metric-rl-user") + + before = REGISTRY.get_sample_value("gateway_rate_limit_hits_total") or 0.0 + + with pytest.raises(HTTPException) as exc_info: + limiter.check("metric-rl-user") + + assert exc_info.value.status_code == 429 + assert (REGISTRY.get_sample_value("gateway_rate_limit_hits_total") or 0.0) - before == 1.0 From 17d3a69812669aaf5aacb060d2b66bb74cf099ef Mon Sep 17 00:00:00 2001 From: Peter Wilson Date: Tue, 15 Sep 2026 20:21:30 +0100 Subject: [PATCH 3/8] refactor(metrics): move the budget exceeded counter into budget_service budget_service is the only code that increments gateway_budget_exceeded, so the counter now lives there, registered on the shared registry, and the three call sites increment it directly. The wrapper function and its unit test go: the counter's real path is covered by the integration test that drives a zero-limit budget to a 403. Part of #1177. --- src/gateway/metrics.py | 11 ----------- src/gateway/services/budget_service.py | 15 +++++++++++---- tests/unit/test_gateway_metrics.py | 9 --------- 3 files changed, 11 insertions(+), 24 deletions(-) diff --git a/src/gateway/metrics.py b/src/gateway/metrics.py index b78cc429e8..68f61ee33a 100644 --- a/src/gateway/metrics.py +++ b/src/gateway/metrics.py @@ -77,12 +77,6 @@ registry=REGISTRY, ) -BUDGET_EXCEEDED = Counter( - "gateway_budget_exceeded", - "Total number of budget exceeded events", - registry=REGISTRY, -) - AUTH_FAILURES = Counter( "gateway_auth_failures", "Total number of authentication failures", @@ -228,11 +222,6 @@ def record_inline_cost_settlement(outcome: str) -> None: INLINE_COST_SETTLEMENTS.labels(outcome=outcome).inc() -def record_budget_exceeded() -> None: - """Record a budget exceeded event.""" - BUDGET_EXCEEDED.inc() - - def record_auth_failure(reason: str) -> None: """Record an authentication failure.""" AUTH_FAILURES.labels(reason=reason).inc() diff --git a/src/gateway/services/budget_service.py b/src/gateway/services/budget_service.py index 2d3cf6dac2..c7867f9403 100644 --- a/src/gateway/services/budget_service.py +++ b/src/gateway/services/budget_service.py @@ -10,13 +10,14 @@ from any_llm import AnyLLM from any_llm.exceptions import AnyLLMError from fastapi import HTTPException, status +from prometheus_client import Counter from sqlalchemy import select, update from sqlalchemy.exc import SQLAlchemyError from sqlalchemy.ext.asyncio import AsyncSession from gateway.core.metered_pricing import estimate_metered_cost from gateway.log_config import logger -from gateway.metrics import record_budget_exceeded +from gateway.metrics import REGISTRY from gateway.models.entities import MAX_COUNT_LIMIT, Budget, BudgetResetLog, ModelPricing, User from gateway.models.money import to_usd from gateway.repositories.users_repository import get_active_user @@ -35,6 +36,12 @@ from gateway.services.scoped_budget_service import settle as settle_scoped from gateway.types.budget_state import BudgetState +BUDGET_EXCEEDED = Counter( + "gateway_budget_exceeded", + "Total number of budget exceeded events", + registry=REGISTRY, +) + # Every counter in this module is a ``NUMERIC(18, 6)`` column, so the constants # the SQL is built from are ``Decimal`` too. A bare ``0.0`` in a CASE arm or a # ``.values()`` would make PostgreSQL resolve the whole expression as double @@ -613,7 +620,7 @@ async def reserve_budget( db, scoped, usd, tokens=held_tokens, requests=held_requests, new_request=new_request ) if refused is not None: - record_budget_exceeded() + BUDGET_EXCEEDED.inc() axis = await blocked_axis( db, refused, @@ -715,7 +722,7 @@ def guards_for(committed: Any, held_amount: Any, cap: Any) -> list[Any]: await db.commit() if not getattr(result, "rowcount", 0): - record_budget_exceeded() + BUDGET_EXCEEDED.inc() # The scoped ceilings admitted this request and are already holding it, so # give every axis back before rejecting. Without this the holds would leak # on every per-user refusal and permanently shrink each ceiling. @@ -1026,7 +1033,7 @@ async def increase_reservation( db, handle.scoped_budgets, additional, tokens=grown_tokens, requests=0, new_request=False ) if refused is not None: - record_budget_exceeded() + BUDGET_EXCEEDED.inc() axis = await blocked_axis( db, refused, diff --git a/tests/unit/test_gateway_metrics.py b/tests/unit/test_gateway_metrics.py index c3250fef38..e4432208b2 100644 --- a/tests/unit/test_gateway_metrics.py +++ b/tests/unit/test_gateway_metrics.py @@ -15,7 +15,6 @@ metrics_endpoint, record_abandoned_attempt, record_auth_failure, - record_budget_exceeded, record_cost, record_inline_cost_settlement, record_tokens, @@ -71,14 +70,6 @@ def test_record_inline_cost_settlement_increments_counter(outcome: str) -> None: assert _sample("gateway_inline_cost_settlements_total", labels) - before == 1.0 -def test_record_budget_exceeded_increments_counter() -> None: - before = _sample("gateway_budget_exceeded_total") - - record_budget_exceeded() - - assert _sample("gateway_budget_exceeded_total") - before == 1.0 - - def test_record_auth_failure_increments_counter() -> None: labels = {"reason": "unit-test-reason"} before = _sample("gateway_auth_failures_total", labels) From 4b054b1176f1d691a8791306e6f08e170929b29d Mon Sep 17 00:00:00 2001 From: Peter Wilson Date: Tue, 15 Sep 2026 20:21:34 +0100 Subject: [PATCH 4/8] refactor(metrics): move the usage log writer metrics into log_writer The two writers are the only code that touches the four usage log metrics, so they now live in log_writer.py, registered on the shared registry, under names that read in that module (ROWS, QUEUE_DEPTH, BATCH_SIZE, FLUSH_DURATION). The lowercase aliases metrics.py exported for the writer are gone with the move. Part of #1177. --- src/gateway/metrics.py | 33 ------------------- src/gateway/services/log_writer.py | 52 ++++++++++++++++++++++-------- 2 files changed, 38 insertions(+), 47 deletions(-) diff --git a/src/gateway/metrics.py b/src/gateway/metrics.py index 68f61ee33a..1a610b4a86 100644 --- a/src/gateway/metrics.py +++ b/src/gateway/metrics.py @@ -84,33 +84,6 @@ registry=REGISTRY, ) -LOG_WRITER_QUEUE_DEPTH = Gauge( - "gateway_usage_log_queue_depth", - "Number of usage log entries waiting to be written", - registry=REGISTRY, -) - -LOG_WRITER_BATCH_SIZE = Histogram( - "gateway_usage_log_batch_size", - "Number of rows per flush batch", - ["writer"], - registry=REGISTRY, -) - -LOG_WRITER_FLUSH_DURATION = Histogram( - "gateway_usage_log_flush_duration_seconds", - "Time spent flushing usage log batches", - ["writer", "result"], - registry=REGISTRY, -) - -LOG_WRITER_ROWS = Counter( - "gateway_usage_log_rows", - "Total usage log rows by outcome", - ["writer", "result"], - registry=REGISTRY, -) - _PROMETHEUS_CONTENT_TYPE = "text/plain; version=0.0.4; charset=utf-8" @@ -225,9 +198,3 @@ def record_inline_cost_settlement(outcome: str) -> None: def record_auth_failure(reason: str) -> None: """Record an authentication failure.""" AUTH_FAILURES.labels(reason=reason).inc() - - -log_writer_queue_depth = LOG_WRITER_QUEUE_DEPTH -log_writer_batch_size = LOG_WRITER_BATCH_SIZE -log_writer_flush_duration = LOG_WRITER_FLUSH_DURATION -log_writer_rows = LOG_WRITER_ROWS diff --git a/src/gateway/services/log_writer.py b/src/gateway/services/log_writer.py index 48b813589c..e19144cba4 100644 --- a/src/gateway/services/log_writer.py +++ b/src/gateway/services/log_writer.py @@ -6,16 +6,40 @@ import time from typing import Protocol +from prometheus_client import Counter, Gauge, Histogram + from gateway.core.database import DATABASE_ERRORS, create_log_session from gateway.log_config import logger -from gateway.metrics import ( - log_writer_batch_size, - log_writer_flush_duration, - log_writer_queue_depth, - log_writer_rows, -) +from gateway.metrics import REGISTRY from gateway.models.entities import UsageLog +QUEUE_DEPTH = Gauge( + "gateway_usage_log_queue_depth", + "Number of usage log entries waiting to be written", + registry=REGISTRY, +) + +BATCH_SIZE = Histogram( + "gateway_usage_log_batch_size", + "Number of rows per flush batch", + ["writer"], + registry=REGISTRY, +) + +FLUSH_DURATION = Histogram( + "gateway_usage_log_flush_duration_seconds", + "Time spent flushing usage log batches", + ["writer", "result"], + registry=REGISTRY, +) + +ROWS = Counter( + "gateway_usage_log_rows", + "Total usage log rows by outcome", + ["writer", "result"], + registry=REGISTRY, +) + class LogWriter(Protocol): async def put(self, log: UsageLog) -> None: ... @@ -37,11 +61,11 @@ async def put(self, log: UsageLog) -> None: # and, under the batch writer, lag the budget gate. db.add(log) await db.commit() - log_writer_rows.labels(writer="single", result="written").inc() + ROWS.labels(writer="single", result="written").inc() except DATABASE_ERRORS as e: # pragma: no cover - defensive logging await db.rollback() logger.error("SingleLogWriter failed: %s", e) - log_writer_rows.labels(writer="single", result="dropped").inc() + ROWS.labels(writer="single", result="dropped").inc() async def start(self) -> None: pass @@ -61,7 +85,7 @@ def __init__(self, max_batch: int = 100, flush_interval: float = 1.0) -> None: async def put(self, log: UsageLog) -> None: await self._queue.put(log) - log_writer_queue_depth.set(self._queue.qsize()) + QUEUE_DEPTH.set(self._queue.qsize()) async def start(self) -> None: self._task = asyncio.create_task(self._run()) @@ -106,7 +130,7 @@ async def _collect_batch(self) -> list[UsageLog]: async def _flush(self, batch: list[UsageLog]) -> None: start = time.monotonic() - log_writer_batch_size.labels(writer="batch").observe(len(batch)) + BATCH_SIZE.labels(writer="batch").observe(len(batch)) try: async with create_log_session() as db: # Spend is reconciled inline via the budget reservation path, not @@ -114,12 +138,12 @@ async def _flush(self, batch: list[UsageLog]) -> None: for log in batch: db.add(log) await db.commit() - log_writer_rows.labels(writer="batch", result="written").inc(len(batch)) - log_writer_flush_duration.labels(writer="batch", result="ok").observe(time.monotonic() - start) + ROWS.labels(writer="batch", result="written").inc(len(batch)) + FLUSH_DURATION.labels(writer="batch", result="ok").observe(time.monotonic() - start) except DATABASE_ERRORS as e: # pragma: no cover - defensive logging logger.error("BatchLogWriter flush failed, dropping %d rows: %s", len(batch), e) - log_writer_rows.labels(writer="batch", result="dropped").inc(len(batch)) - log_writer_flush_duration.labels(writer="batch", result="error").observe(time.monotonic() - start) + ROWS.labels(writer="batch", result="dropped").inc(len(batch)) + FLUSH_DURATION.labels(writer="batch", result="error").observe(time.monotonic() - start) async def _flush_all(self) -> None: batch: list[UsageLog] = [] From 7dec3509ea1016ca9855c6bc6735097ef17a4697 Mon Sep 17 00:00:00 2001 From: Peter Wilson Date: Tue, 15 Sep 2026 20:23:13 +0100 Subject: [PATCH 5/8] refactor(metrics): move the auth failure counter into api.deps api.deps is where nearly every authentication failure is counted, so gateway_auth_failures and its record helper now live there, registered on the shared registry. The three sign-in routes already import their dependencies from api.deps and now take the helper from the same place. The helper's unit test goes: the bearer-token and key-verification tests already assert the counter through the real callers. Part of #1177. --- src/gateway/api/deps.py | 15 ++++++++++++++- src/gateway/api/routes/auth_oauth.py | 3 +-- src/gateway/api/routes/auth_session.py | 3 +-- src/gateway/api/routes/auth_webauthn.py | 3 +-- src/gateway/metrics.py | 12 ------------ tests/unit/test_gateway_metrics.py | 10 ---------- 6 files changed, 17 insertions(+), 29 deletions(-) diff --git a/src/gateway/api/deps.py b/src/gateway/api/deps.py index 65ba1e7366..8ee232652c 100644 --- a/src/gateway/api/deps.py +++ b/src/gateway/api/deps.py @@ -6,6 +6,7 @@ from typing import Annotated from fastapi import Depends, HTTPException, Request, status +from prometheus_client import Counter from sqlalchemy import select, update from sqlalchemy.ext.asyncio import AsyncSession @@ -15,7 +16,7 @@ from gateway.core.database import DATABASE_ERRORS, create_session, get_db from gateway.core.feature import CoreFeature from gateway.log_config import logger -from gateway.metrics import record_auth_failure +from gateway.metrics import REGISTRY from gateway.models.entities import APIKey from gateway.models.tenancy import User as TenancyUser from gateway.ports.billing_port import BillingPort @@ -39,6 +40,18 @@ _config: GatewayConfig | None = None _LAST_USED_UPDATE_INTERVAL_SECONDS = 300 +AUTH_FAILURES = Counter( + "gateway_auth_failures", + "Total number of authentication failures", + ["reason"], + registry=REGISTRY, +) + + +def record_auth_failure(reason: str) -> None: + """Record an authentication failure.""" + AUTH_FAILURES.labels(reason=reason).inc() + def _as_utc(value: datetime | None) -> datetime | None: """Return ``value`` as a timezone-aware datetime in UTC. diff --git a/src/gateway/api/routes/auth_oauth.py b/src/gateway/api/routes/auth_oauth.py index e55c7e8515..8b95eb671c 100644 --- a/src/gateway/api/routes/auth_oauth.py +++ b/src/gateway/api/routes/auth_oauth.py @@ -46,7 +46,7 @@ from sqlalchemy.exc import SQLAlchemyError from sqlalchemy.ext.asyncio import AsyncSession -from gateway.api.deps import IdentityProviderPortDep, get_config, get_db +from gateway.api.deps import IdentityProviderPortDep, get_config, get_db, record_auth_failure from gateway.api.routes._public_auth import throttle_public_auth # The same refusal the password and passkey sign-ins carry, imported rather than @@ -55,7 +55,6 @@ from gateway.api.routes.auth_session import MAINTENANCE_MODE_REFUSAL from gateway.core.config import OAUTH_PROVIDERS, GatewayConfig from gateway.log_config import logger -from gateway.metrics import record_auth_failure from gateway.services.dashboard_session_service import ( apply_session_cookie, create_dashboard_session, diff --git a/src/gateway/api/routes/auth_session.py b/src/gateway/api/routes/auth_session.py index 161e90d19c..7258a86c7c 100644 --- a/src/gateway/api/routes/auth_session.py +++ b/src/gateway/api/routes/auth_session.py @@ -48,10 +48,9 @@ from sqlalchemy.exc import SQLAlchemyError from sqlalchemy.ext.asyncio import AsyncSession -from gateway.api.deps import get_config, get_db, is_valid_master_key +from gateway.api.deps import get_config, get_db, is_valid_master_key, record_auth_failure from gateway.core.config import GatewayConfig from gateway.log_config import logger -from gateway.metrics import record_auth_failure from gateway.models.tenancy import User as TenancyUser from gateway.rate_limit import RateLimiter from gateway.services.dashboard_session_service import ( diff --git a/src/gateway/api/routes/auth_webauthn.py b/src/gateway/api/routes/auth_webauthn.py index 6d5621f7e1..e03152a22b 100644 --- a/src/gateway/api/routes/auth_webauthn.py +++ b/src/gateway/api/routes/auth_webauthn.py @@ -36,7 +36,7 @@ from sqlalchemy.exc import SQLAlchemyError from sqlalchemy.ext.asyncio import AsyncSession -from gateway.api.deps import CurrentIdentity, get_config, get_db +from gateway.api.deps import CurrentIdentity, get_config, get_db, record_auth_failure from gateway.api.routes._public_auth import throttle_public_auth # The refusal is the same refusal, so it is imported rather than restated: the @@ -45,7 +45,6 @@ from gateway.api.routes.auth_session import MAINTENANCE_MODE_REFUSAL from gateway.core.config import GatewayConfig from gateway.log_config import logger -from gateway.metrics import record_auth_failure from gateway.models.tenancy import ( MAX_WEBAUTHN_CREDENTIAL_NAME, WebAuthnCredentialPublic, diff --git a/src/gateway/metrics.py b/src/gateway/metrics.py index 1a610b4a86..226cc12af9 100644 --- a/src/gateway/metrics.py +++ b/src/gateway/metrics.py @@ -77,13 +77,6 @@ registry=REGISTRY, ) -AUTH_FAILURES = Counter( - "gateway_auth_failures", - "Total number of authentication failures", - ["reason"], - registry=REGISTRY, -) - _PROMETHEUS_CONTENT_TYPE = "text/plain; version=0.0.4; charset=utf-8" @@ -193,8 +186,3 @@ def record_abandoned_attempt(provider: str, model: str, reason: str, position: i def record_inline_cost_settlement(outcome: str) -> None: """Record an attached, unattached, or timed-out inline settlement.""" INLINE_COST_SETTLEMENTS.labels(outcome=outcome).inc() - - -def record_auth_failure(reason: str) -> None: - """Record an authentication failure.""" - AUTH_FAILURES.labels(reason=reason).inc() diff --git a/tests/unit/test_gateway_metrics.py b/tests/unit/test_gateway_metrics.py index e4432208b2..cf5d77f8c6 100644 --- a/tests/unit/test_gateway_metrics.py +++ b/tests/unit/test_gateway_metrics.py @@ -14,7 +14,6 @@ _endpoint_label, metrics_endpoint, record_abandoned_attempt, - record_auth_failure, record_cost, record_inline_cost_settlement, record_tokens, @@ -70,15 +69,6 @@ def test_record_inline_cost_settlement_increments_counter(outcome: str) -> None: assert _sample("gateway_inline_cost_settlements_total", labels) - before == 1.0 -def test_record_auth_failure_increments_counter() -> None: - labels = {"reason": "unit-test-reason"} - before = _sample("gateway_auth_failures_total", labels) - - record_auth_failure("unit-test-reason") - - assert _sample("gateway_auth_failures_total", labels) - before == 1.0 - - def test_record_abandoned_attempt_increments_counter() -> None: labels = {"provider": "ab-prov", "model": "ab-model", "reason": "timeout", "position": "0"} before = _sample("gateway_abandoned_attempts_total", labels) From 73f7311acec2b549168e1aea2ae0d9c3e592a322 Mon Sep 17 00:00:00 2001 From: Peter Wilson Date: Tue, 15 Sep 2026 20:23:34 +0100 Subject: [PATCH 6/8] refactor(metrics): move the abandoned attempt counter into _platform Three modules count abandoned attempts: the platform runner, the local attempt walker, and the pipeline. The runner module already owns the upstream error classification the other two import, so gateway_abandoned_attempts and its record helper now live beside it, on the shared registry, and the walker and the pipeline take the helper from there. The helper's tests move next to the runner's own abandonment tests. Part of #1177. --- src/gateway/api/routes/_attempts.py | 2 +- src/gateway/api/routes/_pipeline.py | 3 ++- src/gateway/api/routes/_platform.py | 23 ++++++++++++++++++- src/gateway/metrics.py | 19 ---------------- tests/unit/test_gateway_metrics.py | 25 -------------------- tests/unit/test_run_platform_attempts.py | 29 +++++++++++++++++++++++- 6 files changed, 53 insertions(+), 48 deletions(-) diff --git a/src/gateway/api/routes/_attempts.py b/src/gateway/api/routes/_attempts.py index 452895e954..dd43c2cb44 100644 --- a/src/gateway/api/routes/_attempts.py +++ b/src/gateway/api/routes/_attempts.py @@ -38,10 +38,10 @@ from gateway.api.routes._platform import ( _provider_failure_http_exc, is_provider_billing_error, + record_abandoned_attempt, upstream_exception_shape, ) from gateway.log_config import logger -from gateway.metrics import record_abandoned_attempt from gateway.services.mcp_loop import MaxToolIterationsExceeded from gateway.services.sandbox_backend import SandboxNotReachableError from gateway.services.web_retrieval_backend import WebSearchNotReachableError diff --git a/src/gateway/api/routes/_pipeline.py b/src/gateway/api/routes/_pipeline.py index 54f9cbaf71..aac65fd6c8 100644 --- a/src/gateway/api/routes/_pipeline.py +++ b/src/gateway/api/routes/_pipeline.py @@ -85,6 +85,7 @@ _resolve_platform_mcp_servers, _resolve_platform_web_search, is_provider_billing_error, + record_abandoned_attempt, run_platform_attempts, upstream_error_message, upstream_exception_chain, @@ -117,7 +118,7 @@ ) from gateway.inflight import track_request from gateway.log_config import logger -from gateway.metrics import record_abandoned_attempt, record_cost, record_inline_cost_settlement, record_tokens +from gateway.metrics import record_cost, record_inline_cost_settlement, record_tokens from gateway.model_labeling import relabel_model from gateway.models.entities import APIKey, ModelPricing, UsageLog from gateway.models.guardrails import GuardrailConfig diff --git a/src/gateway/api/routes/_platform.py b/src/gateway/api/routes/_platform.py index fe1bbd2b9f..970f90922b 100644 --- a/src/gateway/api/routes/_platform.py +++ b/src/gateway/api/routes/_platform.py @@ -25,6 +25,7 @@ from fastapi import HTTPException, Request, status from openai import APIConnectionError as _OpenAIAPIConnectionError from openai import APITimeoutError as _OpenAIAPITimeoutError +from prometheus_client import Counter from pydantic import BaseModel, Field, ValidationError from gateway.core.config import GatewayConfig @@ -34,7 +35,7 @@ cache_write_tokens_of, ) from gateway.log_config import logger -from gateway.metrics import record_abandoned_attempt +from gateway.metrics import REGISTRY from gateway.models.mcp import McpServerConfig, ResolvedMcpServer from gateway.services.bedrock_gateway_auth import build_bedrock_client_args from gateway.services.mcp_loop import MaxToolIterationsExceeded @@ -49,6 +50,26 @@ T = TypeVar("T") +ABANDONED_ATTEMPTS = Counter( + "gateway_abandoned_attempts", + "Total upstream attempts abandoned before their first chunk (provider fallback / timeout waste)", + ["provider", "model", "reason", "position"], + registry=REGISTRY, +) + + +def record_abandoned_attempt(provider: str, model: str, reason: str, position: int) -> None: + """Record an upstream attempt abandoned before it produced its first chunk. + + ``reason`` is one of ``timeout`` (the first-chunk wait elapsed), + ``build_error`` (opening the upstream stream failed), or ``upstream_error`` + (the upstream raised before yielding a chunk). ``position`` is the attempt's + index in the resolved routing plan; label cardinality stays bounded by the + plan length. + """ + ABANDONED_ATTEMPTS.labels(provider=provider, model=model, reason=reason, position=str(position)).inc() + + # Status codes returned by the platform's usage-report endpoint that the # gateway should NOT retry. Auth, payment-required, not-found, conflict, gone, # and unprocessable are all permanent rejection signals: retrying would just diff --git a/src/gateway/metrics.py b/src/gateway/metrics.py index 226cc12af9..89302e6c0d 100644 --- a/src/gateway/metrics.py +++ b/src/gateway/metrics.py @@ -63,13 +63,6 @@ registry=REGISTRY, ) -ABANDONED_ATTEMPTS = Counter( - "gateway_abandoned_attempts", - "Total upstream attempts abandoned before their first chunk (provider fallback / timeout waste)", - ["provider", "model", "reason", "position"], - registry=REGISTRY, -) - INLINE_COST_SETTLEMENTS = Counter( "gateway_inline_cost_settlements", "Inline platform cost settlement outcomes on the hybrid response path", @@ -171,18 +164,6 @@ def record_cost(provider: str, model: str, cost: float) -> None: REQUEST_COST_DOLLARS.labels(provider=provider, model=model).observe(cost) -def record_abandoned_attempt(provider: str, model: str, reason: str, position: int) -> None: - """Record an upstream attempt abandoned before it produced its first chunk. - - ``reason`` is one of ``timeout`` (the first-chunk wait elapsed), - ``build_error`` (opening the upstream stream failed), or ``upstream_error`` - (the upstream raised before yielding a chunk). ``position`` is the attempt's - index in the resolved routing plan; label cardinality stays bounded by the - plan length. - """ - ABANDONED_ATTEMPTS.labels(provider=provider, model=model, reason=reason, position=str(position)).inc() - - def record_inline_cost_settlement(outcome: str) -> None: """Record an attached, unattached, or timed-out inline settlement.""" INLINE_COST_SETTLEMENTS.labels(outcome=outcome).inc() diff --git a/tests/unit/test_gateway_metrics.py b/tests/unit/test_gateway_metrics.py index cf5d77f8c6..1c07788a21 100644 --- a/tests/unit/test_gateway_metrics.py +++ b/tests/unit/test_gateway_metrics.py @@ -13,7 +13,6 @@ MetricsMiddleware, _endpoint_label, metrics_endpoint, - record_abandoned_attempt, record_cost, record_inline_cost_settlement, record_tokens, @@ -69,30 +68,6 @@ def test_record_inline_cost_settlement_increments_counter(outcome: str) -> None: assert _sample("gateway_inline_cost_settlements_total", labels) - before == 1.0 -def test_record_abandoned_attempt_increments_counter() -> None: - labels = {"provider": "ab-prov", "model": "ab-model", "reason": "timeout", "position": "0"} - before = _sample("gateway_abandoned_attempts_total", labels) - - record_abandoned_attempt("ab-prov", "ab-model", "timeout", 0) - - assert _sample("gateway_abandoned_attempts_total", labels) - before == 1.0 - - -def test_record_abandoned_attempt_labels_by_reason_and_position() -> None: - """Each (reason, position) pair is its own series so operators can spot which - plan entry and failure phase dominates the fallback waste.""" - build_labels = {"provider": "ab-prov2", "model": "ab-model2", "reason": "build_error", "position": "1"} - upstream_labels = {"provider": "ab-prov2", "model": "ab-model2", "reason": "upstream_error", "position": "2"} - before_build = _sample("gateway_abandoned_attempts_total", build_labels) - before_upstream = _sample("gateway_abandoned_attempts_total", upstream_labels) - - record_abandoned_attempt("ab-prov2", "ab-model2", "build_error", 1) - record_abandoned_attempt("ab-prov2", "ab-model2", "upstream_error", 2) - - assert _sample("gateway_abandoned_attempts_total", build_labels) - before_build == 1.0 - assert _sample("gateway_abandoned_attempts_total", upstream_labels) - before_upstream == 1.0 - - # Every family the gateway registers, as (name, type, label names). A metric # may move to the module that increments it, but the scrape is an external # contract: dashboards, recording rules and alerts outside this repository diff --git a/tests/unit/test_run_platform_attempts.py b/tests/unit/test_run_platform_attempts.py index 780de99cef..e4f8cae7d1 100644 --- a/tests/unit/test_run_platform_attempts.py +++ b/tests/unit/test_run_platform_attempts.py @@ -18,7 +18,13 @@ from fastapi import HTTPException from gateway.api.routes import _platform -from gateway.api.routes._platform import ResolvedAttempt, ResolvedRoute, default_attempt_kwargs, run_platform_attempts +from gateway.api.routes._platform import ( + ResolvedAttempt, + ResolvedRoute, + default_attempt_kwargs, + record_abandoned_attempt, + run_platform_attempts, +) from gateway.core.config import GatewayConfig from gateway.metrics import REGISTRY from gateway.services.mcp_loop import MaxToolIterationsExceeded @@ -36,6 +42,27 @@ def _abandoned_sample(provider: str, model: str, reason: str, position: int) -> ) +def test_record_abandoned_attempt_increments_counter() -> None: + before = _abandoned_sample("ab-prov", "ab-model", "timeout", 0) + + record_abandoned_attempt("ab-prov", "ab-model", "timeout", 0) + + assert _abandoned_sample("ab-prov", "ab-model", "timeout", 0) - before == 1.0 + + +def test_record_abandoned_attempt_labels_by_reason_and_position() -> None: + """Each (reason, position) pair is its own series so operators can spot which + plan entry and failure phase dominates the fallback waste.""" + before_build = _abandoned_sample("ab-prov2", "ab-model2", "build_error", 1) + before_upstream = _abandoned_sample("ab-prov2", "ab-model2", "upstream_error", 2) + + record_abandoned_attempt("ab-prov2", "ab-model2", "build_error", 1) + record_abandoned_attempt("ab-prov2", "ab-model2", "upstream_error", 2) + + assert _abandoned_sample("ab-prov2", "ab-model2", "build_error", 1) - before_build == 1.0 + assert _abandoned_sample("ab-prov2", "ab-model2", "upstream_error", 2) - before_upstream == 1.0 + + def _single_attempt(provider: str, model: str) -> ResolvedAttempt: return ResolvedAttempt( attempt_id="a0", position=0, provider=provider, model=model, api_key="k", managed=False From 3fff87c38d8bd56d5d9402cb8ae78a07aeb57ddf Mon Sep 17 00:00:00 2001 From: Peter Wilson Date: Tue, 15 Sep 2026 20:24:59 +0100 Subject: [PATCH 7/8] refactor(metrics): move the token, cost and settlement metrics into the pipeline The shared request pipeline is the only code that records token counts, request cost and inline cost settlement outcomes, so those three metrics and their record helpers now live in _pipeline.py, registered on the shared registry. The Prometheus Counter is imported under an alias there because the module already uses collections.Counter. metrics.py now holds the registry, the process collector and the HTTP request instrumentation only. Part of #1177. --- src/gateway/api/routes/_pipeline.py | 44 ++++++++++++++++++++++- src/gateway/metrics.py | 41 +-------------------- tests/unit/test_gateway_metrics.py | 47 ------------------------ tests/unit/test_pipeline_metrics.py | 55 +++++++++++++++++++++++++++++ 4 files changed, 99 insertions(+), 88 deletions(-) create mode 100644 tests/unit/test_pipeline_metrics.py diff --git a/src/gateway/api/routes/_pipeline.py b/src/gateway/api/routes/_pipeline.py index aac65fd6c8..87f4c8729f 100644 --- a/src/gateway/api/routes/_pipeline.py +++ b/src/gateway/api/routes/_pipeline.py @@ -61,6 +61,8 @@ from any_llm.types.responses import ResponsesParams from fastapi import BackgroundTasks, HTTPException, Request, Response, status from fastapi.responses import StreamingResponse +from prometheus_client import Counter as PrometheusCounter +from prometheus_client import Histogram from sqlalchemy.exc import SQLAlchemyError from sqlalchemy.ext.asyncio import AsyncSession @@ -118,7 +120,7 @@ ) from gateway.inflight import track_request from gateway.log_config import logger -from gateway.metrics import record_cost, record_inline_cost_settlement, record_tokens +from gateway.metrics import REGISTRY from gateway.model_labeling import relabel_model from gateway.models.entities import APIKey, ModelPricing, UsageLog from gateway.models.guardrails import GuardrailConfig @@ -221,6 +223,46 @@ ResultT = TypeVar("ResultT") ChunkT = TypeVar("ChunkT") +TOKENS = PrometheusCounter( + "gateway_tokens", + "Total number of tokens processed", + ["provider", "model", "type"], + registry=REGISTRY, +) + +REQUEST_COST_DOLLARS = Histogram( + "gateway_request_cost_dollars", + "Request cost in USD", + ["provider", "model"], + registry=REGISTRY, +) + +INLINE_COST_SETTLEMENTS = PrometheusCounter( + "gateway_inline_cost_settlements", + "Inline platform cost settlement outcomes on the hybrid response path", + ["outcome"], + registry=REGISTRY, +) + + +def record_tokens(provider: str, model: str, prompt_tokens: int, completion_tokens: int) -> None: + """Record token usage metrics.""" + if prompt_tokens: + TOKENS.labels(provider=provider, model=model, type="input").inc(prompt_tokens) + if completion_tokens: + TOKENS.labels(provider=provider, model=model, type="output").inc(completion_tokens) + + +def record_cost(provider: str, model: str, cost: float) -> None: + """Record request cost.""" + REQUEST_COST_DOLLARS.labels(provider=provider, model=model).observe(cost) + + +def record_inline_cost_settlement(outcome: str) -> None: + """Record an attached, unattached, or timed-out inline settlement.""" + INLINE_COST_SETTLEMENTS.labels(outcome=outcome).inc() + + # --------------------------------------------------------------------------- # Shared wire-level detail strings. These are client-visible API contract # values; do not edit them without a deprecation plan. diff --git a/src/gateway/metrics.py b/src/gateway/metrics.py index 89302e6c0d..1f6fee7b39 100644 --- a/src/gateway/metrics.py +++ b/src/gateway/metrics.py @@ -1,4 +1,4 @@ -"""Prometheus metrics for the gateway.""" +"""Prometheus registry and HTTP request instrumentation for the gateway.""" from __future__ import annotations @@ -49,27 +49,6 @@ registry=REGISTRY, ) -TOKENS = Counter( - "gateway_tokens", - "Total number of tokens processed", - ["provider", "model", "type"], - registry=REGISTRY, -) - -REQUEST_COST_DOLLARS = Histogram( - "gateway_request_cost_dollars", - "Request cost in USD", - ["provider", "model"], - registry=REGISTRY, -) - -INLINE_COST_SETTLEMENTS = Counter( - "gateway_inline_cost_settlements", - "Inline platform cost settlement outcomes on the hybrid response path", - ["outcome"], - registry=REGISTRY, -) - _PROMETHEUS_CONTENT_TYPE = "text/plain; version=0.0.4; charset=utf-8" @@ -149,21 +128,3 @@ async def send_wrapper(message: Message) -> None: REQUEST_DURATION_SECONDS.labels(method=method, endpoint=endpoint, api_version=api_version).observe( duration ) - - -def record_tokens(provider: str, model: str, prompt_tokens: int, completion_tokens: int) -> None: - """Record token usage metrics.""" - if prompt_tokens: - TOKENS.labels(provider=provider, model=model, type="input").inc(prompt_tokens) - if completion_tokens: - TOKENS.labels(provider=provider, model=model, type="output").inc(completion_tokens) - - -def record_cost(provider: str, model: str, cost: float) -> None: - """Record request cost.""" - REQUEST_COST_DOLLARS.labels(provider=provider, model=model).observe(cost) - - -def record_inline_cost_settlement(outcome: str) -> None: - """Record an attached, unattached, or timed-out inline settlement.""" - INLINE_COST_SETTLEMENTS.labels(outcome=outcome).inc() diff --git a/tests/unit/test_gateway_metrics.py b/tests/unit/test_gateway_metrics.py index 1c07788a21..d6529c8f2b 100644 --- a/tests/unit/test_gateway_metrics.py +++ b/tests/unit/test_gateway_metrics.py @@ -13,9 +13,6 @@ MetricsMiddleware, _endpoint_label, metrics_endpoint, - record_cost, - record_inline_cost_settlement, - record_tokens, ) @@ -24,50 +21,6 @@ def _sample(name: str, labels: dict[str, str] | None = None) -> float: return REGISTRY.get_sample_value(name, labels or {}) or 0.0 -def test_record_tokens_increments_counters() -> None: - input_labels = {"provider": "test-prov", "model": "test-model", "type": "input"} - output_labels = {"provider": "test-prov", "model": "test-model", "type": "output"} - before_in = _sample("gateway_tokens_total", input_labels) - before_out = _sample("gateway_tokens_total", output_labels) - - record_tokens("test-prov", "test-model", 100, 50) - - assert _sample("gateway_tokens_total", input_labels) - before_in == 100.0 - assert _sample("gateway_tokens_total", output_labels) - before_out == 50.0 - - -def test_record_tokens_skips_zero_values() -> None: - labels_in = {"provider": "zero-prov", "model": "zero-model", "type": "input"} - labels_out = {"provider": "zero-prov", "model": "zero-model", "type": "output"} - before_in = _sample("gateway_tokens_total", labels_in) - before_out = _sample("gateway_tokens_total", labels_out) - - record_tokens("zero-prov", "zero-model", 0, 0) - - assert _sample("gateway_tokens_total", labels_in) == before_in - assert _sample("gateway_tokens_total", labels_out) == before_out - - -def test_record_cost_observes_histogram() -> None: - labels = {"provider": "cost-prov", "model": "cost-model"} - before_count = _sample("gateway_request_cost_dollars_count", labels) - - record_cost("cost-prov", "cost-model", 1.23) - - assert _sample("gateway_request_cost_dollars_count", labels) - before_count == 1.0 - assert _sample("gateway_request_cost_dollars_sum", labels) >= 1.23 - - -@pytest.mark.parametrize("outcome", ["attached", "unattached", "timeout"]) -def test_record_inline_cost_settlement_increments_counter(outcome: str) -> None: - labels = {"outcome": outcome} - before = _sample("gateway_inline_cost_settlements_total", labels) - - record_inline_cost_settlement(outcome) - - assert _sample("gateway_inline_cost_settlements_total", labels) - before == 1.0 - - # Every family the gateway registers, as (name, type, label names). A metric # may move to the module that increments it, but the scrape is an external # contract: dashboards, recording rules and alerts outside this repository diff --git a/tests/unit/test_pipeline_metrics.py b/tests/unit/test_pipeline_metrics.py new file mode 100644 index 0000000000..bb23acd5d2 --- /dev/null +++ b/tests/unit/test_pipeline_metrics.py @@ -0,0 +1,55 @@ +"""Unit tests for the token, cost and inline settlement metrics the pipeline records.""" + +import pytest + +from gateway.api.routes._pipeline import record_cost, record_inline_cost_settlement, record_tokens +from gateway.metrics import REGISTRY + + +def _sample(name: str, labels: dict[str, str] | None = None) -> float: + """Read a metric sample value from the registry, returning 0.0 if not found.""" + return REGISTRY.get_sample_value(name, labels or {}) or 0.0 + + +def test_record_tokens_increments_counters() -> None: + input_labels = {"provider": "test-prov", "model": "test-model", "type": "input"} + output_labels = {"provider": "test-prov", "model": "test-model", "type": "output"} + before_in = _sample("gateway_tokens_total", input_labels) + before_out = _sample("gateway_tokens_total", output_labels) + + record_tokens("test-prov", "test-model", 100, 50) + + assert _sample("gateway_tokens_total", input_labels) - before_in == 100.0 + assert _sample("gateway_tokens_total", output_labels) - before_out == 50.0 + + +def test_record_tokens_skips_zero_values() -> None: + labels_in = {"provider": "zero-prov", "model": "zero-model", "type": "input"} + labels_out = {"provider": "zero-prov", "model": "zero-model", "type": "output"} + before_in = _sample("gateway_tokens_total", labels_in) + before_out = _sample("gateway_tokens_total", labels_out) + + record_tokens("zero-prov", "zero-model", 0, 0) + + assert _sample("gateway_tokens_total", labels_in) == before_in + assert _sample("gateway_tokens_total", labels_out) == before_out + + +def test_record_cost_observes_histogram() -> None: + labels = {"provider": "cost-prov", "model": "cost-model"} + before_count = _sample("gateway_request_cost_dollars_count", labels) + + record_cost("cost-prov", "cost-model", 1.23) + + assert _sample("gateway_request_cost_dollars_count", labels) - before_count == 1.0 + assert _sample("gateway_request_cost_dollars_sum", labels) >= 1.23 + + +@pytest.mark.parametrize("outcome", ["attached", "unattached", "timeout"]) +def test_record_inline_cost_settlement_increments_counter(outcome: str) -> None: + labels = {"outcome": outcome} + before = _sample("gateway_inline_cost_settlements_total", labels) + + record_inline_cost_settlement(outcome) + + assert _sample("gateway_inline_cost_settlements_total", labels) - before == 1.0 From fe90cb6fe2d6d36aeaf1f071d1d61a6c29ef1160 Mon Sep 17 00:00:00 2001 From: Peter Wilson Date: Wed, 16 Sep 2026 14:55:21 +0100 Subject: [PATCH 8/8] refactor(metrics): declare metrics through the types gateway.metrics re-exports The six modules that own a metric imported Counter, Gauge and Histogram from prometheus_client directly. Importing them from gateway.metrics instead keeps one module as the gateway's dependency on prometheus_client, so making that dependency optional later changes that module alone. gateway.metrics lists the types in __all__, because mypy strict does not treat an imported name as exported. The PrometheusCounter alias in _pipeline.py stays: that module also imports collections.Counter. --- src/gateway/api/deps.py | 3 +-- src/gateway/api/routes/_pipeline.py | 5 ++--- src/gateway/api/routes/_platform.py | 3 +-- src/gateway/metrics.py | 7 ++++++- src/gateway/rate_limit.py | 3 +-- src/gateway/services/budget_service.py | 3 +-- src/gateway/services/log_writer.py | 4 +--- 7 files changed, 13 insertions(+), 15 deletions(-) diff --git a/src/gateway/api/deps.py b/src/gateway/api/deps.py index 8ee232652c..8b3e56faf5 100644 --- a/src/gateway/api/deps.py +++ b/src/gateway/api/deps.py @@ -6,7 +6,6 @@ from typing import Annotated from fastapi import Depends, HTTPException, Request, status -from prometheus_client import Counter from sqlalchemy import select, update from sqlalchemy.ext.asyncio import AsyncSession @@ -16,7 +15,7 @@ from gateway.core.database import DATABASE_ERRORS, create_session, get_db from gateway.core.feature import CoreFeature from gateway.log_config import logger -from gateway.metrics import REGISTRY +from gateway.metrics import REGISTRY, Counter from gateway.models.entities import APIKey from gateway.models.tenancy import User as TenancyUser from gateway.ports.billing_port import BillingPort diff --git a/src/gateway/api/routes/_pipeline.py b/src/gateway/api/routes/_pipeline.py index 87f4c8729f..44761c29d8 100644 --- a/src/gateway/api/routes/_pipeline.py +++ b/src/gateway/api/routes/_pipeline.py @@ -61,8 +61,6 @@ from any_llm.types.responses import ResponsesParams from fastapi import BackgroundTasks, HTTPException, Request, Response, status from fastapi.responses import StreamingResponse -from prometheus_client import Counter as PrometheusCounter -from prometheus_client import Histogram from sqlalchemy.exc import SQLAlchemyError from sqlalchemy.ext.asyncio import AsyncSession @@ -120,7 +118,8 @@ ) from gateway.inflight import track_request from gateway.log_config import logger -from gateway.metrics import REGISTRY +from gateway.metrics import REGISTRY, Histogram +from gateway.metrics import Counter as PrometheusCounter from gateway.model_labeling import relabel_model from gateway.models.entities import APIKey, ModelPricing, UsageLog from gateway.models.guardrails import GuardrailConfig diff --git a/src/gateway/api/routes/_platform.py b/src/gateway/api/routes/_platform.py index 970f90922b..c6fad67d32 100644 --- a/src/gateway/api/routes/_platform.py +++ b/src/gateway/api/routes/_platform.py @@ -25,7 +25,6 @@ from fastapi import HTTPException, Request, status from openai import APIConnectionError as _OpenAIAPIConnectionError from openai import APITimeoutError as _OpenAIAPITimeoutError -from prometheus_client import Counter from pydantic import BaseModel, Field, ValidationError from gateway.core.config import GatewayConfig @@ -35,7 +34,7 @@ cache_write_tokens_of, ) from gateway.log_config import logger -from gateway.metrics import REGISTRY +from gateway.metrics import REGISTRY, Counter from gateway.models.mcp import McpServerConfig, ResolvedMcpServer from gateway.services.bedrock_gateway_auth import build_bedrock_client_args from gateway.services.mcp_loop import MaxToolIterationsExceeded diff --git a/src/gateway/metrics.py b/src/gateway/metrics.py index 1f6fee7b39..882cda1882 100644 --- a/src/gateway/metrics.py +++ b/src/gateway/metrics.py @@ -1,4 +1,7 @@ -"""Prometheus registry and HTTP request instrumentation for the gateway.""" +"""Prometheus registry, metric types, and HTTP request instrumentation for the gateway. + +The metric types are re-exported so that code declaring a metric need not depend on ``prometheus_client`` directly. +""" from __future__ import annotations @@ -21,6 +24,8 @@ from starlette.requests import Request from starlette.types import ASGIApp, Message, Receive, Scope, Send +__all__ = ["Counter", "Gauge", "Histogram", "MetricsMiddleware", "REGISTRY", "metrics_endpoint"] + REGISTRY = CollectorRegistry() # process_resident_memory_bytes and friends. The gateway keeps its own registry diff --git a/src/gateway/rate_limit.py b/src/gateway/rate_limit.py index 1cc3d3749e..0eb55235dc 100644 --- a/src/gateway/rate_limit.py +++ b/src/gateway/rate_limit.py @@ -6,9 +6,8 @@ from dataclasses import dataclass from fastapi import HTTPException, Request, status -from prometheus_client import Counter -from gateway.metrics import REGISTRY +from gateway.metrics import REGISTRY, Counter RATE_LIMIT_HITS = Counter( "gateway_rate_limit_hits", diff --git a/src/gateway/services/budget_service.py b/src/gateway/services/budget_service.py index c7867f9403..d3deb70ef2 100644 --- a/src/gateway/services/budget_service.py +++ b/src/gateway/services/budget_service.py @@ -10,14 +10,13 @@ from any_llm import AnyLLM from any_llm.exceptions import AnyLLMError from fastapi import HTTPException, status -from prometheus_client import Counter from sqlalchemy import select, update from sqlalchemy.exc import SQLAlchemyError from sqlalchemy.ext.asyncio import AsyncSession from gateway.core.metered_pricing import estimate_metered_cost from gateway.log_config import logger -from gateway.metrics import REGISTRY +from gateway.metrics import REGISTRY, Counter from gateway.models.entities import MAX_COUNT_LIMIT, Budget, BudgetResetLog, ModelPricing, User from gateway.models.money import to_usd from gateway.repositories.users_repository import get_active_user diff --git a/src/gateway/services/log_writer.py b/src/gateway/services/log_writer.py index e19144cba4..baf0db1b42 100644 --- a/src/gateway/services/log_writer.py +++ b/src/gateway/services/log_writer.py @@ -6,11 +6,9 @@ import time from typing import Protocol -from prometheus_client import Counter, Gauge, Histogram - from gateway.core.database import DATABASE_ERRORS, create_log_session from gateway.log_config import logger -from gateway.metrics import REGISTRY +from gateway.metrics import REGISTRY, Counter, Gauge, Histogram from gateway.models.entities import UsageLog QUEUE_DEPTH = Gauge(