From 7020b0a64b17aacc4a8971e3822be94a2dc3610d Mon Sep 17 00:00:00 2001 From: shivam Date: Fri, 17 Apr 2026 13:32:05 -0700 Subject: [PATCH 1/4] fix(proxy): replay ASGI receive after metrics auth to avoid /metrics hang --- .../middleware/prometheus_auth_middleware.py | 27 ++++++++++++++-- .../test_prometheus_auth_middleware.py | 31 +++++++++++++++++++ 2 files changed, 56 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/middleware/prometheus_auth_middleware.py b/litellm/proxy/middleware/prometheus_auth_middleware.py index 5915e4aa07..d0e38f9894 100644 --- a/litellm/proxy/middleware/prometheus_auth_middleware.py +++ b/litellm/proxy/middleware/prometheus_auth_middleware.py @@ -2,6 +2,7 @@ Prometheus Auth Middleware - Pure ASGI implementation """ import json +from typing import List from fastapi import Request from starlette.types import ASGIApp, Receive, Scope, Send @@ -39,8 +40,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[dict] = [] + + async def receive_for_auth() -> dict: + message = await receive() + buffered_messages.append(message) + return message + + request = Request(scope, receive_for_auth) api_key = request.headers.get(_AUTHORIZATION_HEADER) or "" try: @@ -69,5 +79,18 @@ class PrometheusAuthMiddleware: ) return + replay_idx = 0 + + async def receive_replay() -> dict: + 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) diff --git a/tests/test_litellm/proxy/middleware/test_prometheus_auth_middleware.py b/tests/test_litellm/proxy/middleware/test_prometheus_auth_middleware.py index 9fd244d9c3..310ee11573 100644 --- a/tests/test_litellm/proxy/middleware/test_prometheus_auth_middleware.py +++ b/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. From 733ccd6e0920252a99fdbe8a5bad892503026499 Mon Sep 17 00:00:00 2001 From: shivam Date: Fri, 17 Apr 2026 16:36:34 -0700 Subject: [PATCH 2/4] fixed linting --- litellm/proxy/middleware/prometheus_auth_middleware.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/middleware/prometheus_auth_middleware.py b/litellm/proxy/middleware/prometheus_auth_middleware.py index d0e38f9894..d5dec6c4bd 100644 --- a/litellm/proxy/middleware/prometheus_auth_middleware.py +++ b/litellm/proxy/middleware/prometheus_auth_middleware.py @@ -2,7 +2,7 @@ Prometheus Auth Middleware - Pure ASGI implementation """ import json -from typing import List +from typing import Any, List, MutableMapping from fastapi import Request from starlette.types import ASGIApp, Receive, Scope, Send @@ -43,9 +43,9 @@ class PrometheusAuthMiddleware: # 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[dict] = [] + buffered_messages: List[MutableMapping[str, Any]] = [] - async def receive_for_auth() -> dict: + async def receive_for_auth() -> MutableMapping[str, Any]: message = await receive() buffered_messages.append(message) return message @@ -81,7 +81,7 @@ class PrometheusAuthMiddleware: replay_idx = 0 - async def receive_replay() -> dict: + async def receive_replay() -> MutableMapping[str, Any]: nonlocal replay_idx if replay_idx < len(buffered_messages): msg = buffered_messages[replay_idx] From 49ddb0e5bad5d2f3dbb2f5bb4014ece36c295eef Mon Sep 17 00:00:00 2001 From: shivam Date: Fri, 17 Apr 2026 16:43:01 -0700 Subject: [PATCH 3/4] style: apply Black formatting to Prometheus integration modules Reformat prometheus logger, types, and metrics auth middleware so lint CI (black --check) passes. Made-with: Cursor --- litellm/integrations/prometheus.py | 46 +++++++++++-------- .../middleware/prometheus_auth_middleware.py | 1 + litellm/types/integrations/prometheus.py | 12 ++--- 3 files changed, 33 insertions(+), 26 deletions(-) diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index b3bf792e93..d5fde2a486 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -88,7 +88,9 @@ class PrometheusLogger(CustomLogger): _custom_buckets = litellm.prometheus_latency_buckets self.latency_buckets = ( - tuple(_custom_buckets) if _custom_buckets is not None else LATENCY_BUCKETS + tuple(_custom_buckets) + if _custom_buckets is not None + else LATENCY_BUCKETS ) # Create metric factory functions @@ -1097,9 +1099,11 @@ class PrometheusLogger(CustomLogger): ), client_ip=standard_logging_payload["metadata"].get("requester_ip_address"), user_agent=standard_logging_payload["metadata"].get("user_agent"), - stream=str(standard_logging_payload.get("stream")) - if litellm.prometheus_emit_stream_label - else None, + stream=( + str(standard_logging_payload.get("stream")) + if litellm.prometheus_emit_stream_label + else None + ), ) if ( @@ -1767,9 +1771,11 @@ class PrometheusLogger(CustomLogger): client_ip=_metadata.get("requester_ip_address"), user_agent=_metadata.get("user_agent"), model_id=model_id, - stream=str(request_data.get("stream")) - if litellm.prometheus_emit_stream_label - else None, + stream=( + str(request_data.get("stream")) + if litellm.prometheus_emit_stream_label + else None + ), ) _labels = prometheus_label_factory( supported_enum_labels=self.get_labels_for_metric( @@ -2093,9 +2099,9 @@ class PrometheusLogger(CustomLogger): ): try: verbose_logger.debug("setting remaining tokens requests metric") - standard_logging_payload: Optional[ - StandardLoggingPayload - ] = request_kwargs.get("standard_logging_object") + standard_logging_payload: Optional[StandardLoggingPayload] = ( + request_kwargs.get("standard_logging_object") + ) if standard_logging_payload is None: return @@ -2728,9 +2734,7 @@ class PrometheusLogger(CustomLogger): ) return - async def fetch_keys( - page_size: int, page: int - ) -> Tuple[ + async def fetch_keys(page_size: int, page: int) -> Tuple[ List[Union[str, UserAPIKeyAuth, LiteLLM_DeletedVerificationToken]], Optional[int], ]: @@ -2921,9 +2925,11 @@ class PrometheusLogger(CustomLogger): org_alias=org.organization_alias or "", spend=org.spend or 0.0, max_budget=budget_table.max_budget if budget_table else None, - budget_reset_at=getattr(budget_table, "budget_reset_at", None) - if budget_table - else None, + budget_reset_at=( + getattr(budget_table, "budget_reset_at", None) + if budget_table + else None + ), ) async def _set_team_budget_metrics_after_api_request( @@ -3405,10 +3411,10 @@ class PrometheusLogger(CustomLogger): from litellm.constants import PROMETHEUS_BUDGET_METRICS_REFRESH_INTERVAL_MINUTES from litellm.integrations.custom_logger import CustomLogger - prometheus_loggers: List[ - CustomLogger - ] = litellm.logging_callback_manager.get_custom_loggers_for_type( - callback_type=PrometheusLogger + prometheus_loggers: List[CustomLogger] = ( + litellm.logging_callback_manager.get_custom_loggers_for_type( + callback_type=PrometheusLogger + ) ) # we need to get the initialized prometheus logger instance(s) and call logger.initialize_remaining_budget_metrics() on them verbose_logger.debug("found %s prometheus loggers", len(prometheus_loggers)) diff --git a/litellm/proxy/middleware/prometheus_auth_middleware.py b/litellm/proxy/middleware/prometheus_auth_middleware.py index d5dec6c4bd..3b30fd3d63 100644 --- a/litellm/proxy/middleware/prometheus_auth_middleware.py +++ b/litellm/proxy/middleware/prometheus_auth_middleware.py @@ -1,6 +1,7 @@ """ Prometheus Auth Middleware - Pure ASGI implementation """ + import json from typing import Any, List, MutableMapping diff --git a/litellm/types/integrations/prometheus.py b/litellm/types/integrations/prometheus.py index 51a41f97e0..1b36ad5daa 100644 --- a/litellm/types/integrations/prometheus.py +++ b/litellm/types/integrations/prometheus.py @@ -676,9 +676,9 @@ class PrometheusMetricLabels: litellm_managed_batch_created_total = _batch_user_labels - litellm_managed_file_size_bytes: List[ - str - ] = [] # labels: purpose, file_type, model, api_provider, user (custom) + litellm_managed_file_size_bytes: List[str] = ( + [] + ) # labels: purpose, file_type, model, api_provider, user (custom) litellm_managed_batch_duration_seconds = [ UserAPIKeyLabelNames.v1_LITELLM_MODEL_NAME.value, @@ -687,9 +687,9 @@ class PrometheusMetricLabels: litellm_managed_file_created_total = _batch_user_labels - litellm_managed_file_deleted_total: List[ - str - ] = [] # only "result" label, added at metric creation + litellm_managed_file_deleted_total: List[str] = ( + [] + ) # only "result" label, added at metric creation litellm_check_batch_cost_jobs_polled: List[str] = [] From 1170bd55c97612fe996de924b6ad95c87bb229ab Mon Sep 17 00:00:00 2001 From: shivam Date: Fri, 17 Apr 2026 16:58:25 -0700 Subject: [PATCH 4/4] fix(prometheus): remove dead prometheus_label_factory call; apply Black async_post_call_failure_hook passed only supported_enum_labels to prometheus_label_factory; enum_values is required. The call was unused because _inc_labeled_counter builds labels internally. Reformat Prometheus-related modules and passthrough/copilot helpers for CI. Made-with: Cursor --- litellm/integrations/prometheus.py | 18 ++++++++---------- litellm/integrations/prometheus_helpers.py | 3 +-- litellm/llms/github_copilot/authenticator.py | 4 +--- litellm/passthrough/utils.py | 4 +++- litellm/types/integrations/prometheus.py | 2 +- 5 files changed, 14 insertions(+), 17 deletions(-) diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index f4fa1e6c03..723b142dfa 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -51,6 +51,7 @@ if TYPE_CHECKING: else: AsyncIOScheduler = Any + class PrometheusLogger(CustomLogger): # Class variables or attributes @@ -991,9 +992,7 @@ class PrometheusLogger(CustomLogger): amount: float = 1.0, ) -> None: _labels = prometheus_label_factory( - supported_enum_labels=self.get_labels_for_metric( - metric_name=metric_name - ), + supported_enum_labels=self.get_labels_for_metric(metric_name=metric_name), enum_values=enum_values, label_context=label_context, ) @@ -1118,7 +1117,9 @@ class PrometheusLogger(CustomLogger): user_api_key = hash_token(user_api_key) - label_context = PrometheusLabelFactoryContext(enum_values) #amortized per request. + label_context = PrometheusLabelFactoryContext( + enum_values + ) # amortized per request. # increment total LLM requests and spend metric self._increment_top_level_request_and_spend_metrics( @@ -1791,11 +1792,6 @@ class PrometheusLogger(CustomLogger): else None ), ) - _labels = prometheus_label_factory( - supported_enum_labels=self.get_labels_for_metric( - metric_name="litellm_proxy_failed_requests_metric" - ), - ) _label_ctx = PrometheusLabelFactoryContext(enum_values) PrometheusLogger._inc_labeled_counter( self, @@ -3495,7 +3491,9 @@ def _prometheus_labels_from_context( } if UserAPIKeyLabelNames.END_USER.value in filtered_labels: - filtered_labels[UserAPIKeyLabelNames.END_USER.value] = ctx.get_resolved_end_user() + filtered_labels[UserAPIKeyLabelNames.END_USER.value] = ( + ctx.get_resolved_end_user() + ) for sk, val in ctx._custom_by_sanitized_key.items(): if sk in supported_enum_labels: diff --git a/litellm/integrations/prometheus_helpers.py b/litellm/integrations/prometheus_helpers.py index 34f4855863..784ab524dd 100644 --- a/litellm/integrations/prometheus_helpers.py +++ b/litellm/integrations/prometheus_helpers.py @@ -51,8 +51,7 @@ class PrometheusLabelFactoryContext: self.enum_values = enum_values enum_dict = enum_values.model_dump() self._sanitized_enum: Dict[str, Optional[str]] = { - k: _sanitize_prometheus_label_value(v) - for k, v in enum_dict.items() + k: _sanitize_prometheus_label_value(v) for k, v in enum_dict.items() } self._custom_by_sanitized_key: Dict[str, Optional[str]] = {} if enum_values.custom_metadata_labels is not None: diff --git a/litellm/llms/github_copilot/authenticator.py b/litellm/llms/github_copilot/authenticator.py index f4698861ed..9de2987b9f 100644 --- a/litellm/llms/github_copilot/authenticator.py +++ b/litellm/llms/github_copilot/authenticator.py @@ -294,9 +294,7 @@ class Authenticator: access_token_url = os.getenv( "GITHUB_COPILOT_ACCESS_TOKEN_URL", DEFAULT_GITHUB_ACCESS_TOKEN_URL ) - client_id = os.getenv( - "GITHUB_COPILOT_CLIENT_ID", DEFAULT_GITHUB_CLIENT_ID - ) + client_id = os.getenv("GITHUB_COPILOT_CLIENT_ID", DEFAULT_GITHUB_CLIENT_ID) for attempt in range(max_attempts): try: diff --git a/litellm/passthrough/utils.py b/litellm/passthrough/utils.py index 5dde13f007..d39a0dda15 100644 --- a/litellm/passthrough/utils.py +++ b/litellm/passthrough/utils.py @@ -79,7 +79,9 @@ class BasePassthroughUtils: for header_name, header_value in request_headers.items(): if header_name.lower().startswith(PASS_THROUGH_HEADER_PREFIX): # Strip the 'x-pass-' prefix and normalize to lowercase - actual_header_name = header_name[len(PASS_THROUGH_HEADER_PREFIX) :].lower() + actual_header_name = header_name[ + len(PASS_THROUGH_HEADER_PREFIX) : + ].lower() if actual_header_name in _PASS_THROUGH_PROTECTED_HEADERS or any( actual_header_name.startswith(p) for p in _PASS_THROUGH_PROTECTED_HEADER_PREFIXES diff --git a/litellm/types/integrations/prometheus.py b/litellm/types/integrations/prometheus.py index 338c5a79ce..43a287f29b 100644 --- a/litellm/types/integrations/prometheus.py +++ b/litellm/types/integrations/prometheus.py @@ -784,7 +784,7 @@ class UserAPIKeyLabelValues: org_id: Optional[str] = None org_alias: Optional[str] = None - #Added for test compatibility. + # Added for test compatibility. def __init__(self, **kwargs: Any) -> None: """ Match former Pydantic behavior: unknown keys are ignored; ``api_key_hash`` maps to