feat: add SSE keep-alive comments for idle streaming responses (#51034)
Signed-off-by: fuzzifikation <40043768+fuzzifikation@users.noreply.github.com>
Dieter Schwarzmann committed
Aug 24, 2026 at 03:13 UTC
585bb07c7001ce52737ff0d18183d7a4b663cce4
6 files changed
+381
-2
tests/entrypoints/openai/test_cli_args.py
+30
@@ -258,6 +258,36 @@ def test_subcommand_without_serve_args_skips_validation():
258
validate_parsed_serve_args(args)
259
260
261
+@pytest.mark.parametrize(
262
+ "cli_args, expected",
263
+ [
264
+ ([], 0),
265
+ (["--sse-keep-alive-interval", "0"], 0),
266
+ (["--sse-keep-alive-interval", "15"], 15),
267
+ ],
268
+)
269
+def test_sse_keep_alive_interval_valid(serve_parser, cli_args, expected):
270
+ """The default is 0 and zero or positive integers parse exactly."""
271
+ args = serve_parser.parse_args(args=cli_args)
272
+ assert args.sse_keep_alive_interval == expected
273
+ validate_parsed_serve_args(args)
274
+
275
+
276
+@pytest.mark.parametrize("value", ["-1", "-5"])
277
+def test_sse_keep_alive_interval_negative(serve_parser, value):
278
+ """Negative intervals are rejected by validation."""
279
+ args = serve_parser.parse_args(args=["--sse-keep-alive-interval", value])
280
+ with pytest.raises(ValueError):
281
+ validate_parsed_serve_args(args)
282
+
283
+
284
+@pytest.mark.parametrize("value", ["nan", "inf", "-inf", "0.5", "abc"])
285
+def test_sse_keep_alive_interval_non_integer(serve_parser, value):
286
+ """Non-integer values fail argparse's int parsing with SystemExit."""
287
+ with pytest.raises(SystemExit):
288
+ serve_parser.parse_args(args=["--sse-keep-alive-interval", value])
289
+
290
+
291
@pytest.mark.parametrize(
292
"cli_args, expected_middleware",
293
[
tests/entrypoints/openai/test_sse_keep_alive.py
new
+258
@@ -0,0 +1,258 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import asyncio
5
+
6
+import pytest
7
+from starlette.responses import StreamingResponse
8
+
9
+from vllm.entrypoints.openai.sse_keep_alive import (
10
+ SSE_KEEP_ALIVE_COMMENT,
11
+ with_sse_keep_alive,
12
+)
13
+
14
+# --- Disabled path ----------------------------------------------------------
15
+
16
+
17
+def test_disabled_returns_same_generator():
18
+ """For 0, negative, and NaN intervals the exact same generator object must
19
+ be returned (identity), proving no wrapper or per-chunk overhead when
20
+ disabled."""
21
+
22
+ async def gen():
23
+ yield "x"
24
+
25
+ for disabled in (0.0, -5.0, float("nan"), float("inf")):
26
+ g = gen()
27
+ assert with_sse_keep_alive(g, disabled) is g
28
+
29
+
30
+# --- Keep-alive emission ----------------------------------------------------
31
+
32
+
33
+@pytest.mark.asyncio
34
+async def test_keep_alive_during_silent_start():
35
+ """A silent start (queue wait / prefill) emits keep-alive comments before
36
+ the first real chunk."""
37
+ release = asyncio.Event()
38
+
39
+ async def gen():
40
+ await release.wait() # simulates queue wait / prefill
41
+ yield "data: first\n\n"
42
+
43
+ out: list[str] = []
44
+ async for item in with_sse_keep_alive(gen(), 0.01):
45
+ out.append(item)
46
+ if item == SSE_KEEP_ALIVE_COMMENT:
47
+ release.set()
48
+ if item == "data: first\n\n":
49
+ break
50
+
51
+ assert SSE_KEEP_ALIVE_COMMENT in out
52
+ assert out.index(SSE_KEEP_ALIVE_COMMENT) < out.index("data: first\n\n")
53
+
54
+
55
+@pytest.mark.asyncio
56
+async def test_no_keep_alive_when_streaming_is_fast():
57
+ """If chunks arrive faster than the interval, no comment is emitted and all
58
+ data is forwarded."""
59
+
60
+ async def gen():
61
+ for chunk in ["a", "b", "c"]:
62
+ yield chunk
63
+
64
+ out = [c async for c in with_sse_keep_alive(gen(), 0.5)]
65
+ assert out == ["a", "b", "c"]
66
+ assert SSE_KEEP_ALIVE_COMMENT not in out
67
+
68
+
69
+@pytest.mark.asyncio
70
+async def test_keep_alive_between_chunks_does_not_drop_data():
71
+ """Keep-alives interleaved with real chunks must never drop or reorder
72
+ data."""
73
+ release_b = asyncio.Event()
74
+ release_c = asyncio.Event()
75
+
76
+ async def gen():
77
+ yield "data: a\n\n"
78
+ await release_b.wait()
79
+ yield "data: b\n\n"
80
+ await release_c.wait()
81
+ yield "data: c\n\n"
82
+
83
+ out: list[str] = []
84
+ async for item in with_sse_keep_alive(gen(), 0.01):
85
+ out.append(item)
86
+ if item == SSE_KEEP_ALIVE_COMMENT and not release_b.is_set():
87
+ release_b.set()
88
+ elif item == "data: b\n\n" and not release_c.is_set():
89
+ release_c.set()
90
+ elif item == "data: c\n\n":
91
+ break
92
+
93
+ assert SSE_KEEP_ALIVE_COMMENT in out
94
+ data_chunks = [c for c in out if c != SSE_KEEP_ALIVE_COMMENT]
95
+ assert data_chunks == ["data: a\n\n", "data: b\n\n", "data: c\n\n"]
96
+
97
+
98
+@pytest.mark.asyncio
99
+async def test_empty_generator_returns_immediately():
100
+ async def gen():
101
+ if False:
102
+ yield "never"
103
+
104
+ out = [c async for c in with_sse_keep_alive(gen(), 0.01)]
105
+ assert out == []
106
+
107
+
108
+# --- Errors and cancellation ------------------------------------------------
109
+
110
+
111
+@pytest.mark.asyncio
112
+async def test_exception_propagates():
113
+ async def gen():
114
+ yield "data: before\n\n"
115
+ raise RuntimeError("boom")
116
+
117
+ wrapped = with_sse_keep_alive(gen(), 0.01)
118
+ assert await anext(wrapped) == "data: before\n\n"
119
+ with pytest.raises(RuntimeError, match="boom"):
120
+ await anext(wrapped)
121
+
122
+
123
+@pytest.mark.asyncio
124
+async def test_upstream_cancelled_error_propagates():
125
+ async def gen():
126
+ yield "data: before\n\n"
127
+ raise asyncio.CancelledError("upstream cancelled")
128
+
129
+ wrapped = with_sse_keep_alive(gen(), 0.01)
130
+ assert await anext(wrapped) == "data: before\n\n"
131
+ with pytest.raises(asyncio.CancelledError):
132
+ await anext(wrapped)
133
+
134
+
135
+@pytest.mark.asyncio
136
+async def test_aclose_closes_upstream():
137
+ """Closing a wrapper that is mid-stream must complete promptly and close
138
+ the upstream generator (regression for the earlier leaked-task hang)."""
139
+ finalized = asyncio.Event()
140
+
141
+ async def gen():
142
+ try:
143
+ yield "chunk-0"
144
+ await asyncio.Event().wait() # block forever
145
+ finally:
146
+ finalized.set()
147
+
148
+ wrapped = with_sse_keep_alive(gen(), 10.0)
149
+ assert await anext(wrapped) == "chunk-0"
150
+ await asyncio.wait_for(wrapped.aclose(), timeout=2.0)
151
+ assert finalized.is_set()
152
+
153
+
154
+@pytest.mark.asyncio
155
+async def test_client_disconnect_cancels_upstream():
156
+ """Cancelling the consumer mid-stream must cancel the pending anext and
157
+ close the upstream generator so its cleanup runs."""
158
+ finalized = asyncio.Event()
159
+
160
+ async def gen():
161
+ try:
162
+ for i in range(1000):
163
+ yield f"chunk-{i}"
164
+ await asyncio.Event().wait()
165
+ finally:
166
+ finalized.set()
167
+
168
+ wrapped = with_sse_keep_alive(gen(), 10.0)
169
+ assert await anext(wrapped) == "chunk-0"
170
+ task = asyncio.current_task()
171
+ assert task is not None
172
+ task.cancel()
173
+ with pytest.raises(asyncio.CancelledError):
174
+ await anext(wrapped)
175
+ assert finalized.is_set()
176
+
177
+
178
+# --- Response-level integration (the shipped StreamingResponse path) --------
179
+
180
+
181
+@pytest.mark.asyncio
182
+async def test_streaming_response_emits_comment():
183
+ """A keep-alive comment must reach the client through the real
184
+ StreamingResponse used by the routers."""
185
+ release = asyncio.Event()
186
+
187
+ async def gen():
188
+ await release.wait()
189
+ yield "data: hi\n\n"
190
+
191
+ response = StreamingResponse(
192
+ content=with_sse_keep_alive(gen(), 0.01),
193
+ media_type="text/event-stream",
194
+ )
195
+ body: list[bytes] = []
196
+ started = asyncio.Event()
197
+
198
+ async def send(message):
199
+ if message["type"] == "http.response.start":
200
+ started.set()
201
+ elif message["type"] == "http.response.body":
202
+ body.append(message["body"])
203
+ if SSE_KEEP_ALIVE_COMMENT.encode() in message["body"]:
204
+ release.set()
205
+
206
+ async def receive():
207
+ await asyncio.Event().wait()
208
+
209
+ await response(
210
+ {"type": "http", "asgi": {"spec_version": "2.3"}},
211
+ receive,
212
+ send,
213
+ )
214
+ data = b"".join(body).decode()
215
+ assert SSE_KEEP_ALIVE_COMMENT in data
216
+ assert "data: hi" in data
217
+
218
+
219
+@pytest.mark.asyncio
220
+async def test_streaming_response_disconnect_closes_upstream():
221
+ """On the ASGI 2.3 disconnect path used by the pinned Uvicorn, Starlette
222
+ cancels streaming and the wrapper's finally must close the upstream
223
+ generator."""
224
+ finalized = asyncio.Event()
225
+
226
+ async def gen():
227
+ try:
228
+ await asyncio.Event().wait()
229
+ yield "never"
230
+ finally:
231
+ finalized.set()
232
+
233
+ response = StreamingResponse(
234
+ content=with_sse_keep_alive(gen(), 10.0),
235
+ media_type="text/event-stream",
236
+ )
237
+ started = asyncio.Event()
238
+
239
+ async def receive():
240
+ await started.wait()
241
+ return {"type": "http.disconnect"}
242
+
243
+ async def send(message):
244
+ if message["type"] == "http.response.start":
245
+ started.set()
246
+
247
+ await response(
248
+ {"type": "http", "asgi": {"spec_version": "2.3"}},
249
+ receive,
250
+ send,
251
+ )
252
+ assert finalized.is_set()
253
+ pending = [
254
+ t
255
+ for t in asyncio.all_tasks()
256
+ if t is not asyncio.current_task() and not t.done()
257
+ ]
258
+ assert pending == []
vllm/entrypoints/openai/chat_completion/api_router.py
+7
-1
@@ -15,6 +15,7 @@ from vllm.entrypoints.openai.chat_completion.protocol import (
15
)
16
from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat
17
from vllm.entrypoints.openai.engine.protocol import ErrorResponse
18
+from vllm.entrypoints.openai.sse_keep_alive import with_sse_keep_alive
19
from vllm.entrypoints.serve.utils.api_utils import (
20
load_aware_call,
21
validate_json_request,
@@ -71,7 +72,12 @@ async def create_chat_completion(request: ChatCompletionRequest, raw_request: Re
72
headers=metrics_header(metrics_header_format),
73
)
74
74
- return StreamingResponse(content=generator, media_type="text/event-stream")
75
+ args = getattr(raw_request.app.state, "args", None)
76
+ keep_alive_interval = getattr(args, "sse_keep_alive_interval", 0)
77
+ return StreamingResponse(
78
+ content=with_sse_keep_alive(generator, float(keep_alive_interval)),
79
+ media_type="text/event-stream",
80
+ )
81
82
83
@router.post(
vllm/entrypoints/openai/cli_args.py
+14
@@ -137,6 +137,13 @@ class BaseFrontendArgs:
137
"""If set to True, enable tracking server_load_metrics in the app state."""
138
enable_force_include_usage: bool = False
139
"""If set to True, including usage on every request."""
140
+ sse_keep_alive_interval: int = 0
141
+ """Send an SSE keep-alive comment line every this many seconds when a
142
+ `/v1/chat/completions` or `/v1/completions` streaming response is idle
143
+ (queued, prefill, or between tokens), to prevent reverse proxies/tunnels
144
+ with read timeouts from closing the connection. Defaults to 0, which
145
+ disables keep-alive comments entirely.
146
+ """
147
enable_tokenizer_info_endpoint: bool = False
148
"""Enable the `/tokenizer_info` endpoint. May expose chat
149
templates and other tokenizer configuration."""
@@ -426,6 +433,13 @@ def validate_parsed_serve_args(args: argparse.Namespace):
433
if args.enable_log_outputs and not args.enable_log_requests:
434
raise TypeError("Error: --enable-log-outputs requires --enable-log-requests")
435
436
+ # SSE keep-alive interval must be zero (disabled) or a positive integer.
437
+ if getattr(args, "sse_keep_alive_interval", 0) < 0:
438
+ raise ValueError(
439
+ "Error: --sse-keep-alive-interval must be a non-negative integer "
440
+ "(0 disables keep-alive comments)."
441
+ )
442
+
443
if getattr(args, "enable_per_request_metrics", False) and getattr(
444
args, "disable_log_stats", False
445
):
vllm/entrypoints/openai/completion/api_router.py
+7
-1
@@ -13,6 +13,7 @@ from vllm.entrypoints.openai.completion.protocol import (
13
)
14
from vllm.entrypoints.openai.completion.serving import OpenAIServingCompletion
15
from vllm.entrypoints.openai.engine.protocol import ErrorResponse
16
+from vllm.entrypoints.openai.sse_keep_alive import with_sse_keep_alive
17
from vllm.entrypoints.serve.utils.api_utils import (
18
load_aware_call,
19
validate_json_request,
@@ -63,7 +64,12 @@ async def create_completion(request: CompletionRequest, raw_request: Request):
64
headers=metrics_header(metrics_header_format),
65
)
66
66
- return StreamingResponse(content=generator, media_type="text/event-stream")
67
+ args = getattr(raw_request.app.state, "args", None)
68
+ keep_alive_interval = getattr(args, "sse_keep_alive_interval", 0)
69
+ return StreamingResponse(
70
+ content=with_sse_keep_alive(generator, float(keep_alive_interval)),
71
+ media_type="text/event-stream",
72
+ )
73
74
75
def attach_router(app: FastAPI):
vllm/entrypoints/openai/sse_keep_alive.py
new
+65
@@ -0,0 +1,65 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""SSE keep-alive comments for idle streaming responses.
4
+
5
+Reverse proxies and tunnels (Cloudflare Tunnel, NGINX ``proxy_read_timeout``,
6
+AWS ALB, ...) can close a streaming connection when no bytes are sent for a
7
+while. That happens while a request is queued, during prefill, or in gaps
8
+between tokens. ``with_sse_keep_alive`` wraps the final SSE generator and emits
9
+a comment line (``": keep-alive\\n\\n"``) whenever it has been idle for
10
+``interval`` seconds. Comments are ignored by SSE-compliant clients but still
11
+count as bytes for proxy read-timeout logic.
12
+"""
13
+
14
+import asyncio
15
+import contextlib
16
+from collections.abc import AsyncGenerator
17
+
18
+SSE_KEEP_ALIVE_COMMENT = ": keep-alive\n\n"
19
+
20
+
21
+def with_sse_keep_alive(
22
+ generator: AsyncGenerator[str, None],
23
+ interval: float,
24
+) -> AsyncGenerator[str, None]:
25
+ """Emit an SSE keep-alive comment when ``generator`` is idle.
26
+
27
+ A non-positive or non-finite ``interval`` returns ``generator`` unchanged,
28
+ so the default path has no overhead. Otherwise a keep-alive comment is
29
+ yielded whenever no chunk arrives within ``interval`` seconds.
30
+ """
31
+ if not 0 < interval < float("inf"):
32
+ return generator
33
+ return _keep_alive_stream(generator, interval)
34
+
35
+
36
+async def _keep_alive_stream(
37
+ generator: AsyncGenerator[str, None],
38
+ interval: float,
39
+) -> AsyncGenerator[str, None]:
40
+ # A single retained ``anext()`` task advances the generator; asyncio forbids
41
+ # concurrent ``__anext__`` calls, and ``wait_for`` would cancel the task
42
+ # mid-advance on timeout. ``asyncio.wait`` observes both outcomes without
43
+ # cancelling the pending advance.
44
+ next_task: asyncio.Task[str] = asyncio.create_task(anext(generator))
45
+ try:
46
+ while True:
47
+ done, _ = await asyncio.wait({next_task}, timeout=interval)
48
+ if not done:
49
+ yield SSE_KEEP_ALIVE_COMMENT
50
+ continue
51
+ try:
52
+ chunk = next_task.result()
53
+ except StopAsyncIteration:
54
+ return
55
+ yield chunk
56
+ next_task = asyncio.create_task(anext(generator))
57
+ finally:
58
+ next_task.cancel()
59
+ # The task outcome was already handled in the loop; suppress so cleanup
60
+ # never corrupts it (matching ``merge_async_iterators``).
61
+ with contextlib.suppress(BaseException):
62
+ await next_task
63
+ # Close the upstream generator so its cleanup runs deterministically.
64
+ with contextlib.suppress(BaseException):
65
+ await generator.aclose()