From 03500799bee3bd6d544d8e9f17ee961026595faf Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Fri, 27 Feb 2026 17:47:00 -0800 Subject: [PATCH] refactor: clean class with static methods, add tests, fix sentinel pattern --- .../in_flight_requests_middleware.py | 86 ++++++++-------- .../test_in_flight_requests_middleware.py | 98 +++++++++++++++++++ 2 files changed, 144 insertions(+), 40 deletions(-) create mode 100644 tests/test_litellm/proxy/middleware/test_in_flight_requests_middleware.py diff --git a/litellm/proxy/middleware/in_flight_requests_middleware.py b/litellm/proxy/middleware/in_flight_requests_middleware.py index 0c0476373eb..d615640d870 100644 --- a/litellm/proxy/middleware/in_flight_requests_middleware.py +++ b/litellm/proxy/middleware/in_flight_requests_middleware.py @@ -10,50 +10,26 @@ from typing import Optional from starlette.types import ASGIApp, Receive, Scope, Send -_in_flight: int = 0 - -# Lazily created on first request so PROMETHEUS_MULTIPROC_DIR is already set -# by the time we register the metric. -_gauge: Optional[object] = None - - -def _get_gauge() -> Optional[object]: - global _gauge - if _gauge is not None: - return _gauge - try: - from prometheus_client import Gauge - - kwargs = {} - if "PROMETHEUS_MULTIPROC_DIR" in os.environ: - # livesum aggregates across all worker processes in the scrape response - kwargs["multiprocess_mode"] = "livesum" - _gauge = Gauge( - "litellm_in_flight_requests", - "Number of HTTP requests currently in-flight on this uvicorn worker", - **kwargs, - ) - except Exception: - pass - return _gauge - - -def get_in_flight_requests() -> int: - return _in_flight - class InFlightRequestsMiddleware: """ - ASGI middleware that increments a counter when a request arrives - and decrements it when the response is sent (or an error occurs). + ASGI middleware that increments a counter when a request arrives and + decrements it when the response is sent (or an error occurs). - The counter is module-level and therefore scoped to a single uvicorn - worker process — exactly the per-pod granularity we want. + The counter is class-level and therefore scoped to a single uvicorn worker + process — exactly the per-pod granularity we want. Also updates the `litellm_in_flight_requests` Prometheus gauge if - prometheus_client is installed. + prometheus_client is installed. The gauge is lazily initialised on the + first request so that PROMETHEUS_MULTIPROC_DIR is already set by the time + we register the metric. Initialisation is attempted only once — if + prometheus_client is absent the class remembers and never retries. """ + _in_flight: int = 0 + _gauge: Optional[object] = None + _gauge_init_attempted: bool = False + def __init__(self, app: ASGIApp) -> None: self.app = app @@ -62,14 +38,44 @@ class InFlightRequestsMiddleware: await self.app(scope, receive, send) return - global _in_flight - _in_flight += 1 - gauge = _get_gauge() + InFlightRequestsMiddleware._in_flight += 1 + gauge = InFlightRequestsMiddleware._get_gauge() if gauge is not None: gauge.inc() # type: ignore[union-attr] try: await self.app(scope, receive, send) finally: - _in_flight -= 1 + InFlightRequestsMiddleware._in_flight -= 1 if gauge is not None: gauge.dec() # type: ignore[union-attr] + + @staticmethod + def get_count() -> int: + """Return the number of HTTP requests currently in-flight.""" + return InFlightRequestsMiddleware._in_flight + + @staticmethod + def _get_gauge() -> Optional[object]: + if InFlightRequestsMiddleware._gauge_init_attempted: + return InFlightRequestsMiddleware._gauge + InFlightRequestsMiddleware._gauge_init_attempted = True + try: + from prometheus_client import Gauge + + kwargs = {} + if "PROMETHEUS_MULTIPROC_DIR" in os.environ: + # livesum aggregates across all worker processes in the scrape response + kwargs["multiprocess_mode"] = "livesum" + InFlightRequestsMiddleware._gauge = Gauge( + "litellm_in_flight_requests", + "Number of HTTP requests currently in-flight on this uvicorn worker", + **kwargs, + ) + except Exception: + InFlightRequestsMiddleware._gauge = None + return InFlightRequestsMiddleware._gauge + + +def get_in_flight_requests() -> int: + """Module-level convenience wrapper used by the /health/backlog endpoint.""" + return InFlightRequestsMiddleware.get_count() diff --git a/tests/test_litellm/proxy/middleware/test_in_flight_requests_middleware.py b/tests/test_litellm/proxy/middleware/test_in_flight_requests_middleware.py new file mode 100644 index 00000000000..830bca49936 --- /dev/null +++ b/tests/test_litellm/proxy/middleware/test_in_flight_requests_middleware.py @@ -0,0 +1,98 @@ +""" +Tests for InFlightRequestsMiddleware. + +Verifies that in_flight_requests is incremented during a request and +decremented after it completes, including on errors. +""" +import asyncio + +import pytest +from starlette.middleware.base import BaseHTTPMiddleware +from starlette.requests import Request +from starlette.responses import JSONResponse, Response +from starlette.routing import Route +from starlette.testclient import TestClient + +from litellm.proxy.middleware.in_flight_requests_middleware import ( + InFlightRequestsMiddleware, + get_in_flight_requests, +) + + +@pytest.fixture(autouse=True) +def reset_state(): + """Reset class-level state between tests.""" + InFlightRequestsMiddleware._in_flight = 0 + yield + InFlightRequestsMiddleware._in_flight = 0 + + +def _make_app(handler): + from starlette.applications import Starlette + + app = Starlette(routes=[Route("/", handler)]) + app.add_middleware(InFlightRequestsMiddleware) + return app + + +# ── Structure ───────────────────────────────────────────────────────────────── + + +def test_is_not_base_http_middleware(): + """Must be pure ASGI — BaseHTTPMiddleware causes streaming degradation.""" + assert not issubclass(InFlightRequestsMiddleware, BaseHTTPMiddleware) + + +def test_has_asgi_call_protocol(): + assert "__call__" in InFlightRequestsMiddleware.__dict__ + + +# ── Counter behaviour ───────────────────────────────────────────────────────── + + +def test_counter_zero_at_start(): + assert get_in_flight_requests() == 0 + + +def test_counter_increments_inside_handler(): + captured = [] + + async def handler(request: Request) -> Response: + captured.append(InFlightRequestsMiddleware.get_count()) + return JSONResponse({}) + + TestClient(_make_app(handler)).get("/") + assert captured == [1] + + +def test_counter_returns_to_zero_after_request(): + async def handler(request: Request) -> Response: + return JSONResponse({}) + + TestClient(_make_app(handler)).get("/") + assert get_in_flight_requests() == 0 + + +def test_counter_decrements_after_error(): + """Counter must reach 0 even when the handler raises.""" + + async def handler(request: Request) -> Response: + return Response("boom", status_code=500) + + TestClient(_make_app(handler)).get("/") + assert get_in_flight_requests() == 0 + + +def test_non_http_scopes_not_counted(): + """Lifespan / websocket scopes must not touch the counter.""" + + class _InnerApp: + async def __call__(self, scope, receive, send): + pass + + mw = InFlightRequestsMiddleware(_InnerApp()) + + asyncio.get_event_loop().run_until_complete( + mw({"type": "lifespan"}, None, None) # type: ignore[arg-type] + ) + assert get_in_flight_requests() == 0