fix(server): compile the native-backend session test against the lease API

[justinchuby]

Main is red on the required check Rust quality — step "Check the native backend compiles" — and has been since the session-lease change landed. Every open PR fails it, including mine; that is how I found it.

What broke

native_driver_sessions_generate_through_server_path (crates/onnx-genai-server/src/tests.rs) still passes a SessionPlacement to generate and close_session, which now take a SessionLeaseGuard (c8732ee69, refined by 5b5eec2ca). Reproduced locally with the exact CI command:

error[E0308]: mismatched types
    --> crates/onnx-genai-server/src/tests.rs:2618:26
2618 |     driver.close_session(session_id).await.unwrap();
     |            ------------- ^^^^^^^^^^ expected `SessionLeaseGuard`, found `SessionPlacement`

Why it landed green

The test is #[cfg(feature = "native-backend")]. No default-feature lane compiles it, so the signature change was green everywhere that ran, and the one job that does compile it is the one now failing on main for everybody.

This is the same failure mode we have spent the week catching in benchmarks — the arm was not on the route we named — except in a test, where it is worse: a benchmark on the wrong route gives you a number you will eventually distrust; a test that is not compiled on the leg you are citing gives you a green check that you will not. I am not touching CI scope here; that is the CI lane's call, and the gate did catch this, just after the merge rather than before it.

The fix

The way the other driver tests do it (driver.rs:2224, :2334): acquire a lease from SessionLeases, hand it to generate — which consumes it and releases it when the turn ends — and take a second lease to close with.

The re-acquire is not ceremony. It fails if a finished turn ever leaves a session leased, which is the property the lease API exists to provide, so the repaired test now asserts slightly more than the one it replaces.

Validation

• 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 (the exact CI step, ci.yml:613).
• cargo test -p onnx-genai-server --features native-backend --lib native_driver_sessions — 1 passed, and it really ran (test tests::native_driver_sessions_generate_through_server_path ... ok), which is the check this PR is about.

Test-only change; no production code touched. Merging normally via merge_when_green.sh once required checks are green — no --admin, no ruleset bypass.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 80.47%. Comparing base (6caef22) to head (69ec641).
⚠️ Report is 40 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2114      +/-   ##
==========================================
+ Coverage   80.30%   80.47%   +0.17%     
==========================================
  Files         426      413      -13     
  Lines      205244   195882    -9362     
  Branches   205244   195882    -9362     
==========================================
- Hits       164811   157631    -7180     
+ Misses      34781    32721    -2060     
+ Partials     5652     5530     -122     

Flag	Coverage Δ
cli-ort-windows	?
offline	80.47% <ø> (-0.05%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.
see 65 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]

Review: this fix compiles, but the second acquire is a race, and main will see it as a flake. REQUEST_CHANGES — or close in favour of #2115, which already handles it.

Both this and #2115 fix the same defect (#2056 changed close_session to take a SessionLeaseGuard and missed the call site behind #[cfg(feature = "native-backend")]). The difference is one line of ordering, and it decides whether the test is deterministic.

let lease = leases
    .acquire(driver.binding(session), "sess-native")
    .expect("the finished turn released its lease");

The comment on that .expect states a guarantee the driver does not provide. Traced on origin/main:

1. The turn's guard is moved into the command and comes to rest in DriverRoute._lease — "Held for as long as the row is" (driver.rs:273).
2. The row is retired in the ContinuousBatchEvent::Finished arm (driver.rs:1572):

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` with it

routes.remove unbinds the row from the map, but the guard is owned by the local route and is not dropped until the arm ends — which is after Finished has been pushed into the consumer's channel. So the consumer can observe Finished, return from collect_generation_result, and call acquire while the driver thread is still inside that arm holding the guard.

Receiving the result therefore implies the release is imminent, not done. There is no happens-before edge between the two, and the window is exactly the kind that widens on a loaded 2-core runner — so this fails rarely, on someone else's PR, in a lane nobody suspects. That is the expensive failure mode, not the loud one.

The first acquire (before generate) is fine: nothing else holds the session at that point.

Two ways forward, either is fine by me:

• Close this in favour of fix(server): give the native-backend test the session lease the driver now requires #2115, which polls acquire with a bounded retry and documents precisely this ordering. It still fails if the release never happens — it only declines to race the driver when it does — so it keeps the assertion's force.
• Or keep this one and make the second acquisition bounded-retry rather than single-shot. The assertion you want is "the turn's lease is released", and a poll asserts that; a single-shot acquire asserts "it was already released by the time this line ran", which is a stronger claim than the code makes and not the one the test is about.

I have no preference between the two authors' styles — only that main does not gain an intermittent. If this one lands with the retry, I will say the same thing on #2115 and one of them should close.

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

[justinchuby]

Caught, and the finding is the one that mattered: the commit fixed the clippy step and broke the formatting step of the same required job, so Rust quality would have stayed red and I would have "fixed" main into a different failure. Renaming session_id → lease pulled the call back under 100 columns and rustfmt no longer wanted it wrapped. cargo fmt --all -- --check is now clean and the amended head is a4d55df35.

That is a sharper version of the same lesson this PR is about: I validated the step I knew was failing and did not re-run the whole gate. Verifying the specific thing you broke is not the same as verifying the gate, exactly as compiling on the default feature set is not the same as compiling the test you are citing.

On the flake question — thank you for tracing it to the code rather than reasoning about it. drop(lease) at driver.rs:1842 is sequenced before deliver_event(.., Finished(result)) at :1846 on the worker thread, SessionLeaseGuard::drop releases synchronously into the same Arc<SessionLeases> the test holds (lease.rs:246-250 → :190-192), and collect_generation_result returns only on Finished. So the re-acquire is ordered behind the release by the mpsc, not by timing. Good: that means the .expect("the finished turn released its lease") is a real assertion about the documented ordering (driver.rs:1836-1841) rather than a hopeful one, which is why I would like to keep it rather than drop the lease silently.

Also useful that you checked priority 5 categorically — the other cfg-gated tests in that file are metrics (a default feature, so already compiled by the default lane) and two that only call parse_native_device. Fixing one and leaving a second offender would have left main red for the next person and made this PR look like the fix.

Re-verified after the amend: clippy -D warnings on the exact CI command clean, cargo fmt --all -- --check clean, cargo test -p onnx-genai-server --features native-backend --lib native_driver_sessions → 1 passed and named in the output. merge_when_green.sh is watching; no --admin, no ruleset bypass.

[justinchuby]

Superseded by #2115, which waits for DriverRoute to release the turn lease before reacquiring it and has passed native-backend validation.

[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.