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

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