diff --git a/python/sglang/srt/utils/auth.py b/python/sglang/srt/utils/auth.py index 0536d5a35..e7ef0767f 100644 --- a/python/sglang/srt/utils/auth.py +++ b/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, + ) diff --git a/python/sglang/test/test_utils.py b/python/sglang/test/test_utils.py index fe8c3f9c3..298308296 100644 --- a/python/sglang/test/test_utils.py +++ b/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 diff --git a/test/registered/scheduler/test_abort.py b/test/registered/scheduler/test_abort.py index 466891608..27ba6406e 100644 --- a/test/registered/scheduler/test_abort.py +++ b/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):