From 2b47bd3a348bf5ca52a3a4910b2e22c851798576 Mon Sep 17 00:00:00 2001 From: fanghao <971360589@qq.com> Date: Mon, 23 Mar 2026 10:58:49 +0800 Subject: [PATCH] [Bug Fix] Fix non-streaming request abort failure when --enable-metrics is enabled (#20625) --- python/sglang/srt/utils/common.py | 6 ++ .../sglang/srt/utils/http_middleware_patch.py | 67 ++++++++++++++++ .../scheduler/test_abort_with_metrics.py | 78 +++++++++++++++++++ 3 files changed, 151 insertions(+) create mode 100644 python/sglang/srt/utils/http_middleware_patch.py create mode 100644 test/registered/scheduler/test_abort_with_metrics.py diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index 857df6446..a4a90b5b1 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -1344,6 +1344,12 @@ def add_prometheus_track_response_middleware(app): ) ) + # Fix: replace BaseHTTPMiddleware's call_next with a pure ASGI version + # that passes `receive` through, so request.is_disconnected() keeps working. + from sglang.srt.utils.http_middleware_patch import patch_app_http_middleware + + patch_app_http_middleware(app) + @app.middleware("http") async def track_http_status_code(request, call_next): # With recording all requests, we have the risk of high cardinality if requests have arbitrary unhandled paths. diff --git a/python/sglang/srt/utils/http_middleware_patch.py b/python/sglang/srt/utils/http_middleware_patch.py new file mode 100644 index 000000000..a334a0c27 --- /dev/null +++ b/python/sglang/srt/utils/http_middleware_patch.py @@ -0,0 +1,67 @@ +""" +Fix @app.middleware("http") whose BaseHTTPMiddleware call_next replaces +ASGI ``receive``, breaking request.is_disconnected() and preventing +non-streaming request abort on client disconnect. + +patch_app_http_middleware(app) replaces @app.middleware("http") with a +version whose call_next passes ``receive`` through untouched. +""" + +from __future__ import annotations + +from starlette.requests import Request + + +class _SentResponse: + """Response proxy returned after the real response was already sent.""" + + def __init__(self, status_code: int): + self.status_code = status_code + + +class _PureASGIDispatch: + """Pure ASGI middleware providing a fixed call_next that passes + ``receive`` through untouched (unlike BaseHTTPMiddleware).""" + + def __init__(self, app, dispatch): + self.app = app + self.dispatch = dispatch + + async def __call__(self, scope, receive, send): + if scope["type"] != "http": + await self.app(scope, receive, send) + return + + request = Request(scope, receive) + status_code = 500 + + async def call_next(_request): + nonlocal status_code + + async def send_and_capture(message): + nonlocal status_code + if message["type"] == "http.response.start": + status_code = message["status"] + await send(message) + + await self.app(scope, receive, send_and_capture) + return _SentResponse(status_code) + + await self.dispatch(request, call_next) + + +def patch_app_http_middleware(app): + """Replace @app.middleware("http") with a fixed-call_next version.""" + _orig = app.middleware + + def _fixed(middleware_type): + if middleware_type == "http": + + def decorator(fn): + app.add_middleware(_PureASGIDispatch, dispatch=fn) + return fn + + return decorator + return _orig(middleware_type) + + app.middleware = _fixed diff --git a/test/registered/scheduler/test_abort_with_metrics.py b/test/registered/scheduler/test_abort_with_metrics.py new file mode 100644 index 000000000..c2adfdd54 --- /dev/null +++ b/test/registered/scheduler/test_abort_with_metrics.py @@ -0,0 +1,78 @@ +""" +Unit test for _PureASGIDispatch: verify that the ASGI ``receive`` callable +is passed through untouched so that request.is_disconnected() works. + +Background: @app.middleware("http") wraps handlers with BaseHTTPMiddleware +whose call_next() replaces the ASGI ``receive``, breaking +request.is_disconnected() and preventing non-streaming abort on client +disconnect. _PureASGIDispatch fixes this. The existing test_abort.py +already covers the full e2e abort flow. +""" + +import asyncio +import unittest + +from starlette.requests import Request + +from sglang.srt.utils.http_middleware_patch import _PureASGIDispatch +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=10, suite="stage-a-cpu-only") + +_HTTP_SCOPE = { + "type": "http", + "asgi": {"version": "3.0"}, + "http_version": "1.1", + "method": "POST", + "path": "/test", + "query_string": b"", + "root_path": "", + "headers": [], +} + + +class TestPureASGIDispatchReceivePassthrough(CustomTestCase): + """Verify _PureASGIDispatch passes ``receive`` through untouched.""" + + @staticmethod + async def _run_with_receive(receive_msg): + """Invoke _PureASGIDispatch and return request.is_disconnected().""" + result = {} + + async def dispatch(request: Request, call_next): + result["disconnected"] = await request.is_disconnected() + await call_next(request) + + async def inner_app(scope, receive, send): + await send({"type": "http.response.start", "status": 200, "headers": []}) + await send({"type": "http.response.body", "body": b""}) + + middleware = _PureASGIDispatch(inner_app, dispatch=dispatch) + + async def receive(): + return receive_msg + + sent = [] + + async def send(msg): + sent.append(msg) + + await middleware(_HTTP_SCOPE, receive, send) + return result["disconnected"] + + def test_is_disconnected_on_client_disconnect(self): + """receive() -> http.disconnect: is_disconnected() must return True.""" + self.assertTrue( + asyncio.run(self._run_with_receive({"type": "http.disconnect"})) + ) + + def test_not_disconnected_when_connected(self): + """receive() -> http.request: is_disconnected() must return False.""" + self.assertFalse( + asyncio.run(self._run_with_receive({"type": "http.request", "body": b""})) + ) + + +if __name__ == "__main__": + unittest.main()