refactor: clean class with static methods, add tests, fix sentinel pattern

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()

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