@tduncan / Nvidia-Optimization-Stack-VLLM / commits / d6bce42983

[Rust Frontend] Report reasoning tokens in chat completion usage (#54883)

Signed-off-by: Bugen Zhao <i@bugenzhao.com>

Bugen Zhao committed Sep 4, 2026 at 09:47 UTC d6bce42983bc0b2095ad6422dbf1399e219ae572
15 files changed +735 -88
.buildkite/test_areas/rust_frontend.yaml
+6
@@ -15,7 +15,9 @@ steps:
15 - tests/utils.py
16 - tests/benchmarks/test_serve_cli.py
17 - tests/entrypoints/openai/chat_completion/test_chat_completion.py
18 + - tests/entrypoints/openai/chat_completion/test_include_reasoning.py
19 - tests/entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py
20 + - tests/entrypoints/openai/chat_completion/test_serving_chat.py
21
22 # - tests/entrypoints/openai/completion/test_prompt_validation.py
23 - tests/entrypoints/launchers/test_shutdown.py
@@ -27,7 +29,9 @@ steps:
29 - export VLLM_WORKER_MULTIPROC_METHOD=spawn
30 - pytest -v -s benchmarks/test_serve_cli.py -k "not insecure and not (test_bench_serve and not test_bench_serve_chat)"
31 - pytest -v -s entrypoints/openai/chat_completion/test_chat_completion.py -k "not test_invalid_json_schema and not test_invalid_regex and not test_kv_transfer_prompt_token_ids_round_trip and not test_kv_transfer_prompt_token_ids_streaming"
32 + - pytest -v -s entrypoints/openai/chat_completion/test_include_reasoning.py
33 - pytest -v -s entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py -k "not multiple"
34 + - pytest -v -s entrypoints/openai/chat_completion/test_serving_chat.py -k "test_chat_per_request_metrics_follow_server_flag or test_streaming_reasoning_usage_counts_across_deltas or test_completion_tokens_details"
35
36 # - pytest -v -s entrypoints/openai/completion/test_prompt_validation.py -k "not prompt_embeds"
37 - pytest -v -s entrypoints/launchers/test_shutdown.py -k "not engine_failure and not test_abort_timeout_exits_quickly"
@@ -53,7 +57,9 @@ steps:
57 - tests/utils.py
58 - tests/benchmarks/test_serve_cli.py
59 - tests/entrypoints/openai/chat_completion/test_chat_completion.py
60 + - tests/entrypoints/openai/chat_completion/test_include_reasoning.py
61 - tests/entrypoints/openai/chat_completion/test_chat_logit_bias_validation.py
62 + - tests/entrypoints/openai/chat_completion/test_serving_chat.py
63 - tests/entrypoints/openai/completion/test_shutdown.py
64 - tests/entrypoints/openai/test_return_token_ids.py
65 - tests/entrypoints/openai/test_uds.py
rust/src/chat/src/event.rs
+34 -1
@@ -123,6 +123,37 @@ impl [AssistantContentBlock] {
123 }
124 }
125
126 +/// Chat-level token usage for one completed request.
127 +///
128 +/// Extends the engine-level [`TokenUsage`] with reasoning attribution measured
129 +/// by the chat output pipeline. [`TokenUsage`] itself stays engine-level and
130 +/// is not touched by parsing.
131 +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
132 +pub struct ChatTokenUsage {
133 + /// Engine-level token usage.
134 + pub engine: TokenUsage,
135 + /// Number of generated tokens attributed to reasoning; 0 when the
136 + /// configured parser has no reasoning channel.
137 + pub reasoning_tokens: usize,
138 +}
139 +
140 +impl From<TokenUsage> for ChatTokenUsage {
141 + fn from(engine: TokenUsage) -> Self {
142 + Self {
143 + engine,
144 + reasoning_tokens: 0,
145 + }
146 + }
147 +}
148 +
149 +impl Deref for ChatTokenUsage {
150 + type Target = TokenUsage;
151 +
152 + fn deref(&self) -> &Self::Target {
153 + &self.engine
154 + }
155 +}
156 +
157 /// Final structured assistant message assembled from the event stream.
158 #[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
159 pub struct AssistantMessage {
@@ -172,6 +203,8 @@ pub enum ChatEvent {
203 index: usize,
204 kind: AssistantBlockKind,
205 delta: String,
206 + /// Number of generated tokens attributed to this delta, when available.
207 + token_count: Option<usize>,
208 },
209 /// Per-decoded-update sample metadata: logprobs and/or output token IDs.
210 LogprobsDelta {
@@ -201,7 +234,7 @@ pub enum ChatEvent {
234 /// metadata.
235 Done {
236 message: AssistantMessage,
204 - usage: TokenUsage,
237 + usage: ChatTokenUsage,
238 finish_reason: FinishReason,
239 /// Connector-specific KV transfer parameters for disaggregated serving.
240 kv_transfer_params: Option<serde_json::Value>,
rust/src/chat/src/lib.rs
+1 -1
@@ -18,7 +18,7 @@ pub use backend::{
18 pub use error::{Error, Result};
19 pub use event::{
20 AssistantBlockKind, AssistantContentBlock, AssistantMessage, AssistantMessageExt,
21 - AssistantToolCall, ChatEvent,
21 + AssistantToolCall, ChatEvent, ChatTokenUsage,
22 };
23 use futures::{StreamExt, TryStreamExt as _};
24 pub use llm_multimodal::MediaContentPart;
rust/src/chat/src/output/default/unified.rs
+214 -14
@@ -17,7 +17,7 @@ use vllm_text::output::{DecodedText, DecodedTextEvent, SampledDelta};
17
18 use crate::Result;
19 use crate::error::Error;
20 -use crate::event::AssistantBlockKind;
20 +use crate::event::{AssistantBlockKind, ChatTokenUsage};
21 use crate::output::{AssistantEvent, DecodedTextEventStream, generate_tool_call_id};
22
23 /// Per-stream unified parsing state.
@@ -31,6 +31,10 @@ struct UnifiedParserState {
31 /// Supported parsers currently emit at most one active tool call at a time.
32 /// Change this to an indexed map if a model needs interleaved calls later.
33 open_call_index: Option<usize>,
34 + /// Running count of tokens attributed to reasoning. Kept solely to
35 + /// produce the final count on `Done`; per-delta counts are carried by
36 + /// each `TextDelta` event itself.
37 + reasoning_tokens: usize,
38 }
39
40 impl UnifiedParserState {
@@ -40,6 +44,7 @@ impl UnifiedParserState {
44 parser,
45 parser_failed: false,
46 open_call_index: None,
47 + reasoning_tokens: 0,
48 }
49 }
50
@@ -69,10 +74,10 @@ impl UnifiedParserState {
74 return Ok(text_event(AssistantBlockKind::Text, delta.text).into_iter().collect());
75 }
76
72 - let mut output = UnifiedParserOutput::default();
77 // The parser consumes the delta; keep its text for the fallback path
78 // that re-emits it as plain text.
79 let fallback_text = delta.text.clone();
80 + let mut output = UnifiedParserOutput::default();
81 match self.parser.parse_into(delta, &mut output) {
82 Ok(()) => {
83 let mut events = Vec::new();
@@ -138,9 +143,13 @@ impl UnifiedParserState {
143 self.open_call_index = None;
144 push_text_delta(events, AssistantBlockKind::Text, delta);
145 }
141 - UnifiedParserEvent::Reasoning(delta) => {
146 + UnifiedParserEvent::Reasoning(piece) => {
147 self.open_call_index = None;
143 - push_text_delta(events, AssistantBlockKind::Reasoning, delta.text);
148 + // By construction this counts exactly the tokens emitted to
149 + // the client as reasoning, so the parser-failure fallback
150 + // needs no special handling.
151 + self.reasoning_tokens += piece.attributions.len();
152 + push_piece_delta(events, AssistantBlockKind::Reasoning, piece);
153 }
154 UnifiedParserEvent::ToolCall(item) => {
155 self.process_tool_item(item, events)?;
@@ -200,12 +209,17 @@ impl UnifiedParserState {
209 }
210 }
211
203 -/// Build one plain text event if `delta` is non-empty.
212 +/// Build one plain text event with no measured tokens, if `delta` is
213 +/// non-empty.
214 fn text_event(kind: AssistantBlockKind, delta: String) -> Option<AssistantEvent> {
215 if delta.is_empty() {
216 return None;
217 }
208 - Some(AssistantEvent::TextDelta { kind, delta })
218 + Some(AssistantEvent::TextDelta {
219 + kind,
220 + delta,
221 + token_count: None,
222 + })
223 }
224
225 /// Push one plain text delta if it is non-empty.
@@ -215,6 +229,22 @@ fn push_text_delta(events: &mut Vec<AssistantEvent>, kind: AssistantBlockKind, d
229 }
230 }
231
232 +/// Push one attributed delta, counting its tokens, if its text is non-empty.
233 +fn push_piece_delta(
234 + events: &mut Vec<AssistantEvent>,
235 + kind: AssistantBlockKind,
236 + piece: DecodedText,
237 +) {
238 + if piece.text.is_empty() {
239 + return;
240 + }
241 + events.push(AssistantEvent::TextDelta {
242 + kind,
243 + delta: piece.text,
244 + token_count: Some(piece.attributions.len()),
245 + });
246 +}
247 +
248 /// Wrap one decoded-text stream into the internal unified assistant stream.
249 #[try_stream]
250 pub(crate) async fn unified_event_stream(
@@ -263,7 +293,10 @@ pub(crate) async fn unified_event_stream(
293 y.yield_ok(next).await;
294 }
295 y.yield_ok(AssistantEvent::Done {
266 - usage: finished.usage,
296 + usage: ChatTokenUsage {
297 + engine: finished.usage,
298 + reasoning_tokens: state.reasoning_tokens,
299 + },
300 finish_reason: finished.finish_reason,
301 kv_transfer_params: finished.kv_transfer_params,
302 ec_transfer_params: finished.ec_transfer_params,
@@ -285,10 +318,12 @@ mod tests {
318 use vllm_parser::reasoning::ReasoningError;
319 use vllm_parser::tool::{Tool, ToolCallDelta};
320 use vllm_parser::unified::{Gemma4UnifiedParser, UnifiedParserError, UnifiedParserOutput};
321 + use vllm_text::DecodedText;
322 use vllm_tokenizer::test_utils::TestTokenizer;
323 + use vllm_tokenizer::{TokenAnchor, TokenAttribution};
324
290 - use super::{DecodedText, unified_event_stream};
291 - use crate::event::AssistantBlockKind;
325 + use super::unified_event_stream;
326 + use crate::event::{AssistantBlockKind, ChatTokenUsage};
327 use crate::output::AssistantEvent;
328
329 enum ScriptedStep {
@@ -303,6 +338,7 @@ mod tests {
338 steps: VecDeque<ScriptedStep>,
339 reset_text: String,
340 tool_call_id: Option<String>,
341 + finish_output: Option<UnifiedParserOutput>,
342 finish_error_reset_text: Option<String>,
343 }
344
@@ -312,10 +348,16 @@ mod tests {
348 steps: steps.into_iter().collect(),
349 reset_text: String::new(),
350 tool_call_id: Some("call_test".to_string()),
351 + finish_output: None,
352 finish_error_reset_text: None,
353 }
354 }
355
356 + fn with_finish_output(mut self, output: UnifiedParserOutput) -> Self {
357 + self.finish_output = Some(output);
358 + self
359 + }
360 +
361 fn with_finish_error(mut self, reset_text: &str) -> Self {
362 self.finish_error_reset_text = Some(reset_text.to_string());
363 self
@@ -371,7 +413,7 @@ mod tests {
413 },
414 ));
415 }
374 - Ok(UnifiedParserOutput::default())
416 + Ok(self.finish_output.take().unwrap_or_default())
417 }
418
419 fn reset(&mut self) -> String {
@@ -381,7 +423,7 @@ mod tests {
423
424 fn decoded_delta(delta: &str) -> vllm_text::output::DecodedTextEvent {
425 vllm_text::output::DecodedTextEvent::TextDelta {
384 - decoded: vllm_text::DecodedText::unattributed(delta),
426 + decoded: DecodedText::unattributed(delta),
427 sampled: vllm_text::SampledDelta::default(),
428 finished: None,
429 }
@@ -389,7 +431,7 @@ mod tests {
431
432 fn finished_delta(delta: &str) -> vllm_text::output::DecodedTextEvent {
433 vllm_text::output::DecodedTextEvent::TextDelta {
392 - decoded: vllm_text::DecodedText::unattributed(delta),
434 + decoded: DecodedText::unattributed(delta),
435 sampled: vllm_text::SampledDelta::default(),
436 finished: Some(Box::new(vllm_text::output::Finished {
437 usage: vllm_llm::TokenUsage::default(),
@@ -425,6 +467,22 @@ mod tests {
467 output
468 }
469
470 + /// One reasoning piece carrying `tokens` synthetic token attributions.
471 + fn attributed_reasoning(delta: &str, tokens: u32) -> UnifiedParserOutput {
472 + let attributions = (0..tokens)
473 + .map(|token_id| TokenAttribution {
474 + token_id,
475 + anchor: TokenAnchor::Visible { byte_offset: 0 },
476 + })
477 + .collect();
478 + let mut output = UnifiedParserOutput::default();
479 + output.push_reasoning(DecodedText {
480 + text: delta.to_string(),
481 + attributions,
482 + });
483 + output
484 + }
485 +
486 fn tool_call(name: &str, arguments: &str) -> UnifiedParserOutput {
487 UnifiedParserOutput {
488 events: vec![vllm_parser::unified::UnifiedParserEvent::ToolCall(
@@ -455,6 +513,15 @@ mod tests {
513 output
514 }
515
516 + /// Terminal usage with a default engine-level usage and the given final
517 + /// reasoning token count.
518 + fn done_usage(reasoning_tokens: usize) -> ChatTokenUsage {
519 + ChatTokenUsage {
520 + engine: vllm_llm::TokenUsage::default(),
521 + reasoning_tokens,
522 + }
523 + }
524 +
525 #[tokio::test]
526 async fn unified_stream_parses_formatted_tool_call_without_latch() {
527 use vllm_parser::tool::{HermesToolParser, ToolParser as _};
@@ -519,6 +586,7 @@ mod tests {
586 vec![AssistantEvent::TextDelta {
587 kind: AssistantBlockKind::Reasoning,
588 delta: "thinking".to_string(),
589 + token_count: Some(0),
590 }]
591 );
592 }
@@ -565,6 +633,7 @@ mod tests {
633 AssistantEvent::TextDelta {
634 kind: AssistantBlockKind::Reasoning,
635 delta: "thinking".to_string(),
636 + token_count: Some(0),
637 },
638 AssistantEvent::ToolCallStart {
639 id: "call_test".to_string(),
@@ -594,6 +663,7 @@ mod tests {
663 AssistantEvent::TextDelta {
664 kind: AssistantBlockKind::Text,
665 delta: "visible ".to_string(),
666 + token_count: None,
667 },
668 AssistantEvent::ToolCallStart {
669 id: "call_test".to_string(),
@@ -633,6 +703,7 @@ mod tests {
703 AssistantEvent::TextDelta {
704 kind: AssistantBlockKind::Text,
705 delta: " done".to_string(),
706 + token_count: None,
707 },
708 ]
709 );
@@ -655,14 +726,17 @@ mod tests {
726 AssistantEvent::TextDelta {
727 kind: AssistantBlockKind::Text,
728 delta: "committed".to_string(),
729 + token_count: None,
730 },
731 AssistantEvent::TextDelta {
732 kind: AssistantBlockKind::Text,
733 delta: "buffered".to_string(),
734 + token_count: None,
735 },
736 AssistantEvent::TextDelta {
737 kind: AssistantBlockKind::Text,
738 delta: "later".to_string(),
739 + token_count: None,
740 },
741 ]
742 );
@@ -683,9 +757,134 @@ mod tests {
757 AssistantEvent::TextDelta {
758 kind: AssistantBlockKind::Text,
759 delta: "buffered".to_string(),
760 + token_count: None,
761 + },
762 + AssistantEvent::Done {
763 + usage: done_usage(0),
764 + finish_reason: crate::FinishReason::Stop(None),
765 + kv_transfer_params: None,
766 + ec_transfer_params: None,
767 + },
768 + ]
769 + );
770 + }
771 +
772 + #[tokio::test]
773 + async fn unified_stream_counts_reasoning_tokens_across_deltas() {
774 + let events = collect(
775 + ScriptedParser::new([
776 + ScriptedStep::Output(attributed_reasoning("think", 2)),
777 + ScriptedStep::Output(combined(text("visible"), attributed_reasoning("more", 3))),
778 + ]),
779 + vec![decoded_delta("a"), finished_delta("b")],
780 + )
781 + .await;
782 +
783 + assert_eq!(
784 + events,
785 + vec![
786 + AssistantEvent::TextDelta {
787 + kind: AssistantBlockKind::Reasoning,
788 + delta: "think".to_string(),
789 + token_count: Some(2),
790 + },
791 + // Non-reasoning deltas carry no token count.
792 + AssistantEvent::TextDelta {
793 + kind: AssistantBlockKind::Text,
794 + delta: "visible".to_string(),
795 + token_count: None,
796 + },
797 + AssistantEvent::TextDelta {
798 + kind: AssistantBlockKind::Reasoning,
799 + delta: "more".to_string(),
800 + token_count: Some(3),
801 + },
802 + AssistantEvent::Done {
803 + usage: done_usage(5),
804 + finish_reason: crate::FinishReason::Stop(None),
805 + kv_transfer_params: None,
806 + ec_transfer_params: None,
807 + },
808 + ]
809 + );
810 + }
811 +
812 + #[tokio::test]
813 + async fn unified_stream_counts_reasoning_emitted_at_finish() {
814 + let events = collect(
815 + ScriptedParser::new([ScriptedStep::Output(text("visible"))])
816 + .with_finish_output(attributed_reasoning("tail", 2)),
817 + vec![finished_delta("raw")],
818 + )
819 + .await;
820 +
821 + assert_eq!(
822 + events,
823 + vec![
824 + AssistantEvent::TextDelta {
825 + kind: AssistantBlockKind::Text,
826 + delta: "visible".to_string(),
827 + token_count: None,
828 + },
829 + AssistantEvent::TextDelta {
830 + kind: AssistantBlockKind::Reasoning,
831 + delta: "tail".to_string(),
832 + token_count: Some(2),
833 + },
834 + AssistantEvent::Done {
835 + usage: done_usage(2),
836 + finish_reason: crate::FinishReason::Stop(None),
837 + kv_transfer_params: None,
838 + ec_transfer_params: None,
839 + },
840 + ]
841 + );
842 + }
843 +
844 + #[tokio::test]
845 + async fn unified_stream_fallback_keeps_emitted_reasoning_count() {
846 + let events = collect(
847 + ScriptedParser::new([
848 + ScriptedStep::Output(attributed_reasoning("think", 2)),
849 + ScriptedStep::Error {
850 + committed: attributed_reasoning("more", 3),
851 + reset_text: "buffered".to_string(),
852 + },
853 + ]),
854 + vec![
855 + decoded_delta("a"),
856 + decoded_delta("bad"),
857 + finished_delta("later"),
858 + ],
859 + )
860 + .await;
861 +
862 + assert_eq!(
863 + events,
864 + vec![
865 + AssistantEvent::TextDelta {
866 + kind: AssistantBlockKind::Reasoning,
867 + delta: "think".to_string(),
868 + token_count: Some(2),
869 + },
870 + // Committed reasoning still counts before the failure fallback.
871 + AssistantEvent::TextDelta {
872 + kind: AssistantBlockKind::Reasoning,
873 + delta: "more".to_string(),
874 + token_count: Some(3),
875 + },
876 + AssistantEvent::TextDelta {
877 + kind: AssistantBlockKind::Text,
878 + delta: "buffered".to_string(),
879 + token_count: None,
880 + },
881 + AssistantEvent::TextDelta {
882 + kind: AssistantBlockKind::Text,
883 + delta: "later".to_string(),
884 + token_count: None,
885 },
886 AssistantEvent::Done {
688 - usage: vllm_llm::TokenUsage::default(),
887 + usage: done_usage(5),
888 finish_reason: crate::FinishReason::Stop(None),
889 kv_transfer_params: None,
890 ec_transfer_params: None,
@@ -727,9 +926,10 @@ mod tests {
926 kind: AssistantBlockKind::Text,
927 delta: "<|tool_call>call:write_file{content:<|\"|>hello world<|\"|>"
928 .to_string(),
929 + token_count: None,
930 },
931 AssistantEvent::Done {
732 - usage: vllm_llm::TokenUsage::default(),
932 + usage: done_usage(0),
933 finish_reason: crate::FinishReason::Stop(None),
934 kv_transfer_params: None,
935 ec_transfer_params: None,
rust/src/chat/src/output/harmony/mod.rs
+7 -1
@@ -259,6 +259,9 @@ impl HarmonyState {
259 events.push(AssistantEvent::TextDelta {
260 kind,
261 delta: "\n".to_string(),
262 + // TODO: measure reasoning tokens from native
263 + // analysis-channel token counts.
264 + token_count: None,
265 });
266 }
267
@@ -268,6 +271,9 @@ impl HarmonyState {
271 events.push(AssistantEvent::TextDelta {
272 kind,
273 delta: group.text,
274 + // TODO: measure reasoning tokens from native
275 + // analysis-channel token counts.
276 + token_count: None,
277 });
278 return;
279 }
@@ -373,7 +379,7 @@ async fn harmony_assistant_event_stream(
379
380 if let Some(finished) = finished {
381 y.yield_ok(AssistantEvent::Done {
376 - usage: finished.usage,
382 + usage: finished.usage.into(),
383 finish_reason: finished.finish_reason,
384 kv_transfer_params: finished.kv_transfer_params,
385 ec_transfer_params: finished.ec_transfer_params,
rust/src/chat/src/output/harmony/tests.rs
+5 -5
@@ -13,7 +13,7 @@ use vllm_text::output::{
13 use super::*;
14 use crate::output::ChatOutputProcessor;
15 use crate::request::{ChatRequest, ChatTool, ChatToolChoice, ResolvedToolContext};
16 -use crate::{AssistantMessageExt, ChatEvent, FinishReason};
16 +use crate::{AssistantMessageExt, ChatEvent, ChatTokenUsage, FinishReason};
17
18 fn assistant_prefix() -> Vec<u32> {
19 harmony_encoding()
@@ -125,11 +125,11 @@ fn interrupted_final_message_is_preserved() {
125 text: "hello".to_string(),
126 }],
127 },
128 - usage: vllm_llm::TokenUsage {
128 + usage: ChatTokenUsage::from(vllm_llm::TokenUsage {
129 prompt_token_count: 0,
130 output_token_count: 0,
131 cached_token_count: 0,
132 - },
132 + }),
133 finish_reason: FinishReason::stop_eos(),
134 kv_transfer_params: None,
135 ec_transfer_params: None,
@@ -176,11 +176,11 @@ fn interrupted_analysis_message_is_preserved() {
176 text: "think".to_string(),
177 }],
178 },
179 - usage: vllm_llm::TokenUsage {
179 + usage: ChatTokenUsage::from(vllm_llm::TokenUsage {
180 prompt_token_count: 0,
181 output_token_count: 0,
182 cached_token_count: 0,
183 - },
183 + }),
184 finish_reason: FinishReason::stop_eos(),
185 kv_transfer_params: None,
186 ec_transfer_params: None,
rust/src/chat/src/output/mod.rs
+4 -3
@@ -7,12 +7,11 @@ use std::sync::Arc;
7 use futures::Stream;
8 use trait_set::trait_set;
9 use uuid::Uuid;
10 -use vllm_llm::TokenUsage;
10 use vllm_text::output::{DecodedLogprobs, DecodedPromptLogprobs, DecodedTextEvent};
11
12 use crate::FinishReason;
13 use crate::error::Result;
15 -use crate::event::{AssistantBlockKind, ChatEvent};
14 +use crate::event::{AssistantBlockKind, ChatEvent, ChatTokenUsage};
15
16 mod default;
17 mod harmony;
@@ -35,6 +34,8 @@ pub(crate) enum AssistantEvent {
34 TextDelta {
35 kind: AssistantBlockKind,
36 delta: String,
37 + /// Number of generated tokens attributed to this delta, when available.
38 + token_count: Option<usize>,
39 },
40 /// Per-decoded-update sample metadata: logprobs and/or output token IDs.
41 LogprobsDelta {
@@ -47,7 +48,7 @@ pub(crate) enum AssistantEvent {
48 /// `ToolCallStart`.
49 ToolCallArgumentsDelta { delta: String },
50 Done {
50 - usage: TokenUsage,
51 + usage: ChatTokenUsage,
52 finish_reason: FinishReason,
53 /// Connector-specific KV transfer parameters for disaggregated serving.
54 kv_transfer_params: Option<serde_json::Value>,
rust/src/chat/src/output/structured.rs
+81 -16
@@ -16,6 +16,7 @@ use super::{AssistantEvent, AssistantEventStream};
16 use crate::error::Error;
17 use crate::event::{
18 AssistantBlockKind, AssistantContentBlock, AssistantMessage, AssistantToolCall, ChatEvent,
19 + ChatTokenUsage,
20 };
21 use crate::{FinishReason, Result};
22
@@ -80,10 +81,11 @@ impl StructuredEventState {
81 &mut self,
82 kind: AssistantBlockKind,
83 delta: String,
84 + token_count: Option<usize>,
85 ) -> Result<Vec<ChatEvent>> {
86 let mut events = Vec::new();
87 self.close_open_tool_call(&mut events);
86 - self.push_text_delta(kind, delta, &mut events);
88 + self.push_text_delta(kind, delta, token_count, &mut events);
89 Ok(events)
90 }
91
@@ -146,7 +148,7 @@ impl StructuredEventState {
148 /// Close any open block and emit the terminal `Done` event.
149 fn finish(
150 &mut self,
149 - usage: vllm_llm::TokenUsage,
151 + usage: ChatTokenUsage,
152 finish_reason: FinishReason,
153 kv_transfer_params: Option<serde_json::Value>,
154 ec_transfer_params: Option<serde_json::Value>,
@@ -170,6 +172,7 @@ impl StructuredEventState {
172 &mut self,
173 kind: AssistantBlockKind,
174 delta: String,
175 + token_count: Option<usize>,
176 events: &mut Vec<ChatEvent>,
177 ) {
178 if delta.is_empty() {
@@ -184,6 +187,7 @@ impl StructuredEventState {
187 index: open_block.index,
188 kind,
189 delta,
190 + token_count,
191 });
192 }
193 // Otherwise, close the currently open block (if any) and start a
@@ -197,7 +201,12 @@ impl StructuredEventState {
201 text: delta.clone(),
202 });
203 events.push(ChatEvent::BlockStart { index, kind });
200 - events.push(ChatEvent::BlockDelta { index, kind, delta });
204 + events.push(ChatEvent::BlockDelta {
205 + index,
206 + kind,
207 + delta,
208 + token_count,
209 + });
210 }
211 }
212 }
@@ -274,8 +283,12 @@ pub(crate) async fn structured_chat_event_stream(
283 })
284 .await;
285 }
277 - AssistantEvent::TextDelta { kind, delta } => {
278 - for next in state.process_text_delta(kind, delta)? {
286 + AssistantEvent::TextDelta {
287 + kind,
288 + delta,
289 + token_count,
290 + } => {
291 + for next in state.process_text_delta(kind, delta, token_count)? {
292 y.yield_ok(next).await;
293 }
294 }
@@ -321,7 +334,7 @@ mod tests {
334 use super::structured_chat_event_stream;
335 use crate::FinishReason;
336 use crate::error::Error;
324 - use crate::event::{AssistantBlockKind, AssistantMessageExt as _, ChatEvent};
337 + use crate::event::{AssistantBlockKind, AssistantMessageExt as _, ChatEvent, ChatTokenUsage};
338 use crate::output::AssistantEvent;
339
340 #[tokio::test]
@@ -335,11 +348,11 @@ mod tests {
348 delta: r#"{"city":"Paris"}"#.to_string(),
349 }),
350 Ok(AssistantEvent::Done {
338 - usage: vllm_llm::TokenUsage {
351 + usage: ChatTokenUsage::from(vllm_llm::TokenUsage {
352 prompt_token_count: 1,
353 output_token_count: 1,
354 cached_token_count: 0,
342 - },
355 + }),
356 finish_reason: FinishReason::stop_eos(),
357 kv_transfer_params: None,
358 ec_transfer_params: None,
@@ -390,11 +403,11 @@ mod tests {
403 delta: r#"{"b":2}"#.to_string(),
404 }),
405 Ok(AssistantEvent::Done {
393 - usage: vllm_llm::TokenUsage {
406 + usage: ChatTokenUsage::from(vllm_llm::TokenUsage {
407 prompt_token_count: 1,
408 output_token_count: 1,
409 cached_token_count: 0,
397 - },
410 + }),
411 finish_reason: FinishReason::stop_eos(),
412 kv_transfer_params: None,
413 ec_transfer_params: None,
@@ -433,6 +446,7 @@ mod tests {
446 Ok(AssistantEvent::TextDelta {
447 kind: AssistantBlockKind::Text,
448 delta: "before".to_string(),
449 + token_count: None,
450 }),
451 Ok(AssistantEvent::ToolCallStart {
452 id: "call_1".to_string(),
@@ -442,11 +456,11 @@ mod tests {
456 delta: r#"{"city":"Paris"}"#.to_string(),
457 }),
458 Ok(AssistantEvent::Done {
445 - usage: vllm_llm::TokenUsage {
459 + usage: ChatTokenUsage::from(vllm_llm::TokenUsage {
460 prompt_token_count: 1,
461 output_token_count: 1,
462 cached_token_count: 0,
449 - },
463 + }),
464 finish_reason: FinishReason::stop_eos(),
465 kv_transfer_params: None,
466 ec_transfer_params: None,
@@ -492,13 +506,14 @@ mod tests {
506 Ok(AssistantEvent::TextDelta {
507 kind: AssistantBlockKind::Text,
508 delta: "done".to_string(),
509 + token_count: None,
510 }),
511 Ok(AssistantEvent::Done {
497 - usage: vllm_llm::TokenUsage {
512 + usage: ChatTokenUsage::from(vllm_llm::TokenUsage {
513 prompt_token_count: 1,
514 output_token_count: 1,
515 cached_token_count: 0,
501 - },
516 + }),
517 finish_reason: FinishReason::stop_eos(),
518 kv_transfer_params: None,
519 ec_transfer_params: None,
@@ -562,11 +577,11 @@ mod tests {
577 delta: r#"{"b":2}"#.to_string(),
578 }),
579 Ok(AssistantEvent::Done {
565 - usage: vllm_llm::TokenUsage {
580 + usage: ChatTokenUsage::from(vllm_llm::TokenUsage {
581 prompt_token_count: 1,
582 output_token_count: 1,
583 cached_token_count: 0,
569 - },
584 + }),
585 finish_reason: FinishReason::stop_eos(),
586 kv_transfer_params: None,
587 ec_transfer_params: None,
@@ -596,4 +611,54 @@ mod tests {
611 assert_eq!(tool_calls.len(), 1);
612 assert_eq!(tool_calls[0].name, "first");
613 }
614 +
615 + #[tokio::test]
616 + async fn structured_stream_forwards_reasoning_token_count() {
617 + let events = stream::iter(vec![
618 + Ok(AssistantEvent::Start {
619 + prompt_token_ids: vec![].into(),
620 + prompt_logprobs: None,
621 + }),
622 + Ok(AssistantEvent::TextDelta {
623 + kind: AssistantBlockKind::Reasoning,
624 + delta: "think".to_string(),
625 + token_count: Some(3),
626 + }),
627 + Ok(AssistantEvent::Done {
628 + usage: ChatTokenUsage {
629 + engine: vllm_llm::TokenUsage {
630 + prompt_token_count: 1,
631 + output_token_count: 4,
632 + cached_token_count: 0,
633 + },
634 + reasoning_tokens: 3,
635 + },
636 + finish_reason: FinishReason::stop_eos(),
637 + kv_transfer_params: None,
638 + ec_transfer_params: None,
639 + }),
640 + ]);
641 +
642 + let events = structured_chat_event_stream(events, true)
643 + .collect::<Vec<_>>()
644 + .await
645 + .into_iter()
646 + .collect::<crate::Result<Vec<_>>>()
647 + .unwrap();
648 +
649 + assert!(matches!(events[0], ChatEvent::Start { .. }));
650 + assert!(matches!(
651 + events[2],
652 + ChatEvent::BlockDelta {
653 + kind: AssistantBlockKind::Reasoning,
654 + token_count: Some(3),
655 + ..
656 + }
657 + ));
658 + let ChatEvent::Done { usage, .. } = &events[4] else {
659 + panic!("expected done");
660 + };
661 + assert_eq!(usage.reasoning_tokens, 3);
662 + assert_eq!(usage.output_token_count, 4);
663 + }
664 }
rust/src/chat/src/stream.rs
+7 -7
@@ -11,7 +11,7 @@ use vllm_text::{DecodedLogprobs, DecodedPositionLogprobs, DecodedPromptLogprobs}
11
12 use crate::FinishReason;
13 use crate::error::{Error, Result};
14 -use crate::event::{AssistantContentBlock, AssistantMessage, ChatEvent};
14 +use crate::event::{AssistantContentBlock, AssistantMessage, ChatEvent, ChatTokenUsage};
15
16 /// Final structured assistant message plus terminal stream metadata.
17 #[derive(Debug, Clone, PartialEq)]
@@ -21,7 +21,7 @@ pub struct CollectedAssistantMessage {
21 pub prompt_logprobs: Option<DecodedPromptLogprobs>,
22 pub logprobs: Option<DecodedLogprobs>,
23 pub token_ids: Vec<u32>,
24 - pub usage: vllm_llm::TokenUsage,
24 + pub usage: ChatTokenUsage,
25 pub finish_reason: FinishReason,
26 /// Connector-specific KV transfer parameters for disaggregated serving.
27 pub kv_transfer_params: Option<serde_json::Value>,
@@ -140,7 +140,7 @@ mod tests {
140
141 use super::{ChatEventStream, CollectedAssistantMessage};
142 use crate::error::Error;
143 - use crate::event::ChatEvent;
143 + use crate::event::{ChatEvent, ChatTokenUsage};
144
145 #[tokio::test]
146 async fn collect_message_requires_terminal_done_event() {
@@ -195,11 +195,11 @@ mod tests {
195 }),
196 Ok(ChatEvent::Done {
197 message: Default::default(),
198 - usage: vllm_llm::TokenUsage {
198 + usage: ChatTokenUsage::from(vllm_llm::TokenUsage {
199 prompt_token_count: 2,
200 output_token_count: 1,
201 cached_token_count: 0,
202 - },
202 + }),
203 finish_reason: FinishReason::stop_eos(),
204 kv_transfer_params: None,
205 ec_transfer_params: None,
@@ -236,11 +236,11 @@ mod tests {
236 }],
237 }),
238 token_ids: vec![],
239 - usage: vllm_llm::TokenUsage {
239 + usage: ChatTokenUsage::from(vllm_llm::TokenUsage {
240 prompt_token_count: 2,
241 output_token_count: 1,
242 cached_token_count: 0,
243 - },
243 + }),
244 finish_reason: FinishReason::stop_eos(),
245 kv_transfer_params: None,
246 ec_transfer_params: None,
rust/src/chat/tests/chat.rs
+15
@@ -393,6 +393,7 @@ async fn chat_streams_text_events() {
393 index: 0,
394 kind: AssistantBlockKind::Text,
395 delta: "H".to_string(),
396 + token_count: None,
397 }
398 );
399 assert_eq!(
@@ -401,6 +402,7 @@ async fn chat_streams_text_events() {
402 index: 0,
403 kind: AssistantBlockKind::Text,
404 delta: "i".to_string(),
405 + token_count: None,
406 }
407 );
408 assert_eq!(
@@ -500,6 +502,7 @@ async fn chat_stream_waits_for_complete_utf8_before_emitting() {
502 index: 0,
503 kind: AssistantBlockKind::Text,
504 delta: "你".to_string(),
505 + token_count: None,
506 }
507 );
508 assert_eq!(
@@ -586,6 +589,7 @@ async fn chat_stream_flushes_held_text_on_finish() {
589 index: 0,
590 kind: AssistantBlockKind::Text,
591 delta: "ok st".to_string(),
592 + token_count: None,
593 }
594 );
595 assert_eq!(
@@ -773,6 +777,7 @@ async fn chat_stream_preserves_terminal_stop_token_when_requested() {
777 index: 0,
778 kind: AssistantBlockKind::Text,
779 delta: "Hi!".to_string(),
780 + token_count: None,
781 }
782 );
783 assert_eq!(
@@ -880,6 +885,7 @@ async fn chat_stream_separates_reasoning_blocks_automatically() {
885 index: 0,
886 kind: AssistantBlockKind::Reasoning,
887 delta: "reason ".to_string(),
888 + token_count: Some(7),
889 }
890 );
891 assert_eq!(
@@ -888,6 +894,7 @@ async fn chat_stream_separates_reasoning_blocks_automatically() {
894 index: 0,
895 kind: AssistantBlockKind::Reasoning,
896 delta: "more".to_string(),
897 + token_count: Some(4),
898 }
899 );
900 assert_eq!(
@@ -912,6 +919,7 @@ async fn chat_stream_separates_reasoning_blocks_automatically() {
919 index: 1,
920 kind: AssistantBlockKind::Text,
921 delta: "answer".to_string(),
922 + token_count: None,
923 }
924 );
925 assert_eq!(
@@ -927,12 +935,15 @@ async fn chat_stream_separates_reasoning_blocks_automatically() {
935 match next_semantic(&mut stream).await {
936 Some(Ok(ChatEvent::Done {
937 message,
938 + usage,
939 finish_reason,
940 ..
941 })) => {
942 assert_eq!(message.reasoning().unwrap(), "reason more");
943 assert_eq!(message.text(), "answer");
944 assert_eq!(finish_reason, FinishReason::Length);
945 + // 7 tokens for "reason " plus 4 for "more"; markers are excluded.
946 + assert_eq!(usage.reasoning_tokens, 11);
947 }
948 other => panic!("unexpected final event: {other:?}"),
949 }
@@ -1000,6 +1011,8 @@ async fn chat_collectors_return_structured_message_and_visible_text() {
1011 message.usage.output_token_count,
1012 "<think>inner</think>outer".len()
1013 );
1014 + // 5 tokens for "inner"; markers are excluded.
1015 + assert_eq!(message.usage.reasoning_tokens, 5);
1016
1017 let _ = shutdown_tx.send(());
1018 engine_task.await.unwrap();
@@ -1079,6 +1092,7 @@ async fn chat_explicitly_disables_reasoning_parser() {
1092 assert_eq!(message.message.reasoning(), None);
1093 assert_eq!(message.message.text(), "<think>reason more</think>answer");
1094 assert_eq!(message.finish_reason, FinishReason::Length);
1095 + assert_eq!(message.usage.reasoning_tokens, 0);
1096
1097 let _ = shutdown_tx.send(());
1098 engine_task.await.unwrap();
@@ -1379,6 +1393,7 @@ async fn chat_stream_and_collect_preserve_prompt_and_sample_logprobs() {
1393 index: 0,
1394 kind: AssistantBlockKind::Text,
1395 delta: "H".to_string(),
1396 + token_count: None,
1397 }
1398 );
1399 assert_eq!(
rust/src/server/src/routes/openai/chat_completions.rs
+250 -34
@@ -324,7 +324,20 @@ async fn chat_completion_chunk_stream(
324 ));
325 }
326 }
327 - Ok(ChatEvent::BlockDelta { kind, delta, .. }) => {
327 + Ok(ChatEvent::BlockDelta {
328 + kind,
329 + delta,
330 + token_count,
331 + ..
332 + }) => {
333 + // The count is per-delta and follows the decoder clock; it is
334 + // fed even when the reasoning delta itself is hidden from the
335 + // client (`include_reasoning=false`).
336 + if matches!(kind, AssistantBlockKind::Reasoning)
337 + && let Some(token_count) = token_count
338 + {
339 + continuous_usage.add_reasoning_tokens(token_count);
340 + }
341 let include_delta =
342 include_reasoning || !matches!(kind, AssistantBlockKind::Reasoning);
343 if include_delta {
@@ -433,6 +446,9 @@ async fn chat_completion_chunk_stream(
446 final_usage.prompt_token_count,
447 final_usage.output_token_count,
448 );
449 + // Converge the running reasoning count onto the authoritative
450 + // terminal value.
451 + continuous_usage.set_reasoning_tokens(final_usage.reasoning_tokens);
452
453 if let Some(pending_chunk) = pending_chunk.as_mut()
454 && let Some(chunk) = pending_chunk.take_chunk(&envelope)
@@ -806,16 +822,31 @@ mod tests {
822 use futures::{StreamExt as _, stream};
823 use serde_json::json;
824 use vllm_chat::{
809 - AssistantBlockKind, AssistantContentBlock, AssistantToolCall, ChatEvent, FinishReason,
825 + AssistantBlockKind, AssistantContentBlock, AssistantToolCall, ChatEvent, ChatTokenUsage,
826 + FinishReason,
827 };
828 use vllm_engine_core_client::protocol::output::StopReason;
829 use vllm_text::{DecodedLogprobs, DecodedPositionLogprobs, DecodedTokenLogprob};
830
831 use super::{
815 - ApiServerOptions, ResponseOptions, StreamResponseEnvelope, block_delta_chunk,
816 - chat_completion_chunk_stream, final_chunk,
832 + ApiServerOptions, ChatCompletionStreamResponse, ResponseOptions, StreamResponseEnvelope,
833 + block_delta_chunk, chat_completion_chunk_stream, final_chunk,
834 };
835
836 + /// Terminal usage with the given engine-level counts and zero reasoning tokens.
837 + fn done_usage(
838 + prompt_tokens: usize,
839 + output_tokens: usize,
840 + cached_tokens: usize,
841 + ) -> ChatTokenUsage {
842 + vllm_llm::TokenUsage {
843 + prompt_token_count: prompt_tokens,
844 + output_token_count: output_tokens,
845 + cached_token_count: cached_tokens,
846 + }
847 + .into()
848 + }
849 +
850 fn stream_envelope() -> Arc<StreamResponseEnvelope> {
851 Arc::new(StreamResponseEnvelope::new(
852 "chatcmpl-1".to_string(),
@@ -913,6 +944,7 @@ mod tests {
944 index: 0,
945 kind: AssistantBlockKind::Text,
946 delta: "hi".to_string(),
947 + token_count: None,
948 }),
949 Ok(ChatEvent::LogprobsDelta {
950 logprobs: Some(DecodedLogprobs {
@@ -929,11 +961,7 @@ mod tests {
961 }),
962 Ok(ChatEvent::Done {
963 message: Default::default(),
932 - usage: vllm_llm::TokenUsage {
933 - prompt_token_count: 1,
934 - output_token_count: 1,
935 - cached_token_count: 1,
936 - },
964 + usage: done_usage(1, 1, 1),
965 finish_reason: FinishReason::stop_eos(),
966 kv_transfer_params: None,
967 ec_transfer_params: None,
@@ -995,6 +1023,7 @@ mod tests {
1023 index: 0,
1024 kind: AssistantBlockKind::Reasoning,
1025 delta: "think".to_string(),
1026 + token_count: None,
1027 }),
1028 Ok(ChatEvent::LogprobsDelta {
1029 logprobs: Some(DecodedLogprobs {
@@ -1011,11 +1040,7 @@ mod tests {
1040 }),
1041 Ok(ChatEvent::Done {
1042 message: Default::default(),
1014 - usage: vllm_llm::TokenUsage {
1015 - prompt_token_count: 1,
1016 - output_token_count: 1,
1017 - cached_token_count: 0,
1018 - },
1043 + usage: done_usage(1, 1, 0),
1044 finish_reason: FinishReason::stop_eos(),
1045 kv_transfer_params: None,
1046 ec_transfer_params: None,
@@ -1059,19 +1084,17 @@ mod tests {
1084 index: 0,
1085 kind: AssistantBlockKind::Reasoning,
1086 delta: "think".to_string(),
1087 + token_count: None,
1088 }),
1089 Ok(ChatEvent::BlockDelta {
1090 index: 1,
1091 kind: AssistantBlockKind::Text,
1092 delta: "answer".to_string(),
1093 + token_count: None,
1094 }),
1095 Ok(ChatEvent::Done {
1096 message: Default::default(),
1070 - usage: vllm_llm::TokenUsage {
1071 - prompt_token_count: 1,
1072 - output_token_count: 2,
1073 - cached_token_count: 0,
1074 - },
1097 + usage: done_usage(1, 2, 0),
1098 finish_reason: FinishReason::stop_eos(),
1099 kv_transfer_params: None,
1100 ec_transfer_params: None,
@@ -1115,6 +1138,7 @@ mod tests {
1138 index: 0,
1139 kind: AssistantBlockKind::Reasoning,
1140 delta: "think".to_string(),
1141 + token_count: None,
1142 }),
1143 Ok(ChatEvent::LogprobsDelta {
1144 logprobs: Some(DecodedLogprobs {
@@ -1133,6 +1157,7 @@ mod tests {
1157 index: 1,
1158 kind: AssistantBlockKind::Text,
1159 delta: "answer".to_string(),
1160 + token_count: None,
1161 }),
1162 Ok(ChatEvent::LogprobsDelta {
1163 logprobs: Some(DecodedLogprobs {
@@ -1149,11 +1174,7 @@ mod tests {
1174 }),
1175 Ok(ChatEvent::Done {
1176 message: Default::default(),
1152 - usage: vllm_llm::TokenUsage {
1153 - prompt_token_count: 1,
1154 - output_token_count: 2,
1155 - cached_token_count: 0,
1156 - },
1177 + usage: done_usage(1, 2, 0),
1178 finish_reason: FinishReason::stop_eos(),
1179 kv_transfer_params: None,
1180 ec_transfer_params: None,
@@ -1226,6 +1247,7 @@ mod tests {
1247 index: 0,
1248 kind: AssistantBlockKind::Reasoning,
1249 delta: "think".to_string(),
1250 + token_count: None,
1251 }),
1252 Ok(ChatEvent::LogprobsDelta {
1253 logprobs: Some(DecodedLogprobs {
@@ -1267,6 +1289,7 @@ mod tests {
1289 index: 1,
1290 kind: AssistantBlockKind::Text,
1291 delta: "answer".to_string(),
1292 + token_count: None,
1293 }),
1294 Ok(ChatEvent::LogprobsDelta {
1295 logprobs: Some(DecodedLogprobs {
@@ -1283,11 +1306,7 @@ mod tests {
1306 }),
1307 Ok(ChatEvent::Done {
1308 message: Default::default(),
1286 - usage: vllm_llm::TokenUsage {
1287 - prompt_token_count: 1,
1288 - output_token_count: 4,
1289 - cached_token_count: 0,
1290 - },
1309 + usage: done_usage(1, 4, 0),
1310 finish_reason: FinishReason::stop_eos(),
1311 kv_transfer_params: None,
1312 ec_transfer_params: None,
@@ -1365,11 +1384,7 @@ mod tests {
1384 }),
1385 Ok(ChatEvent::Done {
1386 message: Default::default(),
1368 - usage: vllm_llm::TokenUsage {
1369 - prompt_token_count: 1,
1370 - output_token_count: 1,
1371 - cached_token_count: 0,
1372 - },
1387 + usage: done_usage(1, 1, 0),
1388 finish_reason: FinishReason::stop_eos(),
1389 kv_transfer_params: None,
1390 ec_transfer_params: None,
@@ -1410,4 +1425,205 @@ mod tests {
1425 None
1426 );
1427 }
1428 +
1429 + /// Measured reasoning-token count of one chunk's continuous usage.
1430 + fn chunk_reasoning_tokens(chunk: &ChatCompletionStreamResponse) -> Option<usize> {
1431 + chunk
1432 + .usage
1433 + .as_ref()
1434 + .map(|usage| usage.completion_tokens_details.reasoning_tokens)
1435 + }
1436 +
1437 + #[tokio::test]
1438 + async fn chunk_stream_reports_reasoning_tokens_with_continuous_usage() {
1439 + let stream = stream::iter(vec![
1440 + Ok(ChatEvent::Start {
1441 + prompt_token_ids: vec![7, 8, 9].into(),
1442 + prompt_logprobs: None,
1443 + }),
1444 + Ok(ChatEvent::BlockDelta {
1445 + index: 0,
1446 + kind: AssistantBlockKind::Reasoning,
1447 + delta: "think".to_string(),
1448 + token_count: Some(2),
1449 + }),
1450 + Ok(ChatEvent::LogprobsDelta {
1451 + logprobs: None,
1452 + token_ids: vec![20, 21],
1453 + }),
1454 + Ok(ChatEvent::BlockDelta {
1455 + index: 1,
1456 + kind: AssistantBlockKind::Text,
1457 + delta: "answer".to_string(),
1458 + // A stray text-side count must not pollute reasoning usage.
1459 + token_count: None,
1460 + }),
1461 + Ok(ChatEvent::Done {
1462 + message: Default::default(),
1463 + usage: ChatTokenUsage {
1464 + engine: vllm_llm::TokenUsage {
1465 + prompt_token_count: 3,
1466 + output_token_count: 3,
1467 + cached_token_count: 0,
1468 + },
1469 + reasoning_tokens: 2,
1470 + },
1471 + finish_reason: FinishReason::stop_eos(),
1472 + kv_transfer_params: None,
1473 + ec_transfer_params: None,
1474 + }),
1475 + ]);
1476 +
1477 + let chunks = chat_completion_chunk_stream(
1478 + stream,
1479 + "chatcmpl-1".to_string(),
1480 + "model".to_string(),
1481 + 1,
1482 + ApiServerOptions::default(),
1483 + ResponseOptions {
1484 + include_usage: true,
1485 + include_continuous_usage: true,
1486 + include_reasoning: true,
1487 + ..Default::default()
1488 + },
1489 + )
1490 + .collect::<Vec<_>>()
1491 + .await
1492 + .into_iter()
1493 + .collect::<Result<Vec<_>, _>>()
1494 + .expect("stream chunks");
1495 +
1496 + // Role, reasoning delta, text delta, finish-reason chunk, usage chunk.
1497 + assert_eq!(chunks.len(), 5);
1498 + // The role chunk reports reasoning tokens from 0 onwards.
1499 + let role_usage = chunks[0].usage.as_ref().expect("role chunk usage");
1500 + assert_eq!(role_usage.prompt_tokens, 3);
1501 + assert_eq!(role_usage.completion_tokens, Some(0));
1502 + assert_eq!(chunk_reasoning_tokens(&chunks[0]), Some(0));
1503 + // The reasoning delta contributes its per-delta measured count, fed by
1504 + // the `BlockDelta` itself rather than the sampled-clock `LogprobsDelta`.
1505 + assert_eq!(chunk_reasoning_tokens(&chunks[1]), Some(2));
1506 + assert_eq!(
1507 + chunks[1].usage.as_ref().and_then(|usage| usage.completion_tokens),
1508 + Some(2)
1509 + );
1510 + // The text delta's own count does not pollute reasoning usage, and the
1511 + // terminal usage chunk converges onto the authoritative final counts.
1512 + assert_eq!(chunk_reasoning_tokens(&chunks[2]), Some(2));
1513 + let final_usage = chunks[4].usage.as_ref().expect("final usage chunk");
1514 + assert_eq!(final_usage.prompt_tokens, 3);
1515 + assert_eq!(final_usage.completion_tokens, Some(3));
1516 + assert_eq!(chunk_reasoning_tokens(&chunks[4]), Some(2));
1517 + }
1518 +
1519 + #[tokio::test]
1520 + async fn chunk_stream_reports_reasoning_tokens_when_reasoning_hidden() {
1521 + let stream = stream::iter(vec![
1522 + Ok(ChatEvent::Start {
1523 + prompt_token_ids: vec![].into(),
1524 + prompt_logprobs: None,
1525 + }),
1526 + Ok(ChatEvent::BlockDelta {
1527 + index: 0,
1528 + kind: AssistantBlockKind::Reasoning,
1529 + delta: "think".to_string(),
1530 + token_count: Some(2),
1531 + }),
1532 + Ok(ChatEvent::Done {
1533 + message: Default::default(),
1534 + usage: ChatTokenUsage {
1535 + engine: vllm_llm::TokenUsage {
1536 + prompt_token_count: 1,
1537 + output_token_count: 2,
1538 + cached_token_count: 0,
1539 + },
1540 + reasoning_tokens: 2,
1541 + },
1542 + finish_reason: FinishReason::stop_eos(),
1543 + kv_transfer_params: None,
1544 + ec_transfer_params: None,
1545 + }),
1546 + ]);
1547 +
1548 + let chunks = chat_completion_chunk_stream(
1549 + stream,
1550 + "chatcmpl-1".to_string(),
1551 + "model".to_string(),
1552 + 1,
1553 + ApiServerOptions::default(),
1554 + ResponseOptions {
1555 + include_usage: true,
1556 + include_continuous_usage: true,
1557 + ..Default::default()
1558 + },
1559 + )
1560 + .collect::<Vec<_>>()
1561 + .await
1562 + .into_iter()
1563 + .collect::<Result<Vec<_>, _>>()
1564 + .expect("stream chunks");
1565 +
1566 + // The reasoning text stays hidden...
1567 + assert!(
1568 + chunks
1569 + .iter()
1570 + .all(|chunk| chunk.choices.iter().all(|choice| choice.delta.reasoning.is_none()))
1571 + );
1572 + // ...but the count is still reported, from 0 on the role chunk up to
1573 + // the converged final usage.
1574 + assert_eq!(chunk_reasoning_tokens(&chunks[0]), Some(0));
1575 + let final_usage = chunks.last().unwrap().usage.as_ref().expect("final usage chunk");
1576 + assert_eq!(final_usage.completion_tokens, Some(2));
1577 + assert_eq!(final_usage.completion_tokens_details.reasoning_tokens, 2);
1578 + }
1579 +
1580 + #[tokio::test]
1581 + async fn chunk_stream_reports_zero_reasoning_tokens_without_reasoning() {
1582 + let stream = stream::iter(vec![
1583 + Ok(ChatEvent::Start {
1584 + prompt_token_ids: vec![].into(),
1585 + prompt_logprobs: None,
1586 + }),
1587 + Ok(ChatEvent::BlockDelta {
1588 + index: 0,
1589 + kind: AssistantBlockKind::Text,
1590 + delta: "answer".to_string(),
1591 + token_count: None,
1592 + }),
1593 + Ok(ChatEvent::Done {
1594 + message: Default::default(),
1595 + usage: done_usage(1, 1, 0),
1596 + finish_reason: FinishReason::stop_eos(),
1597 + kv_transfer_params: None,
1598 + ec_transfer_params: None,
1599 + }),
1600 + ]);
1601 +
1602 + let chunks = chat_completion_chunk_stream(
1603 + stream,
1604 + "chatcmpl-1".to_string(),
1605 + "model".to_string(),
1606 + 1,
1607 + ApiServerOptions::default(),
1608 + ResponseOptions {
1609 + include_usage: true,
1610 + include_continuous_usage: true,
1611 + ..Default::default()
1612 + },
1613 + )
1614 + .collect::<Vec<_>>()
1615 + .await
1616 + .into_iter()
1617 + .collect::<Result<Vec<_>, _>>()
1618 + .expect("stream chunks");
1619 +
1620 + // Reasoning-token details are always reported; a stream without
1621 + // reasoning simply reports 0 on every usage payload.
1622 + for chunk in &chunks {
1623 + let usage = chunk.usage.as_ref().expect("usage");
1624 + assert_eq!(usage.completion_tokens_details.reasoning_tokens, 0);
1625 + let json = serde_json::to_value(usage).expect("usage serializes");
1626 + assert_eq!(json["completion_tokens_details"]["reasoning_tokens"], 0);
1627 + }
1628 + }
1629 }
rust/src/server/src/routes/openai/utils/types.rs
+73 -3
@@ -7,7 +7,7 @@ use std::slice;
7 use llm_multimodal::ImageDetail;
8 use serde::{Deserialize, Deserializer, Serialize};
9 use serde_json::Value;
10 -use vllm_llm::TokenUsage;
10 +use vllm_chat::ChatTokenUsage;
11
12 // ============================================================================
13 // Constants
@@ -440,14 +440,18 @@ pub struct Usage {
440 pub total_tokens: usize,
441 pub completion_tokens: Option<usize>,
442 pub prompt_tokens_details: Option<PromptTokenUsageInfo>,
443 + /// Reasoning-token breakdown.
444 + /// Always present, 0 when the configured parser has no reasoning channel.
445 + pub completion_tokens_details: CompletionTokenUsageInfo,
446 }
447
448 impl Usage {
446 - /// Create a Usage with prompt-token cache details.
449 + /// Create a Usage with prompt-token cache and reasoning-token details.
450 pub fn from_counts(
451 prompt_tokens: usize,
452 completion_tokens: usize,
453 cached_tokens: Option<usize>,
454 + reasoning_tokens: usize,
455 ) -> Self {
456 Self {
457 prompt_tokens,
@@ -456,14 +460,20 @@ impl Usage {
460 prompt_tokens_details: cached_tokens
461 .filter(|&c| c > 0)
462 .map(|c| PromptTokenUsageInfo { cached_tokens: c }),
463 + completion_tokens_details: CompletionTokenUsageInfo { reasoning_tokens },
464 }
465 }
466
462 - pub fn from_token_usage(usage: TokenUsage, enable_prompt_tokens_details: bool) -> Self {
467 + pub fn from_token_usage(
468 + usage: impl Into<ChatTokenUsage>,
469 + enable_prompt_tokens_details: bool,
470 + ) -> Self {
471 + let usage = usage.into();
472 Self::from_counts(
473 usage.prompt_token_count,
474 usage.output_token_count,
475 enable_prompt_tokens_details.then_some(usage.cached_token_count),
476 + usage.reasoning_tokens,
477 )
478 }
479 }
@@ -474,8 +484,15 @@ pub struct PromptTokenUsageInfo {
484 pub cached_tokens: usize,
485 }
486
487 +/// Mirrors the Python vLLM `CompletionTokenUsageInfo` class.
488 +#[derive(Debug, Clone, Default, Serialize)]
489 +pub struct CompletionTokenUsageInfo {
490 + pub reasoning_tokens: usize,
491 +}
492 +
493 #[cfg(test)]
494 mod usage_tests {
495 + use vllm_chat::ChatTokenUsage;
496 use vllm_llm::TokenUsage;
497
498 use super::Usage;
@@ -512,6 +529,59 @@ mod usage_tests {
529 Some(3)
530 );
531 }
532 +
533 + #[test]
534 + fn token_usage_includes_reasoning_token_details() {
535 + let usage = Usage::from_token_usage(
536 + ChatTokenUsage {
537 + engine: TokenUsage {
538 + prompt_token_count: 5,
539 + output_token_count: 4,
540 + cached_token_count: 0,
541 + },
542 + reasoning_tokens: 3,
543 + },
544 + false,
545 + );
546 +
547 + assert_eq!(usage.completion_tokens_details.reasoning_tokens, 3);
548 + let json = serde_json::to_value(&usage).expect("usage serializes");
549 + assert_eq!(json["completion_tokens_details"]["reasoning_tokens"], 3);
550 + }
551 +
552 + #[test]
553 + fn token_usage_serializes_zero_reasoning_tokens() {
554 + let usage = Usage::from_token_usage(
555 + ChatTokenUsage {
556 + engine: TokenUsage {
557 + prompt_token_count: 5,
558 + output_token_count: 0,
559 + cached_token_count: 0,
560 + },
561 + reasoning_tokens: 0,
562 + },
563 + false,
564 + );
565 +
566 + let json = serde_json::to_value(&usage).expect("usage serializes");
567 + assert_eq!(json["completion_tokens_details"]["reasoning_tokens"], 0);
568 + }
569 +
570 + #[test]
571 + fn token_usage_reports_zero_reasoning_tokens_without_reasoning_parser() {
572 + let usage = Usage::from_token_usage(
573 + TokenUsage {
574 + prompt_token_count: 5,
575 + output_token_count: 2,
576 + cached_token_count: 0,
577 + },
578 + false,
579 + );
580 +
581 + assert_eq!(usage.completion_tokens_details.reasoning_tokens, 0);
582 + let json = serde_json::to_value(&usage).expect("usage serializes");
583 + assert_eq!(json["completion_tokens_details"]["reasoning_tokens"], 0);
584 + }
585 }
586
587 /// OpenAI completions-style logprobs.
rust/src/server/src/routes/openai/utils/usage.rs
+19 -1
@@ -12,6 +12,8 @@ use super::types::Usage;
12 pub(crate) struct ContinuousUsage {
13 prompt_tokens: usize,
14 output_tokens: usize,
15 + /// Running reasoning-token count, following the decoder clock.
16 + reasoning_tokens: usize,
17 }
18
19 impl ContinuousUsage {
@@ -25,6 +27,17 @@ impl ContinuousUsage {
27 self.output_tokens = self.output_tokens.saturating_add(output_tokens);
28 }
29
30 + /// Add the tokens attributed to one reasoning delta to the running count.
31 + pub(crate) fn add_reasoning_tokens(&mut self, token_count: usize) {
32 + self.reasoning_tokens = self.reasoning_tokens.saturating_add(token_count);
33 + }
34 +
35 + /// Converge the running reasoning count onto the authoritative terminal
36 + /// value reported at stream end.
37 + pub(crate) fn set_reasoning_tokens(&mut self, reasoning_tokens: usize) {
38 + self.reasoning_tokens = reasoning_tokens;
39 + }
40 +
41 /// Replace the running counts with the final counts reported by generation.
42 pub(crate) fn set_final_counts(&mut self, prompt_tokens: usize, output_tokens: usize) {
43 self.prompt_tokens = prompt_tokens;
@@ -33,6 +46,11 @@ impl ContinuousUsage {
46
47 /// Build a streaming usage snapshot without prompt cache details.
48 pub(crate) fn to_usage(&self) -> Usage {
36 - Usage::from_counts(self.prompt_tokens, self.output_tokens, None)
49 + Usage::from_counts(
50 + self.prompt_tokens,
51 + self.output_tokens,
52 + None,
53 + self.reasoning_tokens,
54 + )
55 }
56 }
tests/entrypoints/openai/chat_completion/test_include_reasoning.py
+17 -1
@@ -5,7 +5,7 @@
5
6 Verifies that reasoning content is included by default and suppressed
7 when ``include_reasoning=False``, for both streaming and non-streaming
8 -Chat Completions.
8 +Chat Completions, while reasoning tokens remain represented in usage.
9 """
10
11 import openai
@@ -77,6 +77,22 @@ async def test_include_reasoning_false_non_streaming(client: openai.AsyncOpenAI)
77 assert msg.content, "Expected content in response even without reasoning"
78
79
80 +@pytest.mark.asyncio
81 +async def test_reasoning_tokens_in_usage(client: openai.AsyncOpenAI):
82 + """Hidden reasoning still contributes to terminal usage details."""
83 + response = await client.chat.completions.create(
84 + model=MODEL_NAME,
85 + messages=MESSAGES,
86 + max_tokens=64,
87 + extra_body={"include_reasoning": False},
88 + )
89 +
90 + assert response.usage is not None
91 + details = response.usage.completion_tokens_details
92 + assert details is not None
93 + assert details.reasoning_tokens > 0
94 +
95 +
96 @pytest.mark.asyncio
97 async def test_include_reasoning_true_streaming(client: openai.AsyncOpenAI):
98 """Default: reasoning deltas appear in streaming response."""
tests/entrypoints/openai/chat_completion/test_serving_chat.py
+2 -1
@@ -736,7 +736,8 @@ async def test_chat_per_request_metrics_follow_server_flag():
736 request_metadata=RequestResponseMetadata(request_id="chatcmpl-test-id"),
737 )
738 assert disabled_response.metrics is None
739 - assert disabled_response.usage.completion_tokens_details is None
739 + details = disabled_response.usage.completion_tokens_details
740 + assert details is None or details.reasoning_tokens == 0
741
742 enabled_serving = _build_minimal_metrics_serving_chat(
743 enable_per_request_metrics=True