From bac8f4d84c3753b8ec9048b3f55c6c1b40e8e702 Mon Sep 17 00:00:00 2001
From: Copilot <223556219+Copilot@users.noreply.github.com>
Date: Tue, 25 Aug 2026 14:06:38 +0000
Subject: [PATCH] test(server): acquire native session leases through registry

Load the native server fixture through AppState so the driver carries its production model key, then bind, acquire, generate, and close through SessionRegistry lease APIs. Assert generation and close both release their leases.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 crates/onnx-genai-server/src/tests.rs | 68 +++++++++++++++++++--------
 1 file changed, 49 insertions(+), 19 deletions(-)

diff --git a/crates/onnx-genai-server/src/tests.rs b/crates/onnx-genai-server/src/tests.rs
index 39ec75ab11..e6a9de6a04 100644
--- a/crates/onnx-genai-server/src/tests.rs
+++ b/crates/onnx-genai-server/src/tests.rs
@@ -2592,32 +2592,42 @@ async fn stalled_output_route_does_not_block_another_completion() {
 async fn native_driver_sessions_generate_through_server_path() {
     let model_dir = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
         .join("../../tests/fixtures/tiny-native-sub4-engine");
-    let (driver, ()) = EngineDriver::start(
-        move || {
-            Ok((
-                Engine::from_dir(
-                    &model_dir,
-                    EngineConfig {
-                        decode_backend: EngineDecodeBackend::Native,
-                        ..EngineConfig::default()
-                    },
-                )?,
-                (),
-            ))
+    let model_id = "tiny-native-sub4-engine";
+    let state = AppState::load_with_config(
+        &model_dir,
+        Some(model_id.to_string()),
+        ServerConfig {
+            max_queue_depth: 4,
+            engine_config: EngineConfig {
+                decode_backend: EngineDecodeBackend::Native,
+                ..EngineConfig::default()
+            },
+            ..ServerConfig::default()
         },
-        2,
-        4,
     )
     .unwrap();
-    let session_id = driver.create_session().await.unwrap();
+    let handle = handle_for(&state, model_id);
+    let placement = handle.engine.create_session().await.unwrap();
+    let client_id = "native-driver-session";
+    let binding = handle.engine.binding(placement);
+    assert!(
+        state
+            .sessions
+            .insert(client_id.to_string(), binding.clone())
+            .unwrap()
+            .is_none()
+    );
+    let lease = state.sessions.acquire(binding, client_id).unwrap();
+    assert_eq!(leases_held(&state), 1, "generation holds the session lease");
 
     let mut request =
         onnx_genai::GenerateRequest::new(onnx_genai::GeneratePrompt::TokenIds(vec![0]));
     request.options.max_new_tokens = 2;
     request.options.temperature = 0.0;
     request.options.stop_on_eos = false;
-    let generation = driver
-        .generate(Some(session_id), request, None)
+    let generation = handle
+        .engine
+        .generate(Some(lease), request, None)
         .await
         .unwrap();
     let result = timeout(
@@ -2629,8 +2639,28 @@ async fn native_driver_sessions_generate_through_server_path() {
     .expect("native session generation failed");
 
     assert_eq!(result.token_ids, vec![1, 1]);
-    assert_eq!(driver.session_token_count(session_id).await.unwrap(), 3);
-    driver.close_session(session_id).await.unwrap();
+    assert_eq!(
+        leases_held(&state),
+        0,
+        "completed generation releases the session lease"
+    );
+    assert_eq!(
+        handle.engine.session_token_count(placement).await.unwrap(),
+        3
+    );
+
+    let lease = state.sessions.take_for_close(client_id).unwrap();
+    assert_eq!(leases_held(&state), 1, "close holds the session lease");
+    handle.engine.close_session(lease).await.unwrap();
+    assert!(
+        state.sessions.get(client_id).unwrap().is_none(),
+        "closed session is no longer routed"
+    );
+    assert_eq!(
+        leases_held(&state),
+        0,
+        "completed close releases the session lease"
+    );
 }
 
 // ── KV cache dtype CLI/env surface tests ─────────────────────────────────────