feat: Auth prometheus metrics (#4725)

* Add prometheus metrics to jwt store and rate limiter

* Delay prometheus imports

* Use cache to count claims

src/phoenix/server/jwt_store.py

@@ -16,6 +16,7 @@
     ClaimSet,
     Token,
 )
+from phoenix.config import get_env_enable_prometheus
 from phoenix.db import models
 from phoenix.db.enums import UserRole
 from phoenix.server.types import (
@@ -446,6 +447,13 @@ def _to_db(self, claims: RefreshTokenClaims) -> models.RefreshToken:
             expires_at=claims.expiration_time,
         )
 
+    async def _update(self) -> None:
+        await super()._update()
+        if get_env_enable_prometheus():
+            from phoenix.server.prometheus import JWT_STORE_TOKENS_ACTIVE
+
+            JWT_STORE_TOKENS_ACTIVE.set(len(self._claims._cache))
+
 
 class _ApiKeyStore(
     _Store[
@@ -489,3 +497,10 @@ def _to_db(self, claims: ApiKeyClaims) -> models.ApiKey:
             created_at=claims.issued_at,
             expires_at=claims.expiration_time or None,
         )
+
+    async def _update(self) -> None:
+        await super()._update()
+        if get_env_enable_prometheus():
+            from phoenix.server.prometheus import JWT_STORE_API_KEYS_ACTIVE
+
+            JWT_STORE_API_KEYS_ACTIVE.set(len(self._claims._cache))

src/phoenix/server/prometheus.py

@@ -50,6 +50,26 @@
     documentation="Total count of bulk loader exceptions",
 )
 
+RATE_LIMITER_CACHE_SIZE = Gauge(
+    name="rate_limiter_cache_size",
+    documentation="Current size of the rate limiter cache",
+)
+
+RATE_LIMITER_THROTTLES = Counter(
+    name="rate_limiter_throttles_total",
+    documentation="Total count of rate limiter throttles",
+)
+
+JWT_STORE_TOKENS_ACTIVE = Gauge(
+    name="jwt_store_tokens_active",
+    documentation="Current number of refresh tokens in the JWT store",
+)
+
+JWT_STORE_API_KEYS_ACTIVE = Gauge(
+    name="jwt_store_api_keys_active",
+    documentation="Current number of API keys in the JWT store",
+)
+
 
 class PrometheusMiddleware(BaseHTTPMiddleware):
     async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:

src/phoenix/server/rate_limiters.py

@@ -15,6 +15,7 @@
 
 from fastapi import HTTPException, Request
 
+from phoenix.config import get_env_enable_prometheus
 from phoenix.exceptions import PhoenixException
 
 
@@ -55,6 +56,10 @@ def available_tokens(self) -> float:
 
     def make_request_if_ready(self) -> None:
         if self.available_tokens() < 1:
+            if get_env_enable_prometheus():
+                from phoenix.server.prometheus import RATE_LIMITER_THROTTLES
+
+                RATE_LIMITER_THROTTLES.inc()
             raise UnavailableTokensError
         self.tokens -= 1
 
@@ -144,6 +149,10 @@ def make_request(self, key: str) -> None:
         self._cleanup_expired_limiters(request_time)
         rate_limiter = self._fetch_token_bucket(key, request_time)
         rate_limiter.make_request_if_ready()
+        if get_env_enable_prometheus():
+            from phoenix.server.prometheus import RATE_LIMITER_CACHE_SIZE
+
+            RATE_LIMITER_CACHE_SIZE.set(sum(len(partition) for partition in self.cache_partitions))
 
 
 def fastapi_ip_rate_limiter(