@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 3b03a2cf47
Fork 0

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

[Rust Frontend] Support continuous_usage_stats stream option (#43965)

Co-authored-by: Bugen Zhao <i@bugenzhao.com> Signed-off-by: RickyChen / 陳昭儒 <ricky.chen@infinirc.com> Signed-off-by: Bugen Zhao <i@bugenzhao.com>

Chao-Ju Chen committed Jun 12, 2026 at 01:50 UTC 3b03a2cf4772838da622d81315941bb41bcc03ff
9 files changed +335 -51
rust/src/server/src/routes/openai/chat_completions.rs
+46 -25
@@ -37,6 +37,7 @@ use crate::routes::openai::utils::logprobs::{
37 use crate::routes::openai::utils::types::{
38 ChatLogProbs, FunctionCallDelta, FunctionCallResponse, ToolCall, ToolCallDelta, Usage,
39 };
40 +use crate::routes::openai::utils::usage::ContinuousUsage;
41 use crate::routes::openai::utils::validated_json::ValidatedJson;
42 use crate::state::AppState;
43 use crate::utils::{resolve_request_context, unix_timestamp};
@@ -129,6 +130,8 @@ async fn collect_chat_completion(
130 ResponseOptions {
131 // Ignored: non-streaming responses always include usage.
132 include_usage: _,
133 + // Ignored: non-streaming responses are collected before usage is attached.
134 + include_continuous_usage: _,
135 requested_logprobs,
136 include_prompt_logprobs,
137 include_reasoning,
@@ -249,6 +252,7 @@ async fn chat_completion_chunk_stream(
252 }: ApiServerOptions,
253 ResponseOptions {
254 include_usage,
255 + include_continuous_usage,
256 requested_logprobs,
257 // Ignored: chat streaming prompt logprobs are rejected for Python parity.
258 include_prompt_logprobs: _,
@@ -265,33 +269,47 @@ async fn chat_completion_chunk_stream(
269 // starts or ends, omit its token metadata as well as its visible delta.
270 let mut inside_hidden_reasoning = false;
271 let mut suppress_current_update_metadata = false;
272 + let mut continuous_usage = ContinuousUsage::default();
273 +
274 + /// Yield a chunk with optional continuous usage attached.
275 + macro_rules! yield_chunk {
276 + ($chunk:expr) => {{
277 + let mut chunk = $chunk;
278 + if include_continuous_usage {
279 + chunk.usage = Some(continuous_usage.to_usage());
280 + }
281 + y.yield_ok(chunk).await;
282 + }};
283 + }
284
285 // If the client requested logprobs or token_ids, we need to buffer chunks until
286 // we receive the separate `LogprobsDelta` event, so that we can emit one
287 // combined chunk with both the semantic delta and its per-update metadata.
272 - let mut pending_chunk =
273 - (requested_logprobs || return_token_ids).then(PendingChatChunk::default);
288 + // Continuous usage also buffers so the token count from `LogprobsDelta` can
289 + // be attached to the matching semantic chunk.
290 + let mut pending_chunk = (requested_logprobs || return_token_ids || include_continuous_usage)
291 + .then(PendingChatChunk::default);
292
293 while let Some(next) = stream.next().await {
294 match next {
295 Ok(ChatEvent::Start {
296 prompt_token_ids, ..
297 }) => {
298 + continuous_usage.set_prompt_tokens(prompt_token_ids.len());
299 let mut chunk = start_chunk(&request_id, &response_model, created);
300 if return_token_ids {
301 chunk.prompt_token_ids = Some(prompt_token_ids.to_vec());
302 }
284 - y.yield_ok(chunk).await;
303 + yield_chunk!(chunk);
304 // When echo=true, emit the last assistant message content as a delta chunk.
305 if let Some(echo_text) = &echo {
287 - y.yield_ok(block_delta_chunk(
306 + yield_chunk!(block_delta_chunk(
307 &request_id,
308 &response_model,
309 created,
310 AssistantBlockKind::Text,
311 echo_text.clone(),
293 - ))
294 - .await;
312 + ));
313 }
314 }
315 Ok(ChatEvent::BlockDelta { kind, delta, .. }) => {
@@ -301,14 +319,13 @@ async fn chat_completion_chunk_stream(
319 if let Some(pending_chunk) = pending_chunk.as_mut() {
320 pending_chunk.push_block_delta(kind, delta);
321 } else {
304 - y.yield_ok(block_delta_chunk(
322 + yield_chunk!(block_delta_chunk(
323 &request_id,
324 &response_model,
325 created,
326 kind,
327 delta,
310 - ))
311 - .await;
328 + ));
329 }
330 } else {
331 suppress_current_update_metadata = true;
@@ -318,6 +335,8 @@ async fn chat_completion_chunk_stream(
335 logprobs,
336 token_ids,
337 }) => {
338 + let delta_token_count = token_ids.len();
339 + continuous_usage.add_output_tokens(delta_token_count);
340 let include_metadata =
341 !suppress_current_update_metadata && !inside_hidden_reasoning;
342 suppress_current_update_metadata = false;
@@ -339,16 +358,15 @@ async fn chat_completion_chunk_stream(
358 if let Some(chunk) =
359 pending_chunk.take_chunk(&request_id, &response_model, created)
360 {
342 - y.yield_ok(chunk).await;
361 + yield_chunk!(chunk);
362 }
363 } else if let Some(logprobs) = openai_logprobs {
345 - y.yield_ok(logprobs_only_chunk(
364 + yield_chunk!(logprobs_only_chunk(
365 &request_id,
366 &response_model,
367 created,
368 logprobs,
350 - ))
351 - .await;
369 + ));
370 }
371 }
372 Ok(ChatEvent::BlockStart { kind, .. }) => {
@@ -376,15 +394,14 @@ async fn chat_completion_chunk_stream(
394 if let Some(pending_chunk) = pending_chunk.as_mut() {
395 pending_chunk.push_tool_call_start(tool_index, id, name);
396 } else {
379 - y.yield_ok(tool_call_start_chunk(
397 + yield_chunk!(tool_call_start_chunk(
398 &request_id,
399 &response_model,
400 created,
401 tool_index,
402 id,
403 name,
386 - ))
387 - .await;
404 + ));
405 }
406 }
407 Ok(ChatEvent::ToolCallArgumentsDelta { index, delta }) => {
@@ -392,21 +409,20 @@ async fn chat_completion_chunk_stream(
409 if let Some(pending_chunk) = pending_chunk.as_mut() {
410 pending_chunk.push_tool_call_arguments(tool_index, delta);
411 } else {
395 - y.yield_ok(tool_call_arguments_chunk(
412 + yield_chunk!(tool_call_arguments_chunk(
413 &request_id,
414 &response_model,
415 created,
416 tool_index,
417 delta,
401 - ))
402 - .await;
418 + ));
419 }
420 }
421 Ok(ChatEvent::ToolCallEnd { .. }) => {
422 debug!("ending current tool call");
423 }
424 Ok(ChatEvent::Done {
409 - usage,
425 + usage: final_usage,
426 finish_reason,
427 ..
428 }) => {
@@ -414,18 +430,23 @@ async fn chat_completion_chunk_stream(
430 info!(
431 stream = true,
432 model = %response_model,
417 - prompt_tokens = usage.prompt_token_count,
418 - output_tokens = usage.output_token_count,
433 + prompt_tokens = final_usage.prompt_token_count,
434 + output_tokens = final_usage.output_token_count,
435 finish_reason = finish_reason.as_str(),
436 "chat completion finished"
437 );
438 }
439
440 + continuous_usage.set_final_counts(
441 + final_usage.prompt_token_count,
442 + final_usage.output_token_count,
443 + );
444 +
445 if let Some(pending_chunk) = pending_chunk.as_mut()
446 && let Some(chunk) =
447 pending_chunk.take_chunk(&request_id, &response_model, created)
448 {
428 - y.yield_ok(chunk).await;
449 + yield_chunk!(chunk);
450 }
451
452 match final_chunk(
@@ -435,7 +456,7 @@ async fn chat_completion_chunk_stream(
456 finish_reason,
457 saw_tool_calls,
458 ) {
438 - Ok(chunk) => y.yield_ok(chunk).await,
459 + Ok(chunk) => yield_chunk!(chunk),
460 Err(error) => {
461 error!(
462 error = %error.to_error_response().error.message,
@@ -450,7 +471,7 @@ async fn chat_completion_chunk_stream(
471 &request_id,
472 &response_model,
473 created,
453 - Usage::from_token_usage(usage, enable_prompt_tokens_details),
474 + Usage::from_token_usage(final_usage, enable_prompt_tokens_details),
475 ))
476 .await;
477 }
rust/src/server/src/routes/openai/chat_completions/convert.rs
+51 -2
@@ -33,6 +33,8 @@ pub(super) struct PreparedRequest {
33 pub(super) struct ResponseOptions {
34 /// Whether the caller asked for the final streamed usage chunk.
35 pub include_usage: bool,
36 + /// Whether every streamed chunk should carry cumulative usage.
37 + pub include_continuous_usage: bool,
38 /// Whether the caller requested output logprobs on chat choices.
39 pub requested_logprobs: bool,
40 /// Whether the caller requested top-level prompt logprobs.
@@ -82,6 +84,12 @@ pub(super) fn prepare_chat_request(
84 let include_usage = (request.stream_options.as_ref())
85 .and_then(|options| options.include_usage)
86 .unwrap_or(false);
87 + let include_continuous_usage = include_usage
88 + && request
89 + .stream_options
90 + .as_ref()
91 + .and_then(|options| options.continuous_usage_stats)
92 + .unwrap_or(false);
93 let requested_logprobs = request.logprobs;
94
95 // Auto-enable prompt logprobs for non-streaming echo, matching Python vLLM's
@@ -154,6 +162,7 @@ pub(super) fn prepare_chat_request(
162 response_model,
163 options: ResponseOptions {
164 include_usage,
165 + include_continuous_usage,
166 requested_logprobs,
167 include_prompt_logprobs,
168 include_reasoning,
@@ -375,8 +384,8 @@ mod tests {
384 AssistantRole, ChatCompletionMessage, ChatCompletionRequest,
385 };
386 use crate::routes::openai::utils::types::{
378 - ChatMessage, ContentPart, Function, FunctionCallResponse, ImageUrl, MessageContent, Tool,
379 - ToolCall, ToolChoice, ToolChoiceValue, VideoUrl,
387 + ChatMessage, ContentPart, Function, FunctionCallResponse, ImageUrl, MessageContent,
388 + StreamOptions, Tool, ToolCall, ToolChoice, ToolChoiceValue, VideoUrl,
389 };
390 use crate::utils::{ResolvedRequestContext, resolve_request_context};
391
@@ -456,6 +465,46 @@ mod tests {
465 assert_eq!(prepared.chat_request.tool_choice, ChatToolChoice::Auto);
466 }
467
468 + #[test]
469 + fn prepare_chat_request_maps_stream_usage_and_token_format_options() {
470 + let mut request = base_request();
471 + request.return_tokens_as_token_ids = Some(true);
472 + request.stream_options = Some(StreamOptions {
473 + include_usage: Some(true),
474 + continuous_usage_stats: Some(true),
475 + });
476 +
477 + let prepared = prepare_chat_request(
478 + request,
479 + &served(&["Qwen/Qwen1.5-0.5B-Chat"]),
480 + ResolvedRequestContext::default(),
481 + )
482 + .expect("request is valid");
483 +
484 + assert!(prepared.options.include_usage);
485 + assert!(prepared.options.include_continuous_usage);
486 + assert!(prepared.options.return_tokens_as_token_ids);
487 + }
488 +
489 + #[test]
490 + fn prepare_chat_request_gates_continuous_usage_on_include_usage() {
491 + let mut request = base_request();
492 + request.stream_options = Some(StreamOptions {
493 + include_usage: None,
494 + continuous_usage_stats: Some(true),
495 + });
496 +
497 + let prepared = prepare_chat_request(
498 + request,
499 + &served(&["Qwen/Qwen1.5-0.5B-Chat"]),
500 + ResolvedRequestContext::default(),
501 + )
502 + .expect("request is valid");
503 +
504 + assert!(!prepared.options.include_usage);
505 + assert!(!prepared.options.include_continuous_usage);
506 + }
507 +
508 #[test]
509 fn prepare_chat_request_keeps_optional_sampling_fields_unset() {
510 let prepared = prepare_chat_request(
rust/src/server/src/routes/openai/chat_completions/validate.rs
-9
@@ -137,15 +137,6 @@ pub(super) fn validate_request_compat(
137 "repetition_detection is not supported.",
138 )?;
139
140 - if let Some(options) = &request.stream_options
141 - && options.continuous_usage_stats.is_some()
142 - {
143 - bail_invalid_request!(
144 - param = "stream_options",
145 - "continuous_usage_stats is not supported."
146 - );
147 - }
148 -
140 Ok(())
141 }
142
rust/src/server/src/routes/openai/completions.rs
+29 -6
@@ -31,6 +31,7 @@ use crate::routes::openai::completions::types::{
31 CompletionStreamChoice, CompletionStreamResponse,
32 };
33 use crate::routes::openai::utils::types::LogProbs;
34 +use crate::routes::openai::utils::usage::ContinuousUsage;
35 use crate::routes::openai::utils::validated_json::ValidatedJson;
36 use crate::state::AppState;
37 use crate::utils::{resolve_request_context, unix_timestamp};
@@ -127,6 +128,8 @@ async fn collect_completion(
128 ResponseOptions {
129 // Ignored: non-streaming responses always include usage.
130 include_usage: _,
131 + // Ignored: non-streaming responses are collected before usage is attached.
132 + include_continuous_usage: _,
133 echo,
134 requested_logprobs,
135 include_prompt_logprobs,
@@ -218,6 +221,7 @@ async fn completion_chunk_stream(
221 }: ApiServerOptions,
222 ResponseOptions {
223 include_usage,
224 + include_continuous_usage,
225 echo,
226 requested_logprobs,
227 // Ignored: streaming prompt logprobs are rejected for Python parity.
@@ -230,6 +234,18 @@ async fn completion_chunk_stream(
234 pin_mut!(stream);
235 let mut visible_text_len = 0_u32;
236 let mut first_chunk = true;
237 + let mut continuous_usage = ContinuousUsage::default();
238 +
239 + /// Yield a chunk with optional continuous usage attached.
240 + macro_rules! yield_chunk {
241 + ($chunk:expr) => {{
242 + let mut chunk = $chunk;
243 + if include_continuous_usage {
244 + chunk.usage = Some(continuous_usage.to_usage());
245 + }
246 + y.yield_ok(CompletionSseChunk::Chunk(chunk)).await;
247 + }};
248 + }
249
250 while let Some(next) = stream.next().await {
251 match next {
@@ -237,6 +253,7 @@ async fn completion_chunk_stream(
253 prompt_token_ids, ..
254 }) => {
255 debug!("completion stream started");
256 + continuous_usage.set_prompt_tokens(prompt_token_ids.len());
257 if let Some(prompt) = echo.as_ref() {
258 visible_text_len = text_len(prompt);
259 let mut chunk =
@@ -247,7 +264,7 @@ async fn completion_chunk_stream(
264 }
265 first_chunk = false;
266 }
250 - y.yield_ok(CompletionSseChunk::Chunk(chunk)).await;
267 + yield_chunk!(chunk);
268 } else if return_token_ids {
269 // Emit a chunk with prompt_token_ids in the first streaming response
270 let mut chunk =
@@ -256,7 +273,7 @@ async fn completion_chunk_stream(
273 choice.prompt_token_ids = Some(prompt_token_ids.to_vec());
274 }
275 first_chunk = false;
259 - y.yield_ok(CompletionSseChunk::Chunk(chunk)).await;
276 + yield_chunk!(chunk);
277 }
278 }
279 Ok(DecodedTextEvent::TextDelta {
@@ -281,10 +298,12 @@ async fn completion_chunk_stream(
298 None
299 };
300 let mut chunk = delta_chunk(&request_id, &response_model, created, delta, logprobs);
301 + let delta_token_count = token_ids.len();
302 + continuous_usage.add_output_tokens(delta_token_count);
303 if return_token_ids && let Some(choice) = chunk.choices.first_mut() {
304 choice.token_ids = Some(token_ids);
305 }
287 - y.yield_ok(CompletionSseChunk::Chunk(chunk)).await;
306 + yield_chunk!(chunk);
307 visible_text_len = visible_text_len.saturating_add(delta_text_len);
308
309 if let Some(finished) = finished {
@@ -298,13 +317,17 @@ async fn completion_chunk_stream(
317 "completion finished"
318 );
319 }
301 - y.yield_ok(CompletionSseChunk::Chunk(final_chunk(
320 + continuous_usage.set_final_counts(
321 + finished.usage.prompt_token_count,
322 + finished.usage.output_token_count,
323 + );
324 + let final_chunk = final_chunk(
325 &request_id,
326 &response_model,
327 created,
328 finished.finish_reason,
306 - )?))
307 - .await;
329 + )?;
330 + yield_chunk!(final_chunk);
331
332 if include_usage {
333 y.yield_ok(CompletionSseChunk::Usage(usage_chunk(
rust/src/server/src/routes/openai/completions/convert.rs
+58
@@ -25,6 +25,8 @@ pub(super) struct PreparedRequest {
25 pub(super) struct ResponseOptions {
26 /// Whether the caller asked for the final streamed usage chunk.
27 pub include_usage: bool,
28 + /// Whether every streamed chunk should carry cumulative usage.
29 + pub include_continuous_usage: bool,
30 /// Original text prompt that should be echoed back northbound when
31 /// `echo=true`.
32 pub echo: Option<String>,
@@ -74,6 +76,12 @@ pub(super) fn prepare_completion_request(
76 let include_usage = (request.stream_options.as_ref())
77 .and_then(|options| options.include_usage)
78 .unwrap_or(false);
79 + let include_continuous_usage = include_usage
80 + && request
81 + .stream_options
82 + .as_ref()
83 + .and_then(|options| options.continuous_usage_stats)
84 + .unwrap_or(false);
85 let include_prompt_logprobs = prompt_logprobs.is_some();
86 let echo = request.echo.then(|| request.prompt.as_text().cloned()).flatten();
87
@@ -129,6 +137,7 @@ pub(super) fn prepare_completion_request(
137 response_model,
138 options: ResponseOptions {
139 include_usage,
140 + include_continuous_usage,
141 echo,
142 requested_logprobs: request.logprobs,
143 include_prompt_logprobs,
@@ -247,6 +256,55 @@ mod tests {
256 assert!(!prepared.text_request.decode_options.skip_special_tokens);
257 }
258
259 + #[test]
260 + fn prepare_completion_request_maps_stream_usage_and_token_format_options() {
261 + let request: CompletionRequest = serde_json::from_value(json!({
262 + "model": "Qwen/Qwen1.5-0.5B-Chat",
263 + "prompt": "hello",
264 + "stream": true,
265 + "stream_options": {
266 + "include_usage": true,
267 + "continuous_usage_stats": true
268 + },
269 + "return_tokens_as_token_ids": true
270 + }))
271 + .expect("parse request");
272 +
273 + let prepared = prepare_completion_request(
274 + request,
275 + &served(&["Qwen/Qwen1.5-0.5B-Chat"]),
276 + ResolvedRequestContext::default(),
277 + )
278 + .expect("prepare");
279 +
280 + assert!(prepared.options.include_usage);
281 + assert!(prepared.options.include_continuous_usage);
282 + assert!(prepared.options.return_tokens_as_token_ids);
283 + }
284 +
285 + #[test]
286 + fn prepare_completion_request_gates_continuous_usage_on_include_usage() {
287 + let request: CompletionRequest = serde_json::from_value(json!({
288 + "model": "Qwen/Qwen1.5-0.5B-Chat",
289 + "prompt": "hello",
290 + "stream": true,
291 + "stream_options": {
292 + "continuous_usage_stats": true
293 + }
294 + }))
295 + .expect("parse request");
296 +
297 + let prepared = prepare_completion_request(
298 + request,
299 + &served(&["Qwen/Qwen1.5-0.5B-Chat"]),
300 + ResolvedRequestContext::default(),
301 + )
302 + .expect("prepare");
303 +
304 + assert!(!prepared.options.include_usage);
305 + assert!(!prepared.options.include_continuous_usage);
306 + }
307 +
308 #[test]
309 fn prepare_completion_request_accepts_text_echo() {
310 let request: CompletionRequest = serde_json::from_value(json!({
rust/src/server/src/routes/openai/completions/validate.rs
-9
@@ -95,15 +95,6 @@ pub(super) fn validate_request_compat(
95 );
96 }
97
98 - if let Some(options) = &request.stream_options
99 - && options.continuous_usage_stats.is_some()
100 - {
101 - bail_invalid_request!(
102 - param = "stream_options",
103 - "continuous_usage_stats is not supported."
104 - );
105 - }
106 -
98 Ok(())
99 }
100
rust/src/server/src/routes/openai/utils/mod.rs
+1
@@ -2,4 +2,5 @@ pub mod logprobs;
2 pub mod structured_outputs;
3 pub mod token_ids;
4 pub mod types;
5 +pub mod usage;
6 pub mod validated_json;
rust/src/server/src/routes/openai/utils/usage.rs new
+35
@@ -0,0 +1,35 @@
1 +use super::types::Usage;
2 +
3 +/// Tracks cumulative token counts for OpenAI streaming chunks.
4 +///
5 +/// This helper is intentionally only a counter. Callers decide whether to
6 +/// attach `counts()` to each streamed data chunk, while final usage-only chunks
7 +/// should still be built from the authoritative terminal `TokenUsage`.
8 +#[derive(Debug, Clone, Default)]
9 +pub(crate) struct ContinuousUsage {
10 + prompt_tokens: usize,
11 + output_tokens: usize,
12 +}
13 +
14 +impl ContinuousUsage {
15 + /// Record the prompt-token count reported when a stream starts.
16 + pub(crate) fn set_prompt_tokens(&mut self, prompt_tokens: usize) {
17 + self.prompt_tokens = prompt_tokens;
18 + }
19 +
20 + /// Add newly decoded output tokens to the running completion count.
21 + pub(crate) fn add_output_tokens(&mut self, output_tokens: usize) {
22 + self.output_tokens = self.output_tokens.saturating_add(output_tokens);
23 + }
24 +
25 + /// Replace the running counts with the final counts reported by generation.
26 + pub(crate) fn set_final_counts(&mut self, prompt_tokens: usize, output_tokens: usize) {
27 + self.prompt_tokens = prompt_tokens;
28 + self.output_tokens = output_tokens;
29 + }
30 +
31 + /// Build a streaming usage snapshot without prompt cache details.
32 + pub(crate) fn to_usage(&self) -> Usage {
33 + Usage::from_counts(self.prompt_tokens, self.output_tokens, None)
34 + }
35 +}
rust/src/server/src/routes/tests.rs
+115
@@ -151,6 +151,14 @@ fn sse_data_payloads(text: &str) -> Vec<&str> {
151 text.lines().filter_map(|line| line.strip_prefix("data: ")).collect()
152 }
153
154 +fn sse_json_payloads(text: &str) -> Vec<serde_json::Value> {
155 + sse_data_payloads(text)
156 + .into_iter()
157 + .filter(|payload| *payload != "[DONE]")
158 + .map(|payload| serde_json::from_str(payload).expect("sse json payload"))
159 + .collect()
160 +}
161 +
162 type TestFuture<'a> = Pin<Box<dyn Future<Output = ()> + Send + 'a>>;
163
164 fn boxed_test_future<'a>(future: impl Future<Output = ()> + Send + 'a) -> TestFuture<'a> {
@@ -2341,6 +2349,60 @@ async fn include_usage_adds_final_usage_chunk_before_done() {
2349 assert_eq!(usage_chunk["usage"]["total_tokens"], 25);
2350 }
2351
2352 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2353 +#[serial]
2354 +async fn stream_continuous_usage_stats_adds_usage_to_chat_chunks() {
2355 + let (app, engine_task) = test_app_with_stream_output_specs(default_stream_output_specs()).await;
2356 + let response = app
2357 + .clone()
2358 + .call(
2359 + Request::builder()
2360 + .method("POST")
2361 + .uri("/v1/chat/completions")
2362 + .header("content-type", "application/json")
2363 + .body(Body::from(
2364 + json!({
2365 + "model": "Qwen/Qwen1.5-0.5B-Chat",
2366 + "stream": true,
2367 + "stream_options": {
2368 + "include_usage": true,
2369 + "continuous_usage_stats": true
2370 + },
2371 + "messages": [{"role": "user", "content": "hello"}]
2372 + })
2373 + .to_string(),
2374 + ))
2375 + .expect("build request"),
2376 + )
2377 + .await
2378 + .expect("call app");
2379 +
2380 + assert_eq!(response.status(), StatusCode::OK);
2381 +
2382 + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
2383 + engine_task.await.expect("mock engine task");
2384 + let text = String::from_utf8(body.to_vec()).expect("utf8 body");
2385 + let payloads = sse_json_payloads(&text);
2386 +
2387 + assert!(
2388 + payloads.iter().all(|payload| payload.get("usage").is_some()),
2389 + "{text}"
2390 + );
2391 + assert!(
2392 + payloads.iter().any(|payload| {
2393 + payload["choices"].as_array().is_some_and(|choices| !choices.is_empty())
2394 + && payload["usage"]["completion_tokens"] == json!(1)
2395 + }),
2396 + "{text}"
2397 + );
2398 + let usage_chunk = payloads
2399 + .iter()
2400 + .find(|payload| payload["choices"] == json!([]))
2401 + .expect("final usage chunk");
2402 + assert_eq!(usage_chunk["usage"]["prompt_tokens"], 22);
2403 + assert_eq!(usage_chunk["usage"]["completion_tokens"], 3);
2404 +}
2405 +
2406 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
2407 #[serial]
2408 async fn stream_without_include_usage_keeps_existing_shape() {
@@ -3434,6 +3496,59 @@ async fn completions_happy_path_returns_sse_stream() {
3496 assert_eq!(usage_chunk["usage"]["completion_tokens"], 3);
3497 }
3498
3499 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
3500 +#[serial]
3501 +async fn completions_stream_continuous_usage_stats_adds_usage_to_chunks() {
3502 + let (app, engine_task) = test_app_with_stream_output_specs(default_stream_output_specs()).await;
3503 + let response = app
3504 + .clone()
3505 + .call(
3506 + Request::builder()
3507 + .method("POST")
3508 + .uri("/v1/completions")
3509 + .header("content-type", "application/json")
3510 + .body(Body::from(
3511 + json!({
3512 + "model": "Qwen/Qwen1.5-0.5B-Chat",
3513 + "prompt": "hello",
3514 + "stream": true,
3515 + "stream_options": {
3516 + "include_usage": true,
3517 + "continuous_usage_stats": true
3518 + }
3519 + })
3520 + .to_string(),
3521 + ))
3522 + .expect("build request"),
3523 + )
3524 + .await
3525 + .expect("call app");
3526 +
3527 + assert_eq!(response.status(), StatusCode::OK);
3528 +
3529 + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
3530 + engine_task.await.expect("mock engine task");
3531 + let text = String::from_utf8(body.to_vec()).expect("utf8 body");
3532 + let payloads = sse_json_payloads(&text);
3533 +
3534 + assert!(
3535 + payloads.iter().all(|payload| payload.get("usage").is_some()),
3536 + "{text}"
3537 + );
3538 + assert!(
3539 + payloads.iter().any(|payload| {
3540 + payload["choices"].as_array().is_some_and(|choices| !choices.is_empty())
3541 + && payload["usage"]["completion_tokens"] == json!(1)
3542 + }),
3543 + "{text}"
3544 + );
3545 + let usage_chunk = payloads
3546 + .iter()
3547 + .find(|payload| payload["choices"] == json!([]))
3548 + .expect("final usage chunk");
3549 + assert_eq!(usage_chunk["usage"]["completion_tokens"], 3);
3550 +}
3551 +
3552 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
3553 #[serial]
3554 async fn completions_echo_stream_emits_separate_prompt_chunk() {