fix(server): give the native-backend test the session lease the driver now requires

[justinchuby]

Rust quality is a required check and it is red on main. Every open PR in the repo inherits it. This restores it.

What broke

#2056 changed EngineDriver::generate to take Option<SessionLeaseGuard> and close_session to take SessionLeaseGuard. One call site was not updated:

error[E0308]: mismatched types
  --> crates/onnx-genai-server/src/tests.rs:2605:24
   |                   ---- ^^^^^^^^^^ expected `SessionLeaseGuard`, found `SessionPlacement`
  --> crates/onnx-genai-server/src/tests.rs:2618:26
   |            ------------- ^^^^^^^^^^ expected `SessionLeaseGuard`, found `SessionPlacement`

How it reached main — the gate worked and was overridden

I want to state this precisely rather than call it a CI gap, because it wasn't one.

• Rust quality ran on feat(server): refuse overlapping turns on one session at the routing layer #2056's own head as run 32837618182, job 97771768920, and failed at 11:06:27Z with the two errors above — byte-identical to the failure on main now.
• feat(server): refuse overlapping turns on one session at the routing layer #2056 merged at 12:03:16Z (29668c96), ~57 minutes later, with that required check red.

So no semantic conflict, no two-PRs-green-independently story: the check named the exact file and lines, and the merge proceeded past it.

The contributing factor worth naming is why that red was easy to discount: the site is behind #[cfg(feature = "native-backend")], so a default cargo test never compiles it and the failure looks like somebody else's inherited red. Worth knowing that this same call site was repaired once before — #1944, for generate's third argument. Second decay, same cause.

The fix

generate gets a real lease. For the close path I first wrote the obvious thing and it was wrong, so the comment in the diff records why:

collect_generation_result returns on the DriverEvent::Finished event (routes/completions.rs:925). The lease lives in DriverRoute._lease, and the route is dropped by the driver thread after it sends that event (driver.rs:1489). A returned result therefore implies the release is imminent, not that it has happened — a bare re-acquire races the driver thread. Bounded polling instead, matching the existing pattern at tests.rs:5449.

Falsified in both directions

Not merely green:

state	result
fix as committed	1 passed; 0 failed in 0.15s
mutation: hold the lease across the poll	FAILED — the finished turn never released its lease, 5.37s

The 5.37s matters: it is the full 200 x 25ms budget, so the loop demonstrably polls rather than succeeding on its first attempt and reporting a property it never tested.

Commands run

cargo clippy --locked --all-targets -p onnx-genai-engine -p onnx-genai-server \
  --features onnx-genai-engine/native-backend,onnx-genai-server/native-backend -- -D warnings   # clean (CI's exact command)
cargo test --locked -p onnx-genai-server --features native-backend --lib \
  native_driver_sessions_generate_through_server_path                                            # 1 passed
cargo fmt --all -- --check                                                                       # clean

The test executes here rather than only compiling — tests/fixtures/tiny-native-sub4-engine is present in-tree.

Scope

One test function. No production code, no API change. All runs bounded to taskset -c 24-31 under scripts/hostlock.sh; these are compile/test runs, not measurements, and I make no claim about host quietness.

[justinchuby]

Review: APPROVE from me. The poll is not defensive padding — it is the only version of this test that is deterministic, and I verified the ordering it claims.

Your comment says the lease release is imminent rather than done when the result arrives. That is exactly right, and the mechanism is one line further than the comment goes, so recording it here — driver.rs:1572, ContinuousBatchEvent::Finished:

if let Some(mut route) = routes.remove(&handle.id) {
    route.metrics.result(...);
    let _ = deliver_driver_event(&route.events, DriverEvent::Finished(result), DELIVERY_GRACE);
}   // <-- `route` drops HERE, and `_lease` (driver.rs:273) with it

routes.remove unbinds the row, but the guard is owned by the local route binding and drops at the end of the arm — after Finished is already in the consumer's channel. So the consumer can be running acquire while the driver thread still holds the guard. Not merely unordered: the send is ordered before the drop, which biases the window toward the failing side rather than away from it. A single-shot acquire here is a race that loses more often than intuition suggests, and it would lose on a busy runner in someone else's lane.

The bounded poll keeps the assertion's force — it still fails, with "the finished turn never released its lease", if the release never happens; it only declines to race the driver when it does. 200 × 25ms = 5s ceiling, comfortably inside the 5s timeout the test already uses elsewhere and far above any plausible drop latency.

I have posted the same finding on #2114, which fixes the same defect with a single-shot acquire and a comment asserting the guarantee that does not hold. One of the two should close; I have no preference on style, only that main does not gain an intermittent. This one is correct as written.

Two notes, neither blocking:

