Merge pull request #25980 from BerriAI/litellm_metrics_auth

Fix /metrics hang when require_auth_for_metrics_endpoint is true and auth succeeds

litellm/proxy/middleware/prometheus_auth_middleware.py

@@ -3,6 +3,7 @@ Prometheus Auth Middleware - Pure ASGI implementation
 """
 
 import json
+from typing import Any, List, MutableMapping
 
 from fastapi import Request
 from starlette.types import ASGIApp, Receive, Scope, Send
@@ -40,8 +41,17 @@ class PrometheusAuthMiddleware:
 
         # 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)
+            # user_api_key_auth reads the request body, which consumes ASGI `receive`.
+            # Buffer those messages and replay them for the inner app; otherwise a
+            # successful auth would forward an exhausted receive and /metrics hangs.
+            buffered_messages: List[MutableMapping[str, Any]] = []
+
+            async def receive_for_auth() -> MutableMapping[str, Any]:
+                message = await receive()
+                buffered_messages.append(message)
+                return message
+
+            request = Request(scope, receive_for_auth)
             api_key = request.headers.get(_AUTHORIZATION_HEADER) or ""
 
             try:
@@ -70,5 +80,18 @@ class PrometheusAuthMiddleware:
                 )
                 return
 
+            replay_idx = 0
+
+            async def receive_replay() -> MutableMapping[str, Any]:
+                nonlocal replay_idx
+                if replay_idx < len(buffered_messages):
+                    msg = buffered_messages[replay_idx]
+                    replay_idx += 1
+                    return msg
+                return await receive()
+
+            await self.app(scope, receive_replay, send)
+            return
+
         # Pass through to the inner application
         await self.app(scope, receive, send)

tests/test_litellm/proxy/middleware/test_prometheus_auth_middleware.py

@@ -26,6 +26,15 @@ async def fake_valid_auth(request, api_key):
     return
 
 
+async def fake_valid_auth_reads_body(request, api_key, **kwargs):
+    """
+    Like real user_api_key_auth, consumes the ASGI body stream. Regression test
+    for successful auth passing a drained receive to the inner app (hang).
+    """
+    await request.body()
+    return
+
+
 async def fake_invalid_auth(request, api_key):
     print("running fake invalid auth", request, api_key)
     # Simulate invalid auth by raising an exception.
@@ -62,6 +71,28 @@ def app_with_middleware():
     return app
 
 
+def test_valid_auth_metrics_after_body_consumed(app_with_middleware, monkeypatch):
+    """
+    Auth that reads the request body must not cause /metrics to hang on success.
+    """
+    litellm.require_auth_for_metrics_endpoint = True
+    monkeypatch.setattr(
+        "litellm.proxy.middleware.prometheus_auth_middleware.user_api_key_auth",
+        fake_valid_auth_reads_body,
+    )
+
+    client = TestClient(app_with_middleware)
+    headers = {SpecialHeaders.openai_authorization.value: "valid"}
+
+    response = client.get("/metrics", headers=headers)
+    assert response.status_code == 200, response.text
+    assert response.json() == {"msg": "metrics OK"}
+
+    response = client.get("/metrics/", headers=headers)
+    assert response.status_code == 200, response.text
+    assert response.json() == {"msg": "metrics OK"}
+
+
 def test_valid_auth_metrics(app_with_middleware, monkeypatch):
     """
     Test that a request to /metrics (and /metrics/) with valid auth headers passes.