[Rust Frontend][gRPC] Remove GenerateStream producer task so disconnects abort requests (#60922)

Co-authored-by: Codex <noreply@openai.com>
Co-authored-by: Bugen Zhao <i@bugenzhao.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Alec Flowers <aflowers@nvidia.com>
Signed-off-by: Bugen Zhao <i@bugenzhao.com>

rust/Cargo.lock

@@ -5997,7 +5997,6 @@ dependencies = [
  "tls-listener",
  "tokio",
  "tokio-openssl",
- "tokio-stream",
  "tokio-util",
  "tonic",
  "tonic-health",

rust/src/server/Cargo.toml

@@ -35,7 +35,6 @@ thiserror-ext.workspace = true
 tls-listener.workspace = true
 tokio.workspace = true
 tokio-openssl.workspace = true
-tokio-stream.workspace = true
 tokio-util.workspace = true
 tonic.workspace = true
 tonic-health.workspace = true

rust/src/server/src/grpc/inference.rs

@@ -5,16 +5,17 @@ use std::pin::Pin;
 use std::sync::Arc;
 use std::time::Instant;
 
-use futures::{Stream, StreamExt as _};
+use asynk_strim_attr::{TryYielder, try_stream};
+use futures::{Stream, StreamExt as _, pin_mut};
 use thiserror_ext::AsReport as _;
-use tokio::sync::mpsc;
-use tokio_stream::wrappers::ReceiverStream;
 use tonic::{Request, Response, Status};
 use tracing::{Span, info, info_span, warn};
 use tracing_futures::Instrument as _;
 use uuid::Uuid;
 use vllm_llm::current_unix_timestamp_secs;
-use vllm_text::{DecodedTextEvent, Prompt, SampledDelta, TextOutputStreamExt as _, TextRequest};
+use vllm_text::{
+    DecodedTextEvent, Prompt, SampledDelta, TextOutputStream, TextOutputStreamExt as _, TextRequest,
+};
 
 use super::convert::{self, ResponseOpts};
 use super::{InferenceServer, pb};
@@ -304,69 +305,74 @@ impl pb::inference_server::Inference for InferenceServiceImpl {
             .await
             .map_err(|error| log_text_error(&request_span, started_at, "submission", error))?;
 
-        let (tx, rx) = mpsc::channel(32);
-
-        let task_span = request_span.clone();
-        tokio::spawn(
-            async move {
-                futures::pin_mut!(stream);
-                while let Some(event) = stream.next().await {
-                    let response = match event {
-                        Err(error) => Err(log_text_error(&task_span, started_at, "stream", error)),
-                        Ok(DecodedTextEvent::Start {
-                            prompt_token_ids,
-                            prompt_logprobs,
-                        }) => {
-                            let prompt_info = convert::to_prompt_info(
-                                &prompt_token_ids,
-                                prompt_logprobs.as_ref(),
-                                &response_opts,
-                            );
-                            Ok(pb::GenerateResponse {
-                                prompt_info: Some(prompt_info),
-                                outputs: None,
-                            })
-                        }
-                        Ok(DecodedTextEvent::TextDelta {
-                            decoded,
-                            sampled:
-                                SampledDelta {
-                                    token_ids,
-                                    logprobs,
-                                },
-                            finished,
-                        }) => convert::to_sequence_output(
-                            &decoded.text,
-                            &token_ids,
-                            logprobs.as_ref(),
-                            finished.as_deref(),
-                            &response_opts,
-                        )
-                        .map(|outputs| pb::GenerateResponse {
-                            prompt_info: None,
-                            outputs: Some(outputs),
-                        }),
-                    };
-
-                    let failed = response.is_err();
-                    if tx.send(response).await.is_err() || failed {
-                        break;
-                    }
-                }
-                info!(
-                    parent: &task_span,
-                    elapsed_ms = started_at.elapsed().as_millis() as u64,
-                    "gRPC inference stream closed"
-                );
-            }
-            .instrument(request_span),
-        );
-
-        let response_stream = ReceiverStream::new(rx);
+        let response_stream =
+            generate_stream_responses(stream, response_opts, request_span.clone(), started_at)
+                .instrument(request_span);
         Ok(Response::new(Box::pin(response_stream)))
     }
 }
 
+/// Convert one decoded text stream into gRPC streaming responses.
+///
+/// The response body owns the generation stream, so a client disconnect drops
+/// it and aborts the request in the engine, even before any output arrives.
+#[try_stream]
+async fn generate_stream_responses(
+    stream: impl TextOutputStream,
+    response_opts: ResponseOpts,
+    request_span: Span,
+    started_at: Instant,
+    mut y: TryYielder<pb::GenerateResponse, Status>,
+) -> Result<(), Status> {
+    pin_mut!(stream);
+
+    while let Some(event) = stream.next().await {
+        let response = match event {
+            Err(error) => Err(log_text_error(&request_span, started_at, "stream", error)),
+            Ok(DecodedTextEvent::Start {
+                prompt_token_ids,
+                prompt_logprobs,
+            }) => {
+                let prompt_info = convert::to_prompt_info(
+                    &prompt_token_ids,
+                    prompt_logprobs.as_ref(),
+                    &response_opts,
+                );
+                Ok(pb::GenerateResponse {
+                    prompt_info: Some(prompt_info),
+                    outputs: None,
+                })
+            }
+            Ok(DecodedTextEvent::TextDelta {
+                decoded,
+                sampled:
+                    SampledDelta {
+                        token_ids,
+                        logprobs,
+                    },
+                finished,
+            }) => convert::to_sequence_output(
+                &decoded.text,
+                &token_ids,
+                logprobs.as_ref(),
+                finished.as_deref(),
+                &response_opts,
+            )
+            .map(|outputs| pb::GenerateResponse {
+                prompt_info: None,
+                outputs: Some(outputs),
+            }),
+        };
+        y.yield_ok(response?).await;
+    }
+    info!(
+        parent: &request_span,
+        elapsed_ms = started_at.elapsed().as_millis() as u64,
+        "gRPC inference stream closed"
+    );
+    Ok(())
+}
+
 fn text_error_to_status(error: vllm_text::Error) -> Status {
     let message = error.to_report_string();
     if error.is_request_validation_error() {

rust/src/server/src/grpc/tests.rs

@@ -1827,6 +1827,70 @@ async fn grpc_without_keepalive_keeps_unresponsive_connection_open() {
     server_task.abort();
 }
 
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+#[serial]
+async fn dropping_generate_stream_aborts_before_engine_output() {
+    let (submitted_tx, submitted_rx) = tokio::sync::oneshot::channel();
+    let (inference_service, control_service, engine_health, engine_task) =
+        setup_grpc_service_with_engine_script(
+            b"engine-grpc-cancel-pending-output".to_vec(),
+            default_ready_response(),
+            Arc::new(FakeTextBackend),
+            None,
+            move |dealer, _push| {
+                boxed_test_future(async move {
+                    let add = recv_engine_message(dealer).await;
+                    assert_eq!(add[0].as_ref(), &[0x00]);
+                    let request: EngineCoreRequest =
+                        rmp_serde::from_slice(&add[1]).expect("decode request");
+                    submitted_tx.send(()).expect("signal submitted request");
+
+                    // A queued request produces no output before cancellation.
+                    let abort = recv_engine_message(dealer).await;
+                    assert_eq!(abort[0].as_ref(), &[0x01]);
+                    let request_ids: Vec<String> =
+                        rmp_serde::from_slice(&abort[1]).expect("decode abort request IDs");
+                    assert_eq!(request_ids, vec![request.request_id]);
+                })
+            },
+        )
+        .await;
+    let (channel, server_task) = start_grpc_test_server(
+        inference_service,
+        control_service,
+        engine_health,
+        tokio_util::sync::CancellationToken::new(),
+    )
+    .await;
+    let mut client = InferenceClient::new(channel);
+    let stream = client
+        .generate_stream(pb::GenerateRequest {
+            request_id: "test-cancel-pending-output".to_string(),
+            model: "test-model".to_string(),
+            prompt: Some(pb::generate_request::Prompt::Text("hello".to_string())),
+            stopping: Some(pb::StoppingCriteria {
+                max_new_tokens: 10,
+                ..Default::default()
+            }),
+            ..Default::default()
+        })
+        .await
+        .expect("start generation")
+        .into_inner();
+    tokio::time::timeout(Duration::from_secs(5), submitted_rx)
+        .await
+        .expect("timed out waiting for engine submission")
+        .expect("engine submission signal");
+
+    drop(stream);
+
+    let result = tokio::time::timeout(Duration::from_secs(5), engine_task).await;
+    server_task.abort();
+    result
+        .expect("disconnect must abort without waiting for engine output")
+        .expect("mock engine task");
+}
+
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 #[serial]
 async fn control_abort_resolves_external_id_and_empty_is_noop() {