• session_id still names what is now a SessionPlacement. Cosmetic, and main is red while this waits, so not worth a push.
• The error arm discards the conflict (Err(_) => sleep). Fine — the loop's failure message already says the only thing a reader needs.

Context for anyone arriving here: the defect is #2056's, which merged at 12:03:16Z with Rust quality — a required check — already concluded failure at 11:06:31Z, 56m45s earlier, naming this exact file, line, and type mismatch. Because that lane is required and CI runs on the merge result, every open PR is currently red through no fault of its own; that is how I found this. Full timeline on #2056. main has carried it for 2h+.

Verified statically only: I did not compile -p onnx-genai-server --features native-backend locally (~39G of target/ on a shared box), and CI's Check the native backend compiles covers precisely that. The drop ordering above is read from origin/main, not inferred from the test.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 80.70%. Comparing base (74e6cb6) to head (7a7a76f).
⚠️ Report is 18 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2115      +/-   ##
==========================================
+ Coverage   80.29%   80.70%   +0.40%     
==========================================
  Files         426      429       +3     
  Lines      205244   215150    +9906     
  Branches   205244   215150    +9906     
==========================================
+ Hits       164801   173628    +8827     
- Misses      34792    35748     +956     
- Partials     5651     5774     +123     

Flag	Coverage Δ
cli-ort-linux	72.51% <ø> (?)
cli-ort-windows	72.10% <ø> (+0.09%)	⬆️
mlas	85.80% <ø> (?)
offline	80.82% <ø> (+0.31%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.
see 64 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[justinchuby]

Correction to my review, with measurements. The race is real but I described the wrong widener, and I could not reproduce it in 120 runs.

I said the window "widens on a loaded 2-core runner". I went and measured instead of leaving that as an assertion, and it does not hold. Recording it because #2114 was closed partly on my say-so.

What I ran. Built the test CI never runs (see below), then reverted the poll to #2114's single-shot acquire in my worktree and ran the test binary directly:

condition	result
taskset -c 0-15, 60 iterations, single-shot acquire	60 pass / 0 fail
taskset -c 4 (one CPU), 60 iterations, single-shot acquire	60 pass / 0 fail
current main (#2115's poll), same test	passes, 0.19s
non-vacuity control: pre-fix tests.rs, exact CI clippy step	exit 101, 2 × error[E0308] — so the check that says "fixed" is not vacuous

Why my mechanism was wrong. Single-CPU binding makes failure less likely, not more: the driver thread runs to the end of the if let block — dropping the guard — before the consumer is scheduled at all. The window needs genuine parallelism and a preemption landing in the few instructions between deliver_driver_event returning and route dropping. That is not something runner load reliably opens.

The real widener is a full output channel, and it is much bigger than I claimed. deliver_driver_event (driver.rs:1745) is synchronous and, when try_send hits a full channel, retries in a loop with thread::sleep until DELIVERY_GRACE:

let mut pending = match events.try_send(event) { Ok(()) => return Ok(()), ... };
let deadline = Instant::now() + grace;
loop {
    if Instant::now() >= deadline { return Err(DriverDeliveryError::Stalled); }
    thread::sleep(DELIVERY_RETRY_INTERVAL);
    ...
}

The route binding — and its _lease — stays alive across that entire call. So with a consumer that is not draining, the driver holds the turn's lease for up to DELIVERY_GRACE after Finished is queued, not for a few instructions. DRIVER_OUTPUT_BUFFER is 16; this test emits 2 tokens and drains promptly, which is exactly why 120 attempts could not open it.

So the corrected verdict: the ordering hazard I described is real and the poll is still the right code — but the reachability argument I gave for it was wrong, and a reader of my review would have gone looking for the flake on a busy runner and not found it. A test with a slower consumer or a longer generation is where a single-shot acquire would actually bite. #2115 is correct for a reason slightly different from the one I gave.

The finding that outlasts this. While setting the experiment up I checked what CI does with this test:

$ grep -n "native-backend" .github/workflows/*.yml
ci.yml:620   cargo clippy ... -p onnx-genai-engine -p onnx-genai-server --features .../native-backend   # compiles
ci.yml:1053  cargo test --locked -p onnx-genai-engine --features native-backend                          # runs, engine only

No lane runs cargo test -p onnx-genai-server with native-backend. native_driver_sessions_generate_through_server_path is compiled by exactly one step and executed by none. That is why #2056's breakage surfaced as a compile error rather than a test failure, and it is why neither #2114's nor #2115's runtime behaviour was checked by anything before it merged — including the poll loop, whose whole purpose is a runtime property. I ran it locally (passes, 0.19s, tiny fixture, no ORT) precisely because nothing else was going to.

Same family as #2058: a code path exactly one step can see. I will open this separately rather than bury it in a merged PR thread.

Everything above was run under scripts/hostlock.sh with taskset outermost, on a warm cache; the box was released between runs.