fix(server): guarantee AbortMsg delivery on stream cancellation

[Artemowka22]

Summary

FrontendManager.stream_with_cancellation reacts to a client disconnect with a bare asyncio.create_task(self.abort_user(uid)) and re-raises. The task has no owner: request teardown can complete before the abort coroutine ever runs, and an exception inside it degrades to a "Task exception was never retrieved" warning. The scheduler then keeps decoding for a client that is gone — under a long max_tokens this wastes the GPU for minutes and pins KV pages.

Fix (ported from a downstream audit — agisota/freetoken-mlx, docs/AUDIT.md, defect 2):

• await the abort inline and shielded in the cancellation handler, so delivery completes even as cancellation propagates, then re-raise;
• make abort_user idempotent — exactly one AbortMsg per uid even if the handler runs twice (disconnect + server-side cancel).

Test plan

• pytest tests/server/test_stream_cancellation.py (new): cancellation delivers exactly one AbortMsg before teardown finishes; double cancellation stays single-shot; normal completion sends no abort
• full tests/server/ suite passes

[benwilson]

Supporting evidence, and a case this PR may not cover: the non-streaming path also keeps generating after the client is gone.

Environment: FreeToken git 4b94bdc (HEAD 2026-08-30, no local patches), installed with pip install --no-build-isolation "freetoken[accel] @ git+…@4b94bdc"; torch 2.11.0+cu130, CUDA toolkit 13.0, driver 595.71.05; RTX 3090 24 GB (sm_86) on a rented vast.ai host (503 GiB RAM); image pytorch/pytorch:2.11.0-cuda13.0-cudnn9-devel; checkpoint dealignai/Qwen3.8-Flash-Next-ABLITERATED-NVFP4 @ be794b99 (standard compressed-tensors NVFP4 layout, byte-identical quant config to the RadixArk stock checkpoint — not the FTW repack). (--max-running-requests 1)

Repro: a non-stream POST /v1/chat/completions with max_tokens: 8192 and thinking on; SIGTERM the client at ~6,720 generated tokens. The engine log kept printing Decode batch, #running-req: 1, … 7040 … 7104 … for the next ~70 s until the cap, with a second client's request queued behind it the whole time (#queue-req: 1, its first token came 61 s late). With one running slot an abandoned request is effectively a 1-slot outage for its remaining max_tokens. If stream_with_cancellation is the only place a disconnect is observed, the plain-response handler needs the same request.is_disconnected() poll (or Starlette's http.disconnect receive) feeding abort_user.

[Artemowka22]

Thanks for the measured repro — that is the other half of the bug, and with --max-running-requests 1 your 61 s queued-behind-a-ghost number makes the cost very concrete.

The PR now covers it: the non-streaming handlers run the generation drain as a task and poll request.is_disconnected() once a second (_await_watching_disconnect, the non-streaming twin of stream_with_cancellation). On disconnect 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). Covers /v1/chat/completions and each prompt of a non-streaming /v1/completions batch; handler cancellation on server shutdown delivers the abort too, mirroring the streaming path.

Tests cover the disconnect-abort for both endpoints and the no-abort happy path.

[gdevenyi]

Tried on 2 x RTX 6000 Ada (sm_89) serving Qwen3.8-Flash-Next (RadixArk NVFP4) at TP=2, offload backend, merged onto my deploy branch with ten other open PRs, tests run on the box, then put in production.

Merged clean, the three unit tests pass, and a requests stream closed early was aborted within 3 s. But a probe that abandons a streaming request three ways (disconnect_probe.py: raw socket close after a few SSE chunks, requests close, openai-python stream.close()) showed the request still decoding to max_tokens in the other two cases, on a build with this PR (#running-req: 1 in the scheduler log for 45 s, 4,800 tokens), and the pattern flipped between runs. The TCP connection was gone in every case.

Two things this PR's path cannot see on the installed stack (uvicorn 0.52.4, Starlette 1.6.0):

• under ASGI spec 2.4 Starlette does not run listen_for_disconnect while a StreamingResponse streams, and uvicorn's h11 send() returns silently once disconnected is set, so no CancelledError ever reaches stream_with_cancellation; the generator is left suspended at its yield and finalized whenever garbage collection gets to it (that is the flip between runs);
• is_disconnected()'s zero-timeout poll misses the message as well.

Those two turned out to be red herrings on this stack (a minimal uvicorn + Starlette app delivers the disconnect fine, with or without a BaseHTTPMiddleware). The actual cause is in the frontend: on a disconnect the CancelledError unwinds from the innermost await outwards, and that innermost frame is wait_for_ack, whose finally pops ack_map and event_map for the uid. By the time stream_with_cancellation calls abort_user, the claim self.event_map.pop(uid, None) is not None is False, so the AbortMsg is never sent -- and there is no log line either, because the cancellation interrupts the async for, not the is_disconnected() branch. This PR keeps that claim, so it changes nothing for this path.

What fixed it on my branch (three commits, regression tests added): key the abort's idempotency on a set of aborted uids instead of the maps (the scheduler acks an abort for a uid it no longer has), deliver the abort from finally whenever the stream did not finish (also covers a GeneratorExit from a late body close), and a watcher task on request.receive() as a second signal. Verified on production with ignore_eos so the request would otherwise run to max_tokens: all three variants now abort within 0.5 s, each with an Aborting request for user N line in the log. Suggest folding the claim change into this PR: it is the one line that matters.