diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index cdc480f610e..a0f68605229 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1955,15 +1955,10 @@ class CustomStreamWrapper: hasattr(processed_chunk, "usage") and getattr(processed_chunk, "usage", None) is not None ): - # Strip usage from the outgoing chunk so - # model_dump_json(exclude_none=True) drops it. - # The copy in self.chunks retains usage for - # calculate_total_usage(). + # 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 - # After nullifying usage, check if the chunk has any - # remaining content (delta, finish_reason, etc.). - # is_model_response_stream_empty sees usage=None and - # correctly skips it, only checking meaningful fields. is_empty = is_model_response_stream_empty( model_response=cast(ModelResponseStream, processed_chunk) ) 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)" + )