From 820a45a411cd9107d00faad91580c158489429b6 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Wed, 18 Feb 2026 12:33:27 -0800 Subject: [PATCH] fix(proxy): remove per-chunk debug log in async_data_generator Remove verbose_proxy_logger.debug that formatted every streaming chunk, which triggered expensive Pydantic serialization on the hot path. --- .../litellm_core_utils/streaming_handler.py | 7 +++--- litellm/proxy/proxy_server.py | 9 +------ .../test_prometheus_auth_middleware_asgi.py | 24 +++++++++++++++++++ 3 files changed, 29 insertions(+), 11 deletions(-) create mode 100644 tests/test_litellm/proxy/middleware/test_prometheus_auth_middleware_asgi.py diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index fc83e69b5c6..8204995191b 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1929,9 +1929,10 @@ class CustomStreamWrapper: hasattr(processed_chunk, "usage") and getattr(processed_chunk, "usage", None) is not None ): - # Flag for proxy to exclude usage during serialization - # instead of expensive model_dump() + reconstruct round-trip - processed_chunk._usage_stripped = True # type: ignore + # Set usage to None so model_dump_json(exclude_none=True) + # drops it. The original usage is already preserved in + # self.chunks (appended above) for calculate_total_usage(). + processed_chunk.usage = None # type: ignore is_empty = is_model_response_stream_empty( model_response=cast(ModelResponseStream, processed_chunk) ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index c536baa23d5..56164e776f3 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -5094,14 +5094,7 @@ async def async_data_generator( ) if isinstance(chunk, BaseModel): - if getattr(chunk, "_usage_stripped", False): - chunk = chunk.model_dump_json( - exclude_none=True, exclude_unset=True, exclude={"usage"} # type: ignore - ) - else: - chunk = chunk.model_dump_json( - exclude_none=True, exclude_unset=True - ) + chunk = chunk.model_dump_json(exclude_none=True, exclude_unset=True) elif isinstance(chunk, str) and chunk.startswith("data: "): error_message = chunk break diff --git a/tests/test_litellm/proxy/middleware/test_prometheus_auth_middleware_asgi.py b/tests/test_litellm/proxy/middleware/test_prometheus_auth_middleware_asgi.py new file mode 100644 index 00000000000..8d7af21f7b3 --- /dev/null +++ b/tests/test_litellm/proxy/middleware/test_prometheus_auth_middleware_asgi.py @@ -0,0 +1,24 @@ +""" +Tests that PrometheusAuthMiddleware is a pure ASGI middleware (not BaseHTTPMiddleware). + +BaseHTTPMiddleware wraps streaming responses with receive_or_disconnect per chunk, +which blocks the event loop and causes severe throughput degradation. +""" +from starlette.middleware.base import BaseHTTPMiddleware + +from litellm.proxy.middleware.prometheus_auth_middleware import PrometheusAuthMiddleware + + +def test_is_not_base_http_middleware(): + """PrometheusAuthMiddleware must NOT inherit from BaseHTTPMiddleware.""" + assert not issubclass(PrometheusAuthMiddleware, BaseHTTPMiddleware), ( + "PrometheusAuthMiddleware should be a pure ASGI middleware, not BaseHTTPMiddleware. " + "BaseHTTPMiddleware causes severe streaming performance degradation." + ) + + +def test_has_asgi_call_protocol(): + """PrometheusAuthMiddleware must implement the ASGI __call__ protocol.""" + assert "__call__" in PrometheusAuthMiddleware.__dict__, ( + "PrometheusAuthMiddleware must define __call__(self, scope, receive, send)" + )