diff --git a/python/freetoken/server/api_server.py b/python/freetoken/server/api_server.py index a9b16adfd7..83b785252b 100644 --- a/python/freetoken/server/api_server.py +++ b/python/freetoken/server/api_server.py @@ -60,6 +60,10 @@ # shutdown is treated as expected — no ERROR log, no "failed" latch. See run_backend_supervisor. _SHUTTING_DOWN = threading.Event() BACKEND_DEATH_EXIT_GRACE_S = 10.0 +# The longest a stream stays silent before it sends an SSE comment. A long prompt prefills for +# minutes before its first token, and clients time an idle body out (Node's fetch after 300 s). +STREAM_KEEPALIVE_SECONDS = 15.0 +_STREAM_KEEPALIVE = b": keep-alive\n\n" def get_global_state() -> FrontendManager: @@ -362,16 +366,33 @@ async def stream_generate(self, uid: int): logger.debug("Finished streaming response for user %s", uid) async def stream_with_cancellation(self, generator, request: Request, uid: int): + """Relay ``generator``'s SSE chunks; while it is silent (a long prefill), send a comment every + ``STREAM_KEEPALIVE_SECONDS`` and abort the request if the client has gone.""" + chunks = generator.__aiter__() + pending: asyncio.Future | None = None try: - async for chunk in generator: - # detect if the client has disconnected + while True: + if pending is None: + pending = asyncio.ensure_future(chunks.__anext__()) + done, _ = await asyncio.wait({pending}, timeout=STREAM_KEEPALIVE_SECONDS) if await request.is_disconnected(): logger.info("Client disconnected for user %s", uid) raise asyncio.CancelledError + if not done: + yield _STREAM_KEEPALIVE + continue + try: + chunk = pending.result() + except StopAsyncIteration: + return + pending = None yield chunk except asyncio.CancelledError: asyncio.create_task(self.abort_user(uid)) raise + finally: + if pending is not None and not pending.done(): + pending.cancel() async def abort_user(self, uid: int): await asyncio.sleep(0.1) diff --git a/tests/server/test_stream_keepalive.py b/tests/server/test_stream_keepalive.py new file mode 100644 index 0000000000..e86d4365fc --- /dev/null +++ b/tests/server/test_stream_keepalive.py @@ -0,0 +1,115 @@ +"""A stream that stays silent -- the prefill of a long prompt -- keeps its client, and notices when +the client leaves instead of prefilling for nobody.""" + +from __future__ import annotations + +import asyncio +from types import SimpleNamespace + +import pytest + +from freetoken.server import api_server +from freetoken.server.api_server import FrontendManager + +# Converts a stream that never ends into a failure; the streams below finish in milliseconds. +HANG_GUARD_SECONDS = 10.0 + + +class _Request: + def __init__(self) -> None: + self.gone = False + + async def is_disconnected(self) -> bool: + return self.gone + + +def _manager() -> SimpleNamespace: + async def abort_user(uid: int) -> None: + pass + + return SimpleNamespace(abort_user=abort_user) + + +def _is_sse_comment(chunk: bytes) -> bool: + return chunk.startswith(b":") and chunk.endswith(b"\n\n") + + +def test_a_silent_stream_sends_comments_until_its_next_chunk(monkeypatch): + monkeypatch.setattr(api_server, "STREAM_KEEPALIVE_SECONDS", 0.01) + released = asyncio.Event() + + async def chunks(): + yield b"data: first\n\n" + await released.wait() + yield b"data: last\n\n" + + async def consume() -> list[bytes]: + out = [] + request = _Request() + async for chunk in FrontendManager.stream_with_cancellation(_manager(), chunks(), request, 7): + out.append(chunk) + if _is_sse_comment(chunk): + released.set() + return out + + out = asyncio.run(asyncio.wait_for(consume(), HANG_GUARD_SECONDS)) + assert out[0] == b"data: first\n\n" and out[-1] == b"data: last\n\n" + assert len(out) > 2 and all(_is_sse_comment(chunk) for chunk in out[1:-1]) + + +def test_a_client_that_leaves_during_the_silence_aborts_its_request(monkeypatch): + monkeypatch.setattr(api_server, "STREAM_KEEPALIVE_SECONDS", 0.01) + prefill_done = asyncio.Event() + + async def chunks(): + yield b"data: first\n\n" + await prefill_done.wait() + yield b"data: never read\n\n" + + async def scenario() -> tuple[list[bytes], list[int]]: + aborted = asyncio.Event() + seen: list[int] = [] + + async def abort_user(uid: int) -> None: + seen.append(uid) + aborted.set() + + request = _Request() + out: list[bytes] = [] + + async def consume() -> None: + manager = SimpleNamespace(abort_user=abort_user) + async for chunk in FrontendManager.stream_with_cancellation(manager, chunks(), request, 9): + out.append(chunk) + request.gone = True + + consumer = asyncio.ensure_future(consume()) + # Only the wrapper noticing the departure can end this wait; the guard below fails it. + await aborted.wait() + with pytest.raises(asyncio.CancelledError): + await consumer + return out, seen + + out, seen = asyncio.run(asyncio.wait_for(scenario(), HANG_GUARD_SECONDS)) + assert out == [b"data: first\n\n"] + assert seen == [9] + assert not prefill_done.is_set() + + +def test_a_stream_that_never_goes_quiet_carries_no_comments(monkeypatch): + monkeypatch.setattr(api_server, "STREAM_KEEPALIVE_SECONDS", HANG_GUARD_SECONDS) + + async def chunks(): + for index in range(5): + yield f"data: {index}\n\n".encode() + + async def consume() -> list[bytes]: + return [ + chunk + async for chunk in FrontendManager.stream_with_cancellation( + _manager(), chunks(), _Request(), 3 + ) + ] + + out = asyncio.run(asyncio.wait_for(consume(), HANG_GUARD_SECONDS)) + assert out == [f"data: {index}\n\n".encode() for index in range(5)]