test(sidecar): cover clean TensorRT-LLM liveness shutdown

lib/sidecar/testkit/tests/sidecar_mocker_integration.rs

@@ -934,6 +934,45 @@ async fn trtllm_engine_lost_on_liveness_not_serving() {
     finish(&mut fixture, &engine).await;
 }
 
+#[tokio::test]
+async fn trtllm_clean_liveness_end_preserves_in_flight_generation() {
+    let control = Controller::default();
+    let mut fixture =
+        support::trtllm::Fixture::start(control.clone(), FixtureConfig::default()).await;
+    let engine = fixture.engine().await;
+    engine.start(0).await.unwrap();
+
+    let ctx = mock_context();
+    let handle = control.request(ctx.id(), after_token_responses(1, StreamAction::Continue));
+    let mut stream = engine
+        .generate(
+            request("mocker-model", vec![11, 22, 33], 3),
+            GenerateContext::new(ctx, None),
+        )
+        .await
+        .unwrap();
+    let mut outputs = checkpoint_outputs(&mut stream, &handle).await;
+    fixture.end_liveness_watch().await;
+    let error = bounded("engine lost on clean liveness EOF", engine.engine_lost()).await;
+    assert_eq!(
+        error.error_type(),
+        ErrorType::Backend(BackendError::EngineShutdown)
+    );
+    assert!(
+        error.to_string().contains("liveness stream ended"),
+        "{error}"
+    );
+
+    // Native graceful stop ends Watch before draining its active Generate RPCs.
+    assert!(poll!(stream.next()).is_pending());
+    handle.release();
+    outputs.extend(bounded("in-flight generation drains", stream.collect::<Outputs>()).await);
+    terminal(outputs, &handle.tokens(), 3, FinishReason::Length);
+    bounded("drained RPC released", handle.wait(Event::Dropped)).await;
+    fixture.scheduler_idle().await;
+    finish(&mut fixture, &engine).await;
+}
+
 fn http_request() -> PreprocessedRequest {
     let mut req = request("mocker-model", vec![11, 22, 33], 2);
     req.extra_args = Some(json!({"sglang_tito": {"sampling_params": {"max_new_tokens": 2}}}));

lib/sidecar/testkit/tests/support/trtllm.rs

@@ -9,6 +9,7 @@ use std::sync::{Arc, Mutex};
 use dynamo_backend_common::{BackendError, DisaggregationMode, PreprocessedRequest};
 use dynamo_llm::model_card::ModelDeploymentCard;
 use dynamo_mocker::common::protocols::EngineType;
+use dynamo_sidecar_testkit::bounded;
 use dynamo_sidecar_testkit::control::{Controller, Protocol, RequestHandle};
 use dynamo_sidecar_testkit::fixtures::Outputs;
 use dynamo_sidecar_testkit::server::TestServer;
@@ -35,6 +36,7 @@ pub struct Fixture {
     scripted: Arc<Mutex<HashMap<String, Vec<pb::GenerateResponse>>>>,
     pub discovery: Arc<Discovery>,
     health: tonic_health::server::HealthReporter,
+    liveness_started: Arc<tokio::sync::Notify>,
 }
 
 #[derive(Default)]
@@ -49,6 +51,11 @@ pub struct Discovery {
 }
 
 impl Fixture {
+    pub async fn end_liveness_watch(&mut self) {
+        bounded("liveness Watch opened", self.liveness_started.notified()).await;
+        self.health.clear_service_status("liveness").await;
+    }
+
     /// Sets what `Watch` reports for TensorRT-LLM's `liveness` health.
     pub async fn set_liveness_serving(&self, is_serving: bool) {
         let status = if is_serving {
@@ -134,13 +141,20 @@ impl SidecarFixture for Fixture {
             inner: service.clone(),
             discovery: discovery.clone(),
         };
-        let (health, health_service) = tonic_health::server::health_reporter();
+        let health = tonic_health::server::HealthReporter::new();
+        let liveness_started = Arc::new(tokio::sync::Notify::new());
+        let health_service = ControlledHealth {
+            inner: tonic_health::server::HealthService::from_health_reporter(health.clone()),
+            liveness_started: Arc::clone(&liveness_started),
+        };
         health
             .set_service_status("liveness", tonic_health::ServingStatus::Serving)
             .await;
         let server = TestServer::start(move |listener, shutdown| async move {
             tonic::transport::Server::builder()
-                .add_service(health_service)
+                .add_service(tonic_health::pb::health_server::HealthServer::new(
+                    health_service,
+                ))
                 .add_service(InferenceServer::new(controlled))
                 .add_service(ControlServer::new(control_service))
                 .serve_with_incoming_shutdown(TcpListenerStream::new(listener), async {
@@ -158,6 +172,7 @@ impl SidecarFixture for Fixture {
             scripted,
             discovery,
             health,
+            liveness_started,
         }
     }
 
@@ -189,6 +204,35 @@ impl SidecarFixture for Fixture {
     }
 }
 
+struct ControlledHealth {
+    inner: tonic_health::server::HealthService,
+    liveness_started: Arc<tokio::sync::Notify>,
+}
+
+#[tonic::async_trait]
+impl tonic_health::pb::health_server::Health for ControlledHealth {
+    type WatchStream = tonic_health::server::WatchStream;
+
+    async fn check(
+        &self,
+        request: Request<tonic_health::pb::HealthCheckRequest>,
+    ) -> Result<Response<tonic_health::pb::HealthCheckResponse>, Status> {
+        self.inner.check(request).await
+    }
+
+    async fn watch(
+        &self,
+        request: Request<tonic_health::pb::HealthCheckRequest>,
+    ) -> Result<Response<Self::WatchStream>, Status> {
+        let liveness = request.get_ref().service == "liveness";
+        let response = self.inner.watch(request).await?;
+        if liveness {
+            self.liveness_started.notify_one();
+        }
+        Ok(response)
+    }
+}
+
 struct ControlledControl {
     inner: TrtllmMockerService,
     discovery: Arc<Discovery>,