Merge pull request #20434 from ryan-crabbe/perf/prometheus-asgi-middleware

Perf/prometheus asgi middleware

litellm/proxy/middleware/prometheus_auth_middleware.py

@@ -1,25 +1,24 @@
 """
-Prometheus Auth Middleware
+Prometheus Auth Middleware - Pure ASGI implementation
+"""
+import json
 
-Pure ASGI middleware — avoids Starlette's BaseHTTPMiddleware which wraps
-streaming responses with receive_or_disconnect per chunk, blocking the
-event loop and causing severe throughput degradation under concurrent
-streaming load.
-"""
-from starlette.requests import Request
-from starlette.responses import JSONResponse
+from fastapi import Request
 from starlette.types import ASGIApp, Receive, Scope, Send
 
 import litellm
 from litellm.proxy._types import SpecialHeaders
 from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
 
+# Cache the header name at module level to avoid repeated enum attribute access
+_AUTHORIZATION_HEADER = SpecialHeaders.openai_authorization.value  # "Authorization"
+
 
 class PrometheusAuthMiddleware:
     """
-    Middleware to authenticate requests to the metrics endpoint
+    Middleware to authenticate requests to the metrics endpoint.
 
-    By default, auth is not run on the metrics endpoint
+    By default, auth is not run on the metrics endpoint.
 
     Enabled by setting the following in proxy_config.yaml:
 
@@ -33,51 +32,42 @@ class PrometheusAuthMiddleware:
         self.app = app
 
     async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
-        if scope["type"] not in ("http", "websocket"):
+        # Fast path: only inspect HTTP requests; pass through websocket/lifespan immediately
+        if scope["type"] != "http" or "/metrics" not in scope.get("path", ""):
             await self.app(scope, receive, send)
             return
 
-        request = Request(scope, receive)
+        # Only run auth if configured to do so
+        if litellm.require_auth_for_metrics_endpoint is True:
+            # Construct Request only when auth is actually needed
+            request = Request(scope, receive)
+            api_key = request.headers.get(_AUTHORIZATION_HEADER) or ""
 
-        if self._is_prometheus_metrics_endpoint(request):
-            if self._should_run_auth_on_metrics_endpoint() is True:
-                try:
-                    await user_api_key_auth(
-                        request=request,
-                        api_key=request.headers.get(
-                            SpecialHeaders.openai_authorization.value
-                        )
-                        or "",
-                    )
-                except Exception as e:
-                    response = JSONResponse(
-                        status_code=401,
-                        content=f"Unauthorized access to metrics endpoint: {getattr(e, 'message', str(e))}",
-                    )
-                    await response(scope, receive, send)
-                    return
+            try:
+                await user_api_key_auth(request=request, api_key=api_key)
+            except Exception as e:
+                # Send 401 response directly via ASGI protocol
+                error_message = getattr(e, "message", str(e))
+                body = json.dumps(
+                    f"Unauthorized access to metrics endpoint: {error_message}"
+                ).encode("utf-8")
+                await send(
+                    {
+                        "type": "http.response.start",
+                        "status": 401,
+                        "headers": [
+                            [b"content-type", b"application/json"],
+                            [b"content-length", str(len(body)).encode("ascii")],
+                        ],
+                    }
+                )
+                await send(
+                    {
+                        "type": "http.response.body",
+                        "body": body,
+                    }
+                )
+                return
 
+        # Pass through to the inner application
         await self.app(scope, receive, send)
-
-    @staticmethod
-    def _is_prometheus_metrics_endpoint(request: Request):
-        try:
-            if "/metrics" in request.url.path:
-                return True
-            return False
-        except Exception:
-            return False
-
-    @staticmethod
-    def _should_run_auth_on_metrics_endpoint():
-        """
-        Returns True if auth should be run on the metrics endpoint
-
-        False by default, set to True in proxy_config.yaml to enable
-
-        ```yaml
-        litellm_settings:
-            require_auth_for_metrics_endpoint: true
-        ```
-        """
-        return litellm.require_auth_for_metrics_endpoint

tests/test_litellm/proxy/middleware/test_prometheus_auth_middleware.py

@@ -127,3 +127,46 @@ def test_no_auth_metrics_when_disabled(app_with_middleware, monkeypatch):
     response = client.get("/metrics")
     assert response.status_code == 200, response.text
     assert response.json() == {"msg": "metrics OK"}
+
+
+def test_non_metrics_requests_pass_through(app_with_middleware):
+    """
+    Test that non-metrics endpoints pass through the middleware unaffected.
+    """
+    litellm.require_auth_for_metrics_endpoint = True
+
+    client = TestClient(app_with_middleware)
+
+    response = client.get("/chat/completions")
+    assert response.status_code == 200, response.text
+    assert response.json() == {"msg": "chat completions OK"}
+
+    response = client.get("/embeddings")
+    assert response.status_code == 200, response.text
+    assert response.json() == {"msg": "embeddings OK"}
+
+
+def test_non_metrics_requests_dont_trigger_auth(app_with_middleware, monkeypatch):
+    """
+    Test that non-metrics requests never trigger auth, even when auth is enabled
+    and the auth function would reject the request.
+    """
+    litellm.require_auth_for_metrics_endpoint = True
+
+    def should_not_be_called(*args, **kwargs):
+        raise Exception("Auth should not be called for non-metrics requests")
+
+    monkeypatch.setattr(
+        "litellm.proxy.middleware.prometheus_auth_middleware.user_api_key_auth",
+        should_not_be_called,
+    )
+
+    client = TestClient(app_with_middleware)
+
+    response = client.get("/chat/completions")
+    assert response.status_code == 200, response.text
+    assert response.json() == {"msg": "chat completions OK"}
+
+    response = client.get("/embeddings")
+    assert response.status_code == 200, response.text
+    assert response.json() == {"msg": "embeddings OK"}