fix(ArcInfer/ArcGraph): cross-request token leak — parked burst tokens no longer outlive their sequence

[heydryft]

Stacked on #217 (arcgraph/device-decode-loop-rel); merge after it (D20). Until #217 merges, this PR shows its commits too — the leak fix itself is 136636973, four files.

The bug

device_decode_burst runs every burst with eos_token_id = -1 (normal.rs:2193), so the device loop always parks its full burst in the thread-local pending queue. The engine drains one token per step. A sequence that stops mid-burst — stop string, max_tokens, client cancel, forward error, panic, eviction — left its undrained tail parked, and the next sequence's first sample_sequence (including the one right after its prefill) consumed those tokens as its own: one user's tokens inside another user's response. Worse, with anything parked, device_decode_burst's pending > 0 short-circuit (normal.rs:2168) returns stale graph logits without launching anything — for a new request's prefill step too.

clear_pending_tokens documented the completion contract, but its only production call sites were fallback paths (normal.rs decline/fall-back arms, sampling.rs ineligible arm, kill_device_loop); DeviceDecodeLoop::reset() — the other caller — has zero production call sites. The existing unit test covered the function, not that anything calls it.

The fix: one funnel, not N scattered calls

Every path that takes a sequence out of the running set passes through Sequence::set_state. That is where the stand-down now lives — an exhaustive match, so a future state variant forces a decision:

completion path	mechanism	had leak?	covered by
EOS / stop-token / stop-string / length / model-length / tool-calls / low-confidence	finish_or_add_toks_to_seq -> set_state(Done(reason)) (sampling.rs:267/:277, :74)	yes	funnel
streaming send-failure -> cancel	set_state(Done(Canceled)) (sampling.rs:246)	yes	funnel
client disconnect sweep	default_scheduler.rs:500, paged scheduler.rs:449 -> Done(Canceled)	yes	funnel
TERMINATE_ALL_NEXT_STEP	default_scheduler.rs:556 -> Done(Canceled)	yes	funnel
forward error / panic (step_catching_panics -> handle_pipeline_forward_error!)	utils/mod.rs:403 -> set_state(Error); also handle_seq_error_stateaware_ok! (utils/mod.rs:272)	yes	funnel
PagedAttention eviction / ignore	scheduler.rs:288/:299 -> FinishedIgnored; engine mod.rs:839	yes	funnel
preemption / swap	paged scheduler.rs:532 -> Waiting; Swapped	yes (self-leak: stale tokens after re-prefill)	funnel, fail-closed
DefaultScheduler bucketing waitlist	moves a running seq aside "without a state modification" (default_scheduler.rs:263)	leak possible only with >=2 sequences mid-transition	NOT coverable by any completion funnel — reported below

Why the dependency became non-optional

The funnel could have been #[cfg(feature = "cuda")] — and then it could not be compiled, run, or mutation-tested on any machine that reviews it, which is the exact class that let this ship (and the recorded polarity-bug pattern). arc-cuda-graph's host surface (pending queue, burst driver, admission) is plain Rust already exercised by the CPU clippy lane, and it adds zero new transitive deps to mistralrs-core. So: dep promoted to non-optional (arc-cuda-graph/cuda still rides the cuda feature), and device_loop_pre_sampled_token loses its cfg split — the real take path now compiles and runs everywhere. On non-CUDA builds nothing ever parks, so behaviour is unchanged.

The test that fails on the leak — mutation-checked both ways

pipeline/sampling.rs::device_loop_cross_sequence_leak_tests drives the real sample_sequence: A (greedy, eligible) samples once, its burst parks [21, 22, 23, 24] (eos -1 => full length), A drains 21 — a built-in control proving the take path is live in the test build, so the assertions cannot be vacuous — then A stops mid-burst via Done(Length(2)), and B samples with its own logits peaking at 9.

• Wired (this PR): queue at B's first sample = 0; B's first token = 9 (its own argmax). 2 passed.
• Mutated (funnel arm removed from set_state): queue assertion fails left: 3, right: 0 — A's undrained tail survives; with that assertion skipped to let execution reach B: "cross-request token leak: sequence B's first sampled token is 22, one of sequence A's parked burst tokens." The second test fails the same way for Done(Canceled) (left: 4, right: 0) and every other finished transition.

A second test pins the whole funnel map: Done(Canceled), Done(StopString), Error, FinishedAborted, FinishedIgnored, Waiting, Swapped each clear the queue.

Same-shape audit (report, not fixed here)

