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.

[Artemowka22]

Confirmed, and thank you — the unwind analysis is exactly right and your probe caught what our disconnect tests could not: the CancelledError lands on wait_for_ack's innermost await, its finally pops both maps on the way out, and by the time stream_with_cancellation calls abort_user the event_map.pop(uid) claim is already False, so the AbortMsg silently never goes out (and no log line, since the cancellation interrupts the async for, not the is_disconnected branch). We had even hit the same trap from the other side — the non-streaming helper deliberately aborts before cancelling its drain task for exactly this reason — without realizing the streaming path had no way to win that ordering.

Folded into the PR (new commit):

• The claim change you suggested: abort idempotency is keyed on an insertion-ordered aborted_uids dict, FIFO-bounded to 65,536 like the scheduler's own abort tombstones, instead of on the maps. A duplicate or late AbortMsg is safe on your evidence and ours — the scheduler acks aborts for uids it no longer has.
• Delivery from finally: stream_with_cancellation now delivers the shielded abort whenever the stream did not run to completion — one path for server-side cancellation, a late body close (GeneratorExit, which the old except CancelledError never saw), and generator errors. Your stream.close() variant is the GeneratorExit case.
• This also un-breaks the accounting drain's abort barrier for requests whose maps were already emptied.

Regression tests added: cancellation delivered while the drain generator is suspended in wait_for_ack with both maps already popped (fails on the previous commit, passes now), aclose() mid-stream, and two concurrent aborts collapsing to one AbortMsg.

The request.receive() watcher from your third commit we left out for now, deliberately: your own measurements show the claim + finally changes cover all three disconnect variants once the server delivers either the cancellation or the close, and the watcher is a second signal for stacks that deliver neither promptly. If the maintainers want it in this PR rather than as a follow-up, happy to add it — your branch clearly proves it out.

Validation on our side is CPU-only (macOS: the server suite, 544 passed, plus the full-tree run diffed against the clean base — identical failure lists, all environmental). If you get a chance to point disconnect_probe.py at a build with the new commit, that would be the real proof on the stack that matters.