fix(server): guarantee AbortMsg delivery on stream cancellation

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))

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}",

tests/server/test_stream_cancellation.py

@@ -0,0 +1,217 @@
+"""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())
+
+
+# --------------------------------------------------------------------------- #
+# 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 == []