1. DeviceDecodeLoop per-sequence device state outlives completion. Ring cursors, ring head, the device position buffer the commit kernel advances, and the last committed id in the pinned graph input buffer all live in the pipeline-lifetime loop_slot; DeviceDecodeLoop::reset() exists for exactly this and has no production caller. A second sequence engaging the loop would decode on-device as a continuation of the first. Correctly resetting requires the new sequence's prefill to re-prime graph-mode KV/positions — capture-lane work in normal.rs, outside this PR's fence.
2. Mid-flight interleave (batch grows 1->2, or the bucketing waitlist above): sequence A alive with parked tokens, B samples first — B is eligible and takes A's token; no completion transition ever fires. Fix design: owner-tag the queue (pre-hook publishes seq.id() with eligibility; push_pending_tokens stamps the owner; take requires a match, mismatch clears fail-closed). Contained to arc-cuda-graph + sampling.rs; not shipped here to keep this PR single-mechanism and its mutation semantics clean.
3. pending > 0 short-circuit lacks a seq_len == 1 guard (normal.rs:2168): any stale pending token converts a prefill forward into a no-op returning stale decode logits. Unreachable across sequences with the funnel in place; cheap armor for the normal.rs owner.
4. Eligibility flag: cleared by the same stand_down. Sampler RNG: none per-sequence (greedy, fixed seed). DEVICE_LOOP_KILLED: process-wide by design.

Premise corrections against the tree

• ARC_GRAPH_DEVICE_LOOP is not registered in capability_reachability.rs — no entry exists to keep or update; the crate is covered only by the dead-symbol ratchet (which passes).
• The stream-binding change for capture is already implemented as an opt-in, not merely "visible from source": ARC_CAPTURE_STREAM=1 -> mistralrs_for_server_builder.rs:1238 init_device -> Device::new_cuda_with_stream(0) (a real capturable stream); default off via CAPTURE_STREAM_DEFAULT = false (:1182), auto-refused under NCCL or paged-attn. Named blast radius (in-source): NCCL conflict, PagedAttention device expectations (device_map::get_all_similar_devices), cudarc event tracking disabled, persistent cuBLAS workspace. Flipping the default is one line + its test, gated on a hardware-verified replay (ARC_GRAPH_VERIFY_REPLAYS). Note this entry point covers the server builder only; other entry points still bind the legacy NULL stream and graph.rs:143-167 logs ArcGraph INERT.

+1-step replay divergence (report-only; capture stays default-off)

Wiring is proven innocent (same-step delta = 0.000e0), so the divergence is state carried across the replay boundary. On-box probe to localize it: run with ARC_CAPTURE_STREAM=1 ARC_V4_CAPTURE_PROBE=1 ARC_CANDLE_ALLOC_CACHE=1 ARC_GRAPH_VERIFY_ALWAYS=1 and the layer bisect (arc_layer_bisect / arc_layer_trace_begin, already wired at normal.rs:2802): at step N (replay) and step N+1 (first divergent step), trace eager vs replay per-layer checksums — the first layer index whose delta goes nonzero at N+1 names the tensor whose address or value went stale across the boundary (per-step inputs: positions, KV write index, mask window, input id). The existing ARC_GRAPH_VERIFY_REPLAYS comparison already computes both forwards; extending it to dump per-layer sums at N and N+1 is the whole probe. D14: nothing here is validated until the box runs it.

Verification

• cargo check -p mistralrs-core (macOS, no CUDA): clean.
• cargo test -p mistralrs-core --lib: 695 passed. --test capability_reachability: 12 passed (ratchet clean). --test doc_citations: 11 passed. cargo test -p arc-cuda-graph: 76 passed.
• Scoped clippy lane (arc-bench arc-engine arc-cuda-graph arc-cli mistralrs-quant --tests --examples -D warnings): clean.
• Compile-unverified locally: the cuda-feature build of mistralrs-core with the reworked feature list (macOS cannot build it) — relies on the cargo check (cuda, workspace) and nvcc sm_80/sm_90 lanes. No .cu or cfg(cuda) code was modified.
• D14: leak behaviour is validated by host tests both ways; no GPU run occurred. The device loop itself remains opt-in default-off.

🤖 Generated with Claude Code

https://claude.ai/code/session_01WK8VgocBUrr5djjE1ZfCNf

[github-actions]

Code Metrics Report

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
 Language              Files        Lines         Code     Comments       Blanks
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
 C Header                  5          305          210           52           43
 CSS                       2         1181         1036           34          111
 CUDA                     81        29475        20161         6275         3039
 Dockerfile                1           39           22            8            9
 JavaScript               16         3546         2676          482          388
 Jinja2                    7          694          656            5           33
 JSON                     75         4896         4893            0            3
 Makefile                  1            6            5            0            1
 Metal Shading Lan|       33        12224         9431         1142         1651
 PowerShell                1          300          227           30           43
 Python                  147        15371        12677          824         1870
 Shell                    42        10304         6858         2769          677
 Plain Text                4         3801            0         2479         1322
 TOML                     33         1498         1294           54          150
 YAML                      3           25           23            2            0
─────────────────────────────────────────────────────────────────────────────────
 HTML                      4         2687         2604           43           40
 |- CSS                    2          543          479           37           27
 |- JavaScript             1         1233         1215           12            6
 (Total)                             4463         4298           92           73
─────────────────────────────────────────────────────────────────────────────────
 Jupyter Notebooks         4          122           83           23           16
 |- Markdown               1           60           30           22            8
 |- Python                 1          122          113            1            8
 (Total)                              304          226           46           32
─────────────────────────────────────────────────────────────────────────────────
 Markdown                213        48932            0        38081        10851
 |- BASH                  72         1655         1203          331          121
 |- C                      4           19           19            0            0
 |- CUDA                   2           84           56           16           12
 |- JSON                  19          779          779            0            0
 |- PowerShell             1            1            1            0            0
 |- Python                23         1008          787          113          108
 |- Rust                  68         2063         1727           78          258
 |- TOML                   6          207          164            0           43
 |- YAML                   5           41           36            5            0
 (Total)                            54789         4772        38624        11393
─────────────────────────────────────────────────────────────────────────────────
 Rust                    681       340988       293252        18096        29640
 |- Markdown             504        32562          471        28063         4028
 (Total)                           373550       293723        46159        33668
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
 Total                  1353       516771       363188        99077        54506
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

[heydryft]

Follow-up commit a30bec297 — owner-tagged queue, loop reset, prefill guard, registry entry

(Also: rebased onto release/openrouter-ready after #217's squash-merge — the PR diff is now exactly these two commits.)

1. Owner tag (structural, closes the funnel-invisible hole). The pending queue is stamped with the owning sequence's id at park time (published by device_loop_pre_sampled_token with eligibility); take_pending_token(seq_id) returns Taken | Empty | Foreign{dropped} and a foreign queue is dropped whole, logged with the box instrument line "ArcGraph: dropped N foreign parked tokens". This covers what no completion funnel can: the DefaultScheduler bucketing waitlist moves a running sequence aside "without a state modification" (its own doc), and batch 1→2 growth is the same shape — more common now that #219 landed. The funnel stays as the primary defence.

2. Forward-side guards (device_decode_burst, cfg(cuda) — compile-unverified locally, nvcc/cuda lanes gate it): the pending > 0 short-circuit now requires seq_len == 1 (a parked queue must never replace a prefill) and pending_owned_by_current(). And if the short-circuit fired but the scheduler swapped sequences before the sample, the aliased-logits marker makes that sample fail the step loudly (InternalError → SequenceState::Error → funnel) instead of host-sampling another sequence's stale graph output — a plausible-token silent corruption otherwise.

3. DeviceDecodeLoop::reset() wired — scoped deliberately: every stand_down (the funnel's included) bumps a thread-local device-state generation; run() compares and resets ring head, cursors and fault word before a newer-generation burst. Reset at engagement, not completion, because (a) the funnel lives in Sequence and cannot reach the pipeline-owned loop, and (b) between bursts is the only point where no ring write can still be in flight on the stream. Reset-on-every-completion at the funnel would be both unreachable and racy.

4. Registry entry added: ARC_GRAPH_DEVICE_LOOP is now in capability_reachability.rs (Symbol device_loop_enabled, parent ArcInfer / ArcGraph, Live) — closing the refuted premise from the first commit.

Interleave test, mutation-checked both ways

a_waitlisted_sequence_with_no_state_change_cannot_leak_into_the_next: A parks [21,22,23,24], drains 21 (liveness control), then is waitlisted with no set_state; B samples.

• Wired: B's first token = 9 (its own argmax), foreign tail dropped whole (queue = 0), A resumes on its own logits. 4/4 tests pass.
• Owner comparison removed: "cross-request token leak through the waitlist interleave: B's first sampled token is 22, one of A's parked burst tokens."

Plus an_aliased_step_served_for_another_sequence_fails_instead_of_sampling_stale_logits (double-fault path errors; owner's own marker path takes its token normally) and 4 new arc-cuda-graph unit tests (foreign drop, ownership tracking, one-shot marker, generation bump).

Verification

cargo check clean (core + arc-cuda-graph, macOS no-CUDA); mistralrs-core lib 712 passed; arc-cuda-graph 80 passed; capability_reachability 11 passed (new entry live, dead-symbol ratchet clean); doc_citations 12 passed; scoped clippy lane clean. Compile-unverified locally: the normal.rs cfg(cuda) arm and the DeviceDecodeLoop generation check — rely on the cuda/nvcc CI lanes. D14: not validated on hardware; the gate remains opt-in default-off.

Still open for the enable-on-box decision: DeviceDecodeLoop reset covers ring/cursors/fault; the pinned positions/input-ids buffers are host-rewritten each step by the capture lane, but that interaction (re-engagement after a fallback mid-sequence) is only provable on the box — cheapest probe is the existing ARC_GRAPH_VERIFY_ALWAYS=1 lane plus one two-request session with ARC_GRAPH_DEVICE_LOOP=1, grepping for "ArcGraph: dropped".