fix: preserve disconnect events in api key middleware (#17253)

python/sglang/srt/utils/auth.py

@@ -155,31 +155,54 @@ def add_api_key_middleware(
     """Add middleware for three endpoint auth levels: normal/admin_optional/admin_force."""
     # Import lazily so `decide_request_auth()` can be unit-tested without FastAPI installed.
     from fastapi.responses import ORJSONResponse
+    from starlette.requests import Request
 
-    @app.middleware("http")
-    async def authentication(request, call_next):
-        path = request.url.path
-        authz = request.headers.get("Authorization")
-        level = _get_auth_level_from_app_and_scope(request.app, request.scope)
-        decision = decide_request_auth(
-            method=request.method,
-            path=path,
-            authorization_header=authz,
-            api_key=api_key,
-            admin_api_key=admin_api_key,
-            auth_level=level,
-        )
+    class _ApiKeyASGIMiddleware:
+        """ASGI-native middleware to preserve client disconnect events."""
 
-        if not decision.allowed:
-            return ORJSONResponse(
-                content={
-                    "error": (
-                        "Unauthorized"
-                        if decision.error_status_code == 401
-                        else "Forbidden"
-                    )
-                },
-                status_code=decision.error_status_code,
+        def __init__(self, app, *, api_key, admin_api_key, fastapi_app):
+            self.app = app
+            self.api_key = api_key
+            self.admin_api_key = admin_api_key
+            self.fastapi_app = fastapi_app
+
+        async def __call__(self, scope, receive, send):
+            if scope["type"] != "http":
+                await self.app(scope, receive, send)
+                return
+
+            request = Request(scope, receive=receive)
+            path = request.url.path
+            authz = request.headers.get("Authorization")
+            level = _get_auth_level_from_app_and_scope(self.fastapi_app, scope)
+            decision = decide_request_auth(
+                method=request.method,
+                path=path,
+                authorization_header=authz,
+                api_key=self.api_key,
+                admin_api_key=self.admin_api_key,
+                auth_level=level,
             )
 
-        return await call_next(request)
+            if not decision.allowed:
+                response = ORJSONResponse(
+                    content={
+                        "error": (
+                            "Unauthorized"
+                            if decision.error_status_code == 401
+                            else "Forbidden"
+                        )
+                    },
+                    status_code=decision.error_status_code,
+                )
+                await response(scope, receive, send)
+                return
+
+            await self.app(scope, receive, send)
+
+    app.add_middleware(
+        _ApiKeyASGIMiddleware,
+        api_key=api_key,
+        admin_api_key=admin_api_key,
+        fastapi_app=app,
+    )

python/sglang/test/test_utils.py

@@ -1662,6 +1662,7 @@ def run_and_check_memory_leak(
     disable_overlap,
     chunked_prefill_size,
     assert_has_abort,
+    api_key: Optional[str] = None,
 ):
     other_args = [
         "--chunked-prefill-size",
@@ -1689,6 +1690,7 @@ def run_and_check_memory_leak(
         timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
         other_args=other_args,
         return_stdout_stderr=(stdout, stderr),
+        api_key=api_key,
     )
 
     # Launch a thread to stream the output

test/registered/scheduler/test_abort.py

@@ -63,6 +63,51 @@ class TestAbort(CustomTestCase):
         )
 
 
+class TestAbortWithApiKey(CustomTestCase):
+    def workload_func(self, base_url, model, api_key: str):
+        def process_func():
+            def run_one(_):
+                prompt = """
+                System: You are a helpful assistant.
+                User: What is the capital of France?
+                Assistant: The capital of France is
+                """
+
+                response = requests.post(
+                    f"{base_url}/generate",
+                    json={
+                        "text": prompt,
+                        "sampling_params": {
+                            "temperature": 0,
+                            "max_new_tokens": 2048,
+                        },
+                    },
+                    headers={"Authorization": f"Bearer {api_key}"},
+                )
+                response.json()
+
+            with ThreadPoolExecutor(16) as executor:
+                list(executor.map(run_one, list(range(16))))
+
+        p = multiprocessing.Process(target=process_func)
+        p.start()
+        time.sleep(0.5)
+        p.terminate()
+        time.sleep(10)
+
+    def test_memory_leak_with_api_key(self):
+        api_key = "test-api-key"
+        run_and_check_memory_leak(
+            lambda base_url, model: self.workload_func(base_url, model, api_key),
+            disable_radix_cache=False,
+            enable_mixed_chunk=False,
+            disable_overlap=False,
+            chunked_prefill_size=8192,
+            assert_has_abort=True,
+            api_key=api_key,
+        )
+
+
 class TestAbortAll(CustomTestCase):
     @classmethod
     def setUpClass(cls):