From 8856ae68f6a553a35b3f3b225a838833464ab995 Mon Sep 17 00:00:00 2001 From: mac Date: Wed, 26 Aug 2026 23:13:27 +0300 Subject: [PATCH 1/2] server: guarantee abort delivery on stream cancellation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit stream_with_cancellation reacted to a client disconnect with an unowned asyncio.create_task(abort_user(uid)): request teardown could outrun delivery and a failure inside the task degraded to a never-retrieved-exception warning, so the scheduler kept decoding for a client that was gone. Await the abort inline behind asyncio.shield (a second cancellation cannot kill the delivery task), and make abort_user claim the uid first — exactly one AbortMsg even if cancellation runs twice, and none at all when the stream already finished normally. Found via the freetoken-mlx downstream audit (docs/AUDIT.md, defect 2). --- python/freetoken/server/api_server.py | 13 ++- tests/server/test_stream_cancellation.py | 120 +++++++++++++++++++++++ 2 files changed, 128 insertions(+), 5 deletions(-) create mode 100644 tests/server/test_stream_cancellation.py diff --git a/python/freetoken/server/api_server.py b/python/freetoken/server/api_server.py index 3e2acc854..add3c6605 100644 --- a/python/freetoken/server/api_server.py +++ b/python/freetoken/server/api_server.py @@ -364,15 +364,18 @@ async def stream_with_cancellation(self, generator, request: Request, uid: int): raise asyncio.CancelledError yield chunk except asyncio.CancelledError: - asyncio.create_task(self.abort_user(uid)) + try: + await asyncio.shield(self.abort_user(uid)) + except Exception: # noqa: BLE001 + logger.exception("Failed to deliver abort for user %s", uid) raise async def abort_user(self, uid: int): + claimed = self.event_map.pop(uid, None) is not None + self.ack_map.pop(uid, None) + if not claimed: + return await asyncio.sleep(0.1) - if uid in self.ack_map: - del self.ack_map[uid] - if uid in self.event_map: - del self.event_map[uid] self.stats.on_abort(uid) logger.warning("Aborting request for user %s", uid) await self.send_one(AbortMsg(uid=uid)) diff --git a/tests/server/test_stream_cancellation.py b/tests/server/test_stream_cancellation.py new file mode 100644 index 000000000..120c50456 --- /dev/null +++ b/tests/server/test_stream_cancellation.py @@ -0,0 +1,120 @@ +"""Cancellation-path tests for FrontendManager.stream_with_cancellation / abort_user.""" + +from __future__ import annotations + +import asyncio +from types import SimpleNamespace + +import pytest + +from freetoken.message import AbortMsg +from freetoken.server.api_server import FrontendManager + + +class _Stats: + def __init__(self): + self.aborts = [] + + def on_abort(self, uid): + self.aborts.append(uid) + + +def _state(send_impl=None): + st = SimpleNamespace( + ack_map={7: [object()]}, + event_map={7: asyncio.Event()}, + stats=_Stats(), + ) + sent = [] + + async def default_send(msg): + sent.append(msg) + + st.send_one = send_impl or default_send + st.sent = sent + return st + + +class _Request: + def __init__(self, disconnected=False): + self._disconnected = disconnected + + async def is_disconnected(self): + return self._disconnected + + +async def _consume(state, request, uid=7): + state.abort_user = lambda request_uid: FrontendManager.abort_user(state, request_uid) + async for _ in FrontendManager.stream_with_cancellation(state, _never(), request, uid): + pass + + +async def _never(): + await asyncio.sleep(3600) + yield b"" # pragma: no cover + + +def test_cancellation_sends_one_abort_and_cleans_maps_inline(): + async def run(): + state = _state() + task = asyncio.create_task(_consume(state, _Request())) + await asyncio.sleep(0.01) + task.cancel() + + with pytest.raises(asyncio.CancelledError): + await task + + assert asyncio.all_tasks() - {asyncio.current_task()} == set() + assert len(state.sent) == 1 + assert isinstance(state.sent[0], AbortMsg) + assert state.sent[0].uid == 7 + assert state.ack_map == {} + assert state.event_map == {} + assert state.stats.aborts == [7] + + asyncio.run(run()) + + +def test_abort_delivery_failure_preserves_cancellation(): + async def boom(msg): + raise RuntimeError("zmq down") + + async def run(): + state = _state(send_impl=boom) + task = asyncio.create_task(_consume(state, _Request())) + await asyncio.sleep(0.01) + task.cancel() + + with pytest.raises(asyncio.CancelledError): + await task + + asyncio.run(run()) + + +def test_abort_user_is_idempotent(): + async def run(): + state = _state() + + await FrontendManager.abort_user(state, 7) + assert len(state.sent) == 1 + + await FrontendManager.abort_user(state, 7) + assert len(state.sent) == 1 + assert state.stats.aborts == [7] + + asyncio.run(run()) + + +def test_normal_completion_sends_no_abort(): + async def run(): + state = _state() + + async def one(): + yield b"data: x\n\n" + + async for _ in FrontendManager.stream_with_cancellation(state, one(), _Request(), 7): + pass + assert state.sent == [] + assert 7 in state.ack_map # wait_for_ack owns normal-path cleanup, not the stream + + asyncio.run(run()) From 274a2ce5212969a2d8b3a149bcbedf04d63c0814 Mon Sep 17 00:00:00 2001 From: mac Date: Tue, 1 Sep 2026 21:11:57 +0300 Subject: [PATCH 2/2] server: abort abandoned non-streaming requests on client disconnect Follow-up to the review evidence on #222 (benwilson): the non-streaming path kept generating after the client was gone -- with --max-running-requests 1 an abandoned request is a full outage for its remaining max_tokens (measured repro: ~70 s of dead decode, the next client's first token 61 s late). stream_with_cancellation was the only place a disconnect was observed. Give the plain handlers its non-streaming twin: _await_watching_disconnect() runs the generation drain as a task and polls request.is_disconnected() once a second; when the client goes away it delivers the same shielded abort_user (claim + AbortMsg first, so the drain task's own cleanup cannot swallow the claim), then winds the drain down and answers 499 (client closed request -- for the access log; the wire is dead). Handler cancellation (server shutdown) delivers the abort too, mirroring the streaming path. Covers /v1/chat/completions and each prompt of a non-streaming /v1/completions batch. Requests with request=None (adapter-internal callers) are unaffected. Co-Authored-By: Claude Fable 5 --- python/freetoken/server/openai_api.py | 104 ++++++++++++++++++++--- tests/server/test_stream_cancellation.py | 97 +++++++++++++++++++++ 2 files changed, 187 insertions(+), 14 deletions(-) diff --git a/python/freetoken/server/openai_api.py b/python/freetoken/server/openai_api.py index b4becd263..f05305484 100644 --- a/python/freetoken/server/openai_api.py +++ b/python/freetoken/server/openai_api.py @@ -2,6 +2,7 @@ import asyncio import json +import logging import time import uuid from collections.abc import AsyncIterator, Callable @@ -39,6 +40,8 @@ submit_generation, ) +logger = logging.getLogger(__name__) + #: The wire superset plus "off", DeepSeek's disable synonym that #: effort_toggle_kwargs has always honored. _ACCEPTED_EFFORTS = (*KNOWN_REASONING_EFFORTS, "off") @@ -144,6 +147,65 @@ async def v1_models(): )]) +# How often a non-streaming handler looks at the transport while the engine +# generates. The streaming path checks per chunk; one second bounds an abandoned +# request's extra decode work without measurable polling overhead. +_DISCONNECT_POLL_SECONDS = 1.0 + + +async def _await_watching_disconnect(awaitable, request: Request | None, state: Any, uid: int): + """Run an in-flight generation awaitable while watching the transport — the + non-streaming twin of stream_with_cancellation. Returns the awaitable's result, + or None when the client disconnected first: the same shielded abort_user is + delivered so an abandoned request stops burning decode slots instead of running + to max_tokens (with --max-running-requests 1 that is a full outage for its + remaining budget).""" + gen = asyncio.ensure_future(awaitable) + if request is None: + return await gen + try: + while True: + done, _ = await asyncio.wait({gen}, timeout=_DISCONNECT_POLL_SECONDS) + if done: + return gen.result() + if await request.is_disconnected(): + break + except asyncio.CancelledError: + # The handler itself was cancelled (server shutdown): deliver the abort, + # then let the cancellation propagate — mirrors stream_with_cancellation. + try: + await asyncio.shield(state.abort_user(uid)) + except Exception: # noqa: BLE001 + logger.exception("Failed to deliver abort for user %s", uid) + gen.cancel() + raise + # Client gone. Claim + AbortMsg first (abort_user is a no-op once the ack + # loop's own cleanup has run), then wind down the drain task; its result is + # undeliverable either way. + try: + await asyncio.shield(state.abort_user(uid)) + except Exception: # noqa: BLE001 + logger.exception("Failed to deliver abort for user %s", uid) + gen.cancel() + try: + await gen + except asyncio.CancelledError: + pass + except Exception: # noqa: BLE001 — the response is undeliverable, only cleanup matters + pass + return None + + +def _client_disconnected_response() -> JSONResponse: + # 499 — nginx's "client closed request". The client is gone; this status is + # for the access log, not the wire. + return create_error_response( + "client disconnected before the response was ready", + status_code=499, + err_type="client_disconnected", + ) + + async def handle_chat_completion( req: ChatCompletionRequest, request: Request | None, @@ -199,9 +261,13 @@ async def handle_chat_completion( return StreamingResponse(chunks, media_type="text/event-stream") try: - result = await generate_full(uid, spec, state, source="/v1/chat/completions") + result = await _await_watching_disconnect( + generate_full(uid, spec, state, source="/v1/chat/completions"), request, state, uid + ) except GenerationError as exc: return create_error_response(str(exc), code=exc.code) + if result is None: + return _client_disconnected_response() message: dict[str, Any] = {"role": "assistant", "content": result.content} if result.reasoning: message["reasoning_content"] = result.reasoning @@ -411,19 +477,29 @@ async def handle_completion( for index, prompt in enumerate(prompts): uid = state.new_user() await state.send_one(TokenizeMsg(uid=uid, text=prompt, sampling_params=_resolve_sampling(req, model_sampling))) - text = "" - finish_reason = "stop" - async for ack in state.wait_for_ack(uid): - if getattr(ack, "error", None): - return create_error_response(ack.error) - prompt_tokens += ack.prompt_tokens_delta - completion_tokens += ack.completion_tokens_delta - cached_tokens += ack.cached_tokens - text += ack.incremental_output - if ack.finished: - finish_reason = getattr(ack, "finish_reason", None) or "stop" - break - choices.append({"index": index, "text": text, "finish_reason": finish_reason, "logprobs": None}) + + async def _drain_one(uid: int = uid, index: int = index): + nonlocal prompt_tokens, completion_tokens, cached_tokens + text = "" + finish_reason = "stop" + async for ack in state.wait_for_ack(uid): + if getattr(ack, "error", None): + return create_error_response(ack.error) + prompt_tokens += ack.prompt_tokens_delta + completion_tokens += ack.completion_tokens_delta + cached_tokens += ack.cached_tokens + text += ack.incremental_output + if ack.finished: + finish_reason = getattr(ack, "finish_reason", None) or "stop" + break + return {"index": index, "text": text, "finish_reason": finish_reason, "logprobs": None} + + choice = await _await_watching_disconnect(_drain_one(), request, state, uid) + if choice is None: + return _client_disconnected_response() + if isinstance(choice, JSONResponse): + return choice + choices.append(choice) return { "id": f"cmpl-{uuid.uuid4().hex}", diff --git a/tests/server/test_stream_cancellation.py b/tests/server/test_stream_cancellation.py index 120c50456..ded9178bf 100644 --- a/tests/server/test_stream_cancellation.py +++ b/tests/server/test_stream_cancellation.py @@ -118,3 +118,100 @@ async def one(): assert 7 in state.ack_map # wait_for_ack owns normal-path cleanup, not the stream asyncio.run(run()) + + +# --------------------------------------------------------------------------- # +# Non-streaming handlers: a disconnected client must abort the engine request +# (same shielded abort_user as the streaming path) instead of letting it run to +# max_tokens while later requests queue behind it. +# --------------------------------------------------------------------------- # + +from freetoken.server import openai_api +from freetoken.server.api_models import ChatCompletionRequest, CompletionRequest + + +def _ack(text, finished=False): + return SimpleNamespace( + error=None, + incremental_output=text, + finished=finished, + prompt_tokens_delta=0, + completion_tokens_delta=1, + cached_tokens=0, + finish_reason="stop" if finished else None, + matched_stop=None, + logprobs=None, + ) + + +class _ApiState: + """State fake for the openai_api handlers: hangs forever when given no acks.""" + + def __init__(self, acks=None): + self.config = SimpleNamespace( + model_path="/models/unit-model", + served_model_name="unit-model", + tool_call_parser="llama3", + reasoning_parser=None, + ) + self.acks = acks + self.sent = [] + self.aborted = [] + + def new_user(self): + return 7 + + async def send_one(self, msg): + self.sent.append(msg) + + async def wait_for_ack(self, uid): + if self.acks is None: + await asyncio.sleep(3600) + for ack in self.acks or []: + yield ack + + async def abort_user(self, uid): + self.aborted.append(uid) + + +def _chat_req(**kwargs): + payload = {"model": "m", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 8} + payload.update(kwargs) + return ChatCompletionRequest(**payload) + + +def test_chat_non_stream_disconnect_delivers_abort(monkeypatch): + monkeypatch.setattr(openai_api, "_DISCONNECT_POLL_SECONDS", 0.01) + state = _ApiState(acks=None) # generation never finishes on its own + + resp = asyncio.run( + openai_api.handle_chat_completion(_chat_req(), _Request(disconnected=True), state, {}) + ) + + assert resp.status_code == 499 + assert state.aborted == [7] + + +def test_completion_non_stream_disconnect_delivers_abort(monkeypatch): + monkeypatch.setattr(openai_api, "_DISCONNECT_POLL_SECONDS", 0.01) + state = _ApiState(acks=None) + req = CompletionRequest(model="m", prompt="hello", max_tokens=8) + + resp = asyncio.run( + openai_api.handle_completion(req, _Request(disconnected=True), state, {}) + ) + + assert resp.status_code == 499 + assert state.aborted == [7] + + +def test_non_stream_connected_client_gets_result_without_abort(monkeypatch): + monkeypatch.setattr(openai_api, "_DISCONNECT_POLL_SECONDS", 0.01) + state = _ApiState(acks=[_ack("Hi", finished=True)]) + + result = asyncio.run( + openai_api.handle_chat_completion(_chat_req(), _Request(disconnected=False), state, {}) + ) + + assert result["choices"][0]["message"]["content"] == "Hi" + assert state.aborted == []