@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 182c67daf1
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 2, 07:15 UTC

[Rust Frontend] Support streaming `generate` endpoint (#43779)

Signed-off-by: xunzhuo <xunzhuo@vllm-semantic-router.ai> Co-authored-by: Bugen Zhao <i@bugenzhao.com>

Xunzhuo committed Jun 2, 2026 at 03:30 UTC 182c67daf195bf787a32508bafdf3ae56561cc00
5 files changed +563 -15
rust/src/server/src/routes/inference/generate.rs
+217 -5
@@ -3,23 +3,33 @@ mod types;
3 mod validate;
4
5 use std::collections::HashMap;
6 +use std::convert::Infallible;
7 +use std::result::Result;
8 use std::sync::Arc;
9
10 +use asynk_strim_attr::{TryYielder, try_stream};
11 use axum::Json;
12 use axum::extract::State;
13 use axum::http::HeaderMap;
14 +use axum::response::sse::{Event, Sse};
15 use axum::response::{IntoResponse, Response};
16 +use futures::{Stream, StreamExt as _, pin_mut};
17 use thiserror_ext::AsReport as _;
13 -use tracing::info;
18 +use tracing::{error, info, trace};
19 use tracing_futures::Instrument as _;
20 use vllm_engine_core_client::protocol::logprobs::{Logprobs, PositionLogprobs};
16 -use vllm_llm::{CollectedGenerateOutput, GenerateOutputStreamExt as _};
21 +use vllm_llm::{
22 + CollectedGenerateOutput, FinishReason, GenerateOutput, GenerateOutputStreamExt as _,
23 +};
24
25 use self::convert::prepare_generate_request;
19 -use self::types::{GenerateLogprob, GenerateRequest, GenerateResponse, GenerateResponseChoice};
20 -use crate::error::{ApiError, server_error};
26 +use self::types::{
27 + GenerateLogprob, GenerateRequest, GenerateResponse, GenerateResponseChoice,
28 + GenerateResponseStreamChoice, GenerateStreamResponse,
29 +};
30 +use crate::error::{ApiError, bail_server_error, server_error};
31 use crate::routes::openai::utils::logprobs::clamp_logprob;
22 -use crate::routes::openai::utils::types::{ChatLogProbs, ChatLogProbsContent, TopLogProb};
32 +use crate::routes::openai::utils::types::{ChatLogProbs, ChatLogProbsContent, TopLogProb, Usage};
33 use crate::routes::openai::utils::validated_json::ValidatedJson;
34 use crate::state::AppState;
35 use crate::utils::resolve_request_context;
@@ -46,6 +56,7 @@ pub async fn generate(
56 let log_request = state.enable_log_requests;
57 let include_logprobs = prepared.include_logprobs;
58 let include_prompt_logprobs = prepared.include_prompt_logprobs;
59 + let stream = prepared.stream;
60
61 let raw_stream = match state
62 .chat
@@ -64,6 +75,20 @@ pub async fn generate(
75 }
76 };
77
78 + if stream {
79 + let chunk_stream = generate_chunk_stream(
80 + raw_stream,
81 + prepared.request_id,
82 + log_request,
83 + prepared.include_usage,
84 + prepared.include_continuous_usage,
85 + include_logprobs,
86 + );
87 + let sse_stream = generate_sse_stream(chunk_stream).instrument(request_span);
88 +
89 + return Sse::new(sse_stream).into_response();
90 + }
91 +
92 let collected = match raw_stream.collect_output().instrument(request_span.clone()).await {
93 Ok(collected) => collected,
94 Err(error) => {
@@ -98,6 +123,102 @@ pub async fn generate(
123 Json(response).into_response()
124 }
125
126 +#[try_stream]
127 +async fn generate_chunk_stream(
128 + stream: impl Stream<Item = vllm_llm::Result<GenerateOutput>>,
129 + request_id: String,
130 + log_request: bool,
131 + include_usage: bool,
132 + include_continuous_usage: bool,
133 + include_logprobs: bool,
134 + mut y: TryYielder<GenerateStreamResponse, ApiError>,
135 +) -> Result<(), ApiError> {
136 + pin_mut!(stream);
137 + let mut prompt_tokens: Option<u32> = None;
138 + let mut output_tokens = 0_u32;
139 +
140 + while let Some(next) = stream.next().await {
141 + match next {
142 + Ok(output) => {
143 + if prompt_tokens.is_none() {
144 + prompt_tokens =
145 + output.prompt_info.as_ref().map(|info| info.prompt_token_ids.len() as u32);
146 + }
147 + let usage_prompt_tokens = prompt_tokens.unwrap_or_default();
148 +
149 + let token_ids = output.token_ids;
150 + output_tokens = output_tokens.saturating_add(token_ids.len() as u32);
151 + let finish_reason = output.finish_reason;
152 +
153 + if matches!(finish_reason.as_ref(), Some(FinishReason::Error)) {
154 + bail_server_error!("Internal server error");
155 + }
156 +
157 + if let Some(finish_reason) = finish_reason.as_ref()
158 + && log_request
159 + {
160 + info!(
161 + stream = true,
162 + prompt_tokens = usage_prompt_tokens,
163 + output_tokens,
164 + finish_reason = finish_reason.as_str(),
165 + "generate finished"
166 + );
167 + }
168 +
169 + if token_ids.is_empty() && finish_reason.is_none() {
170 + continue;
171 + }
172 +
173 + let logprobs = if include_logprobs && !token_ids.is_empty() {
174 + let logprobs = output.logprobs.as_ref().ok_or_else(|| {
175 + server_error!(
176 + "raw generate stream requested logprobs but generation returned none"
177 + )
178 + })?;
179 + Some(raw_logprobs_to_openai_chat(logprobs)?)
180 + } else {
181 + None
182 + };
183 +
184 + y.yield_ok(GenerateStreamResponse {
185 + request_id: request_id.clone(),
186 + choices: vec![GenerateResponseStreamChoice {
187 + index: 0,
188 + logprobs,
189 + finish_reason: finish_reason.map(|reason| reason.as_str().to_string()),
190 + token_ids,
191 + }],
192 + usage: include_continuous_usage
193 + .then(|| Usage::from_counts(usage_prompt_tokens, output_tokens)),
194 + })
195 + .await;
196 + }
197 + Err(error) => {
198 + error!(
199 + error = %error.as_report(),
200 + "raw generate stream failed"
201 + );
202 + bail_server_error!("{}", error.to_report_string());
203 + }
204 + }
205 + }
206 +
207 + if include_usage {
208 + y.yield_ok(GenerateStreamResponse {
209 + request_id,
210 + choices: Vec::new(),
211 + usage: Some(Usage::from_counts(
212 + prompt_tokens.unwrap_or_default(),
213 + output_tokens,
214 + )),
215 + })
216 + .await;
217 + }
218 +
219 + Ok(())
220 +}
221 +
222 fn collect_generate(
223 collected: CollectedGenerateOutput,
224 request_id: String,
@@ -213,3 +334,94 @@ fn position_to_logprob_map(position: &PositionLogprobs) -> HashMap<u32, Generate
334 fn format_token_id(token_id: u32) -> String {
335 format!("token_id:{token_id}")
336 }
337 +
338 +/// Convert one raw-generate chunk stream into SSE events.
339 +#[try_stream]
340 +async fn generate_sse_stream(
341 + stream: impl Stream<Item = Result<GenerateStreamResponse, ApiError>>,
342 + mut y: TryYielder<Event, Infallible>,
343 +) -> Result<(), Infallible> {
344 + pin_mut!(stream);
345 +
346 + while let Some(next) = stream.next().await {
347 + match next {
348 + Ok(chunk) => y.yield_ok(to_sse_event(&chunk)).await,
349 + Err(error) => {
350 + y.yield_ok(to_error_sse_event(&error)).await;
351 + break;
352 + }
353 + }
354 + }
355 +
356 + y.yield_ok(done_sse_event()).await;
357 + Ok(())
358 +}
359 +
360 +fn to_sse_event(chunk: &GenerateStreamResponse) -> Event {
361 + let payload = serde_json::to_string(chunk).expect("generate chunk must serialize to JSON");
362 + trace!(payload, "generate emitting chunk");
363 + Event::default().data(payload)
364 +}
365 +
366 +fn to_error_sse_event(error: &ApiError) -> Event {
367 + let payload = serde_json::to_string(&error.to_error_response())
368 + .expect("ErrorResponse must serialize to JSON");
369 + trace!(payload, "generate emitting error");
370 + Event::default().data(payload)
371 +}
372 +
373 +fn done_sse_event() -> Event {
374 + trace!("generate emitting done");
375 + Event::default().data("[DONE]")
376 +}
377 +
378 +#[cfg(test)]
379 +mod tests {
380 + use std::sync::Arc;
381 +
382 + use futures::{TryStreamExt as _, stream};
383 + use vllm_llm::GeneratePromptInfo;
384 +
385 + use super::*;
386 +
387 + #[tokio::test]
388 + async fn generate_chunk_stream_captures_late_prompt_info() {
389 + let stream = stream::iter(vec![
390 + Ok(GenerateOutput {
391 + request_id: String::new(),
392 + prompt_info: None,
393 + token_ids: Vec::new(),
394 + logprobs: None,
395 + finish_reason: None,
396 + kv_transfer_params: None,
397 + }),
398 + Ok(GenerateOutput {
399 + request_id: String::new(),
400 + prompt_info: Some(GeneratePromptInfo {
401 + prompt_token_ids: Arc::from([11_u32, 22_u32]),
402 + prompt_logprobs: None,
403 + }),
404 + token_ids: vec![33],
405 + logprobs: None,
406 + finish_reason: Some(FinishReason::stop_eos()),
407 + kv_transfer_params: None,
408 + }),
409 + ]);
410 +
411 + let chunks: Vec<_> =
412 + generate_chunk_stream(stream, "raw-stream".to_string(), false, true, true, false)
413 + .try_collect()
414 + .await
415 + .expect("collect chunks");
416 +
417 + assert_eq!(chunks.len(), 2);
418 + assert_eq!(
419 + chunks[0].usage.as_ref().expect("chunk usage").prompt_tokens,
420 + 2
421 + );
422 + assert_eq!(
423 + chunks[1].usage.as_ref().expect("final usage").prompt_tokens,
424 + 2
425 + );
426 + }
427 +}
rust/src/server/src/routes/inference/generate/convert.rs
+42
@@ -10,6 +10,9 @@ use crate::utils::{ResolvedRequestContext, merge_kv_transfer_params};
10 pub struct PreparedRequest {
11 pub request_id: String,
12 pub text_request: TextRequest,
13 + pub stream: bool,
14 + pub include_usage: bool,
15 + pub include_continuous_usage: bool,
16 pub include_logprobs: bool,
17 pub include_prompt_logprobs: bool,
18 }
@@ -23,6 +26,18 @@ pub fn prepare_generate_request(
26 ) -> Result<PreparedRequest, ApiError> {
27 validate::validate_request_compat(&request, served_model_names)?;
28
29 + let stream = request.stream;
30 + let include_usage = request
31 + .stream_options
32 + .as_ref()
33 + .and_then(|options| options.include_usage)
34 + .unwrap_or(false);
35 + let include_continuous_usage = include_usage
36 + && request
37 + .stream_options
38 + .as_ref()
39 + .and_then(|options| options.continuous_usage_stats)
40 + .unwrap_or(false);
41 let include_logprobs = request.sampling_params.logprobs.is_some();
42 let include_prompt_logprobs = request.sampling_params.prompt_logprobs.is_some();
43 let mut sampling_params = request.sampling_params;
@@ -47,6 +62,9 @@ pub fn prepare_generate_request(
62 Ok(PreparedRequest {
63 request_id: ctx.request_id,
64 text_request,
65 + stream,
66 + include_usage,
67 + include_continuous_usage,
68 include_logprobs,
69 include_prompt_logprobs,
70 })
@@ -109,4 +127,28 @@ mod tests {
127 Some(json!({"connector": "x"}))
128 );
129 }
130 +
131 + #[test]
132 + fn prepare_generate_request_gates_continuous_usage_on_include_usage() {
133 + let request: GenerateRequest = serde_json::from_value(json!({
134 + "model": "Qwen/Qwen1.5-0.5B-Chat",
135 + "token_ids": [11, 22],
136 + "stream": true,
137 + "stream_options": {
138 + "continuous_usage_stats": true
139 + },
140 + "sampling_params": {}
141 + }))
142 + .expect("parse request");
143 +
144 + let prepared = prepare_generate_request(
145 + request,
146 + &["Qwen/Qwen1.5-0.5B-Chat".to_string()],
147 + ResolvedRequestContext::default(),
148 + )
149 + .expect("prepare");
150 +
151 + assert!(!prepared.include_usage);
152 + assert!(!prepared.include_continuous_usage);
153 + }
154 }
rust/src/server/src/routes/inference/generate/types.rs
+21 -1
@@ -5,7 +5,7 @@ use serde_json::{Map, Value};
5 use validator::Validate;
6 use vllm_text::SamplingParams;
7
8 -use crate::routes::openai::utils::types::{ChatLogProbs, Normalizable};
8 +use crate::routes::openai::utils::types::{ChatLogProbs, Normalizable, StreamOptions, Usage};
9
10 /// vLLM-compatible request type for the token-in/token-out generate API.
11 #[serde_with::skip_serializing_none]
@@ -17,6 +17,7 @@ pub struct GenerateRequest {
17 pub sampling_params: SamplingParams,
18 #[serde(default)]
19 pub stream: bool,
20 + pub stream_options: Option<StreamOptions>,
21 pub cache_salt: Option<String>,
22 #[serde(default)]
23 pub priority: i32,
@@ -37,6 +38,25 @@ pub(super) struct GenerateResponseChoice {
38 pub token_ids: Vec<u32>,
39 }
40
41 +/// Mirrors the Python vLLM `GenerateResponseStreamChoice` class.
42 +#[serde_with::skip_serializing_none]
43 +#[derive(Debug, Clone, Serialize)]
44 +pub(super) struct GenerateResponseStreamChoice {
45 + pub index: u32,
46 + pub logprobs: Option<ChatLogProbs>,
47 + pub finish_reason: Option<String>,
48 + pub token_ids: Vec<u32>,
49 +}
50 +
51 +/// Mirrors the Python vLLM `GenerateStreamResponse` class.
52 +#[serde_with::skip_serializing_none]
53 +#[derive(Debug, Clone, Serialize)]
54 +pub(super) struct GenerateStreamResponse {
55 + pub request_id: String,
56 + pub choices: Vec<GenerateResponseStreamChoice>,
57 + pub usage: Option<Usage>,
58 +}
59 +
60 /// Mirrors the Python vLLM `GenerateResponse` class.
61 #[serde_with::skip_serializing_none]
62 #[derive(Debug, Clone, Serialize)]
rust/src/server/src/routes/inference/generate/validate.rs
+19 -3
@@ -13,8 +13,11 @@ pub(super) fn validate_request_compat(
13 return Err(ApiError::model_not_found(model.clone()));
14 }
15
16 - if request.stream {
17 - bail_invalid_request!(param = "stream", "stream=true is not supported.");
16 + if request.stream_options.is_some() && !request.stream {
17 + bail_invalid_request!(
18 + param = "stream_options",
19 + "stream_options are only supported when stream=true."
20 + );
21 }
22
23 if request.token_ids.is_empty() {
@@ -65,11 +68,24 @@ mod tests {
68 }
69
70 #[test]
68 - fn validate_request_compat_rejects_streaming() {
71 + fn validate_request_compat_accepts_streaming() {
72 let request = GenerateRequest {
73 stream: true,
74 ..base_request()
75 };
76 + assert!(validate_request_compat(&request, &served(&["Qwen/Qwen1.5-0.5B-Chat"])).is_ok());
77 + }
78 +
79 + #[test]
80 + fn validate_request_compat_rejects_stream_options_without_streaming() {
81 + let request: GenerateRequest = serde_json::from_value(json!({
82 + "model": "Qwen/Qwen1.5-0.5B-Chat",
83 + "token_ids": [11, 22],
84 + "stream": false,
85 + "stream_options": {"include_usage": true},
86 + "sampling_params": {}
87 + }))
88 + .expect("parse request");
89 assert!(validate_request_compat(&request, &served(&["Qwen/Qwen1.5-0.5B-Chat"])).is_err());
90 }
91
rust/src/server/src/routes/tests.rs
+264 -6
@@ -2425,8 +2425,72 @@ async fn non_stream_raw_generate_returns_token_output_envelope() {
2425
2426 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2427 #[serial]
2428 -async fn raw_generate_rejects_streaming() {
2429 - let mut app = test_app().await;
2428 +async fn stream_raw_generate_returns_sse_chunks_and_usage() {
2429 + let ipc = IpcNamespace::new().expect("create ipc namespace");
2430 + let handshake_address = ipc.handshake_endpoint();
2431 + let engine_id = b"engine-raw-generate-stream".to_vec();
2432 +
2433 + let engine_task = MockEngineTask::new(spawn_mock_engine_task(
2434 + handshake_address.clone(),
2435 + engine_id.clone(),
2436 + |dealer, push| {
2437 + boxed_test_future(async move {
2438 + let add = recv_engine_message(dealer).await;
2439 + let request: EngineCoreRequest =
2440 + rmp_serde::from_slice(&add[1]).expect("decode request");
2441 + assert_eq!(request.prompt_token_ids.as_deref(), Some(&[11, 22][..]));
2442 + assert_eq!(request.external_req_id.as_deref(), Some("raw-stream"));
2443 +
2444 + send_outputs(
2445 + push,
2446 + EngineCoreOutputs {
2447 + engine_index: 0,
2448 + outputs: vec![
2449 + request_output_with_logprobs(
2450 + &request.request_id,
2451 + vec![33],
2452 + None,
2453 + None,
2454 + Some(sample_logprobs_for_token(33, 34)),
2455 + None,
2456 + ),
2457 + request_output_with_logprobs(
2458 + &request.request_id,
2459 + vec![44],
2460 + Some(EngineCoreFinishReason::Stop),
2461 + None,
2462 + Some(sample_logprobs_for_token(44, 45)),
2463 + None,
2464 + ),
2465 + ],
2466 + scheduler_stats: None,
2467 + timestamp: 0.0,
2468 + utility_output: None,
2469 + finished_requests: None,
2470 + wave_complete: None,
2471 + start_wave: None,
2472 + },
2473 + )
2474 + .await;
2475 + })
2476 + },
2477 + ));
2478 +
2479 + let client = EngineCoreClient::connect(
2480 + EngineCoreClientConfig::new_single(handshake_address)
2481 + .with_model_name("test-model")
2482 + .with_local_input_output_addresses(
2483 + Some(ipc.input_endpoint()),
2484 + Some(ipc.output_endpoint()),
2485 + ),
2486 + )
2487 + .await
2488 + .expect("connect client");
2489 + let chat = ChatLlm::from_shared_backend(Llm::new(client), Arc::new(FakeChatBackend::new()));
2490 + let mut app = build_router(Arc::new(AppState::new(
2491 + vec!["Qwen/Qwen1.5-0.5B-Chat".to_string()],
2492 + chat,
2493 + )));
2494
2495 let response = app
2496 .call(
@@ -2437,9 +2501,17 @@ async fn raw_generate_rejects_streaming() {
2501 .body(Body::from(
2502 json!({
2503 "model": "Qwen/Qwen1.5-0.5B-Chat",
2504 + "request_id": "raw-stream",
2505 "token_ids": [11, 22],
2506 "stream": true,
2442 - "sampling_params": {}
2507 + "stream_options": {
2508 + "include_usage": true,
2509 + "continuous_usage_stats": true
2510 + },
2511 + "sampling_params": {
2512 + "max_tokens": 2,
2513 + "logprobs": 1
2514 + }
2515 })
2516 .to_string(),
2517 ))
@@ -2448,10 +2520,196 @@ async fn raw_generate_rejects_streaming() {
2520 .await
2521 .expect("call app");
2522
2451 - assert_eq!(response.status(), StatusCode::BAD_REQUEST);
2523 + assert_eq!(response.status(), StatusCode::OK);
2524 + assert_eq!(
2525 + response.headers().get("content-type").and_then(|value| value.to_str().ok()),
2526 + Some("text/event-stream")
2527 + );
2528 +
2529 let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
2453 - let json: serde_json::Value = serde_json::from_slice(&body).expect("decode json");
2454 - assert_eq!(json["error"]["param"], "stream");
2530 + engine_task.await.expect("mock engine task");
2531 + let text = String::from_utf8(body.to_vec()).expect("utf8 body");
2532 + let payloads = sse_data_payloads(&text);
2533 + assert_eq!(payloads.len(), 4, "{text}");
2534 +
2535 + let first: serde_json::Value = serde_json::from_str(payloads[0]).expect("first chunk json");
2536 + assert_eq!(first["request_id"], "raw-stream");
2537 + assert_eq!(first["choices"][0]["index"], 0);
2538 + assert_eq!(first["choices"][0]["token_ids"], json!([33]));
2539 + assert_eq!(
2540 + first["choices"][0]["logprobs"]["content"][0]["token"],
2541 + "token_id:33"
2542 + );
2543 + assert_eq!(first["usage"]["prompt_tokens"], 2);
2544 + assert_eq!(first["usage"]["completion_tokens"], 1);
2545 +
2546 + let second: serde_json::Value = serde_json::from_str(payloads[1]).expect("second chunk json");
2547 + assert_eq!(second["choices"][0]["token_ids"], json!([44]));
2548 + assert_eq!(second["choices"][0]["finish_reason"], "stop");
2549 + assert_eq!(second["usage"]["completion_tokens"], 2);
2550 +
2551 + let usage: serde_json::Value = serde_json::from_str(payloads[2]).expect("usage chunk json");
2552 + assert_eq!(usage["choices"], json!([]));
2553 + assert_eq!(usage["usage"]["prompt_tokens"], 2);
2554 + assert_eq!(usage["usage"]["completion_tokens"], 2);
2555 + assert_eq!(usage["usage"]["total_tokens"], 4);
2556 + assert_eq!(payloads[3], "[DONE]");
2557 +}
2558 +
2559 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2560 +#[serial]
2561 +async fn stream_raw_generate_emits_final_usage_without_continuous_usage() {
2562 + let (mut app, engine_task) = test_app_with_stream_output_specs(vec![
2563 + (vec![33], None),
2564 + (vec![44], Some(EngineCoreFinishReason::Stop)),
2565 + ])
2566 + .await;
2567 +
2568 + let response = app
2569 + .call(
2570 + Request::builder()
2571 + .method("POST")
2572 + .uri("/inference/v1/generate")
2573 + .header("content-type", "application/json")
2574 + .body(Body::from(
2575 + json!({
2576 + "model": "Qwen/Qwen1.5-0.5B-Chat",
2577 + "request_id": "raw-stream-final-usage",
2578 + "token_ids": [11, 22],
2579 + "stream": true,
2580 + "stream_options": {
2581 + "include_usage": true
2582 + },
2583 + "sampling_params": {
2584 + "max_tokens": 2
2585 + }
2586 + })
2587 + .to_string(),
2588 + ))
2589 + .expect("build request"),
2590 + )
2591 + .await
2592 + .expect("call app");
2593 +
2594 + assert_eq!(response.status(), StatusCode::OK);
2595 +
2596 + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
2597 + engine_task.await.expect("mock engine task");
2598 + let text = String::from_utf8(body.to_vec()).expect("utf8 body");
2599 + let payloads = sse_data_payloads(&text);
2600 + assert_eq!(payloads.len(), 4, "{text}");
2601 +
2602 + let first: serde_json::Value = serde_json::from_str(payloads[0]).expect("first chunk json");
2603 + assert_eq!(first["choices"][0]["token_ids"], json!([33]));
2604 + assert!(first.get("usage").is_none());
2605 +
2606 + let second: serde_json::Value = serde_json::from_str(payloads[1]).expect("second chunk json");
2607 + assert_eq!(second["choices"][0]["token_ids"], json!([44]));
2608 + assert_eq!(second["choices"][0]["finish_reason"], "stop");
2609 + assert!(second.get("usage").is_none());
2610 +
2611 + let usage: serde_json::Value = serde_json::from_str(payloads[2]).expect("usage chunk json");
2612 + assert_eq!(usage["choices"], json!([]));
2613 + assert_eq!(usage["usage"]["prompt_tokens"], 2);
2614 + assert_eq!(usage["usage"]["completion_tokens"], 2);
2615 + assert_eq!(usage["usage"]["total_tokens"], 4);
2616 + assert_eq!(payloads[3], "[DONE]");
2617 +}
2618 +
2619 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2620 +#[serial]
2621 +async fn stream_raw_generate_emits_empty_finish_chunk() {
2622 + let (mut app, engine_task) = test_app_with_stream_output_specs(vec![
2623 + (vec![33], None),
2624 + (vec![], Some(EngineCoreFinishReason::Stop)),
2625 + ])
2626 + .await;
2627 +
2628 + let response = app
2629 + .call(
2630 + Request::builder()
2631 + .method("POST")
2632 + .uri("/inference/v1/generate")
2633 + .header("content-type", "application/json")
2634 + .body(Body::from(
2635 + json!({
2636 + "model": "Qwen/Qwen1.5-0.5B-Chat",
2637 + "request_id": "raw-stream-empty-finish",
2638 + "token_ids": [11, 22],
2639 + "stream": true,
2640 + "sampling_params": {
2641 + "max_tokens": 2
2642 + }
2643 + })
2644 + .to_string(),
2645 + ))
2646 + .expect("build request"),
2647 + )
2648 + .await
2649 + .expect("call app");
2650 +
2651 + assert_eq!(response.status(), StatusCode::OK);
2652 +
2653 + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
2654 + engine_task.await.expect("mock engine task");
2655 + let text = String::from_utf8(body.to_vec()).expect("utf8 body");
2656 + let payloads = sse_data_payloads(&text);
2657 + assert_eq!(payloads.len(), 3, "{text}");
2658 +
2659 + let first: serde_json::Value = serde_json::from_str(payloads[0]).expect("first chunk json");
2660 + assert_eq!(first["choices"][0]["token_ids"], json!([33]));
2661 + assert!(first["choices"][0].get("finish_reason").is_none());
2662 +
2663 + let second: serde_json::Value = serde_json::from_str(payloads[1]).expect("second chunk json");
2664 + assert_eq!(second["choices"][0]["token_ids"], json!([]));
2665 + assert_eq!(second["choices"][0]["finish_reason"], "stop");
2666 + assert_eq!(payloads[2], "[DONE]");
2667 +}
2668 +
2669 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2670 +#[serial]
2671 +async fn stream_raw_generate_error_finish_returns_sse_error() {
2672 + let (mut app, engine_task) =
2673 + test_app_with_stream_output_specs(vec![(vec![], Some(EngineCoreFinishReason::Error))])
2674 + .await;
2675 +
2676 + let response = app
2677 + .call(
2678 + Request::builder()
2679 + .method("POST")
2680 + .uri("/inference/v1/generate")
2681 + .header("content-type", "application/json")
2682 + .body(Body::from(
2683 + json!({
2684 + "model": "Qwen/Qwen1.5-0.5B-Chat",
2685 + "request_id": "raw-stream-error",
2686 + "token_ids": [11, 22],
2687 + "stream": true,
2688 + "stream_options": {
2689 + "include_usage": true
2690 + },
2691 + "sampling_params": {
2692 + "max_tokens": 2
2693 + }
2694 + })
2695 + .to_string(),
2696 + ))
2697 + .expect("build request"),
2698 + )
2699 + .await
2700 + .expect("call app");
2701 +
2702 + assert_eq!(response.status(), StatusCode::OK);
2703 +
2704 + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
2705 + engine_task.await.expect("mock engine task");
2706 + let text = String::from_utf8(body.to_vec()).expect("utf8 body");
2707 +
2708 + assert!(text.contains("\"type\":\"server_error\""), "{text}");
2709 + assert!(text.contains("Internal server error"), "{text}");
2710 + assert!(!text.contains("\"finish_reason\":\"error\""), "{text}");
2711 + assert!(!text.contains("\"usage\":"), "{text}");
2712 + assert!(text.trim_end().ends_with("data: [DONE]"), "{text}");
2713 }
2714
2715 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]