fix(stargate-bench): report output usage from native mock streams

Honor include_usage with final generated-token counts and preserve the existing stream shape when usage is not requested. Verify the real mock-to-driver path with a Bazel test using declared executable inputs.

Refs: #1817

Author: barrygreengus

src/libraries/rust/stargate/crates/mock-dynamo/src/openai.rs

@@ -72,12 +72,19 @@ const CANARY_ANSWER: &str = "2";
 #[derive(Deserialize)]
 pub(crate) struct ChatRequest {
     pub(crate) stream: Option<bool>,
+    stream_options: Option<ChatStreamOptions>,
     pub(crate) model: Option<String>,
     pub(crate) max_tokens: Option<usize>,
     #[serde(default)]
     pub(crate) messages: Vec<serde_json::Value>,
 }
 
+#[derive(Deserialize)]
+struct ChatStreamOptions {
+    #[serde(default)]
+    include_usage: bool,
+}
+
 #[derive(Deserialize)]
 pub(crate) struct ResponsesRequest {
     pub(crate) stream: Option<bool>,
@@ -105,7 +112,10 @@ struct ChatCompletionChunk<'a> {
     id: &'a str,
     object: &'static str,
     model: &'a str,
-    choices: [ChunkChoice<'a>; 1],
+    choices: &'a [ChunkChoice<'a>],
+    // Omitted unless requested, then null until the final usage chunk.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    usage: Option<Option<ChatUsage>>,
 }
 
 #[derive(Serialize)]
@@ -174,7 +184,7 @@ struct StreamResponseConfig {
 
 #[derive(Clone, Copy, PartialEq, Eq)]
 enum StreamKind {
-    Chat { canary: bool },
+    Chat { canary: bool, include_usage: bool },
     Responses { created_at: u64 },
 }
 
@@ -251,7 +261,13 @@ pub(crate) async fn chat_completions(
             output_tokens,
             kv_cache_access,
             request_slot,
-            kind: StreamKind::Chat { canary },
+            kind: StreamKind::Chat {
+                canary,
+                include_usage: req
+                    .stream_options
+                    .as_ref()
+                    .is_some_and(|options| options.include_usage),
+            },
         });
     }
 
@@ -495,29 +511,52 @@ pub(crate) enum ChatStreamChunk<'a> {
     Role,
     Content(&'a str),
     Stop,
-}
-
-pub(crate) fn chat_chunk_json(id: &str, model: &str, chunk: ChatStreamChunk<'_>) -> String {
-    let (role, content, finish_reason) = match chunk {
-        ChatStreamChunk::Role => (Some("assistant"), None, None),
-        ChatStreamChunk::Content(content) => (None, Some(content), None),
-        ChatStreamChunk::Stop => (None, None, Some("stop")),
-    };
+    Usage {
+        input_tokens: usize,
+        output_tokens: usize,
+    },
+}
+
+pub(crate) fn chat_chunk_json(
+    id: &str,
+    model: &str,
+    chunk: ChatStreamChunk<'_>,
+    include_usage: bool,
+) -> String {
+    let mut usage = include_usage.then_some(None);
+    let choice = match chunk {
+        ChatStreamChunk::Role => Some((Some("assistant"), None, None)),
+        ChatStreamChunk::Content(content) => Some((None, Some(content), None)),
+        ChatStreamChunk::Stop => Some((None, None, Some("stop"))),
+        ChatStreamChunk::Usage {
+            input_tokens,
+            output_tokens,
+        } => {
+            usage = Some(Some(ChatUsage {
+                prompt_tokens: input_tokens,
+                completion_tokens: output_tokens,
+                total_tokens: input_tokens.saturating_add(output_tokens),
+            }));
+            None
+        }
+    }
+    .map(|(role, content, finish_reason)| ChunkChoice {
+        index: 0,
+        delta: Delta { role, content },
+        finish_reason,
+    });
     serde_json::to_string(&ChatCompletionChunk {
         id,
         object: "chat.completion.chunk",
         model,
-        choices: [ChunkChoice {
-            index: 0,
-            delta: Delta { role, content },
-            finish_reason,
-        }],
+        choices: choice.as_slice(),
+        usage,
     })
     .expect("chat stream event should serialize")
 }
 
-fn chat_sse_event(id: &str, model: &str, chunk: ChatStreamChunk<'_>) -> Event {
-    Event::default().data(chat_chunk_json(id, model, chunk))
+fn chat_sse_event(id: &str, model: &str, chunk: ChatStreamChunk<'_>, include_usage: bool) -> Event {
+    Event::default().data(chat_chunk_json(id, model, chunk, include_usage))
 }
 
 fn stream_response(config: StreamResponseConfig) -> Response {
@@ -557,22 +596,22 @@ fn stream_response(config: StreamResponseConfig) -> Response {
 
         state.emit_counters(&request_id, &model, input_tokens, 0, false);
 
-        if matches!(kind, StreamKind::Chat { .. }) {
-            yield Ok(chat_sse_event(&id, &model, ChatStreamChunk::Role));
+        if let StreamKind::Chat { include_usage, .. } = kind {
+            yield Ok(chat_sse_event(&id, &model, ChatStreamChunk::Role, include_usage));
         }
 
         for i in 0..output_tokens {
             if i > 0 {
                 tokio::time::sleep(token_delay(&state, &request_id, i)).await;
             }
-            let token = if matches!(kind, StreamKind::Chat { canary: true }) {
+            let token = if matches!(kind, StreamKind::Chat { canary: true, .. }) {
                 CANARY_ANSWER
             } else {
                 DUMMY_TOKENS[i % DUMMY_TOKENS.len()]
             };
             let event = match kind {
-                StreamKind::Chat { .. } => {
-                    chat_sse_event(&id, &model, ChatStreamChunk::Content(token))
+                StreamKind::Chat { include_usage, .. } => {
+                    chat_sse_event(&id, &model, ChatStreamChunk::Content(token), include_usage)
                 }
                 StreamKind::Responses { .. } => {
                     output_text.push_str(token);
@@ -593,7 +632,7 @@ fn stream_response(config: StreamResponseConfig) -> Response {
         }
 
         let completed = match kind {
-            StreamKind::Chat { .. } => chat_sse_event(&id, &model, ChatStreamChunk::Stop),
+            StreamKind::Chat { include_usage, .. } => chat_sse_event(&id, &model, ChatStreamChunk::Stop, include_usage),
             StreamKind::Responses { created_at } => responses_sse_event(
                 "response.completed",
                 &serde_json::json!({
@@ -628,7 +667,10 @@ fn stream_response(config: StreamResponseConfig) -> Response {
 
         state.emit_counters(&request_id, &model, input_tokens, output_tokens, true);
 
-        if matches!(kind, StreamKind::Chat { .. }) {
+        if let StreamKind::Chat { include_usage, .. } = kind {
+            if include_usage {
+                yield Ok(chat_sse_event(&id, &model, ChatStreamChunk::Usage { input_tokens, output_tokens }, true));
+            }
             yield Ok(Event::default().data("[DONE]"));
         }
     };

src/libraries/rust/stargate/crates/mock-dynamo/src/tests.rs

@@ -27,12 +27,10 @@ use axum::routing::{get, post, put};
 use tokio::io::{AsyncReadExt, AsyncWriteExt};
 
 fn request() -> ChatRequest {
-    ChatRequest {
-        stream: Some(true),
-        model: Some("dummy-model".to_string()),
-        max_tokens: Some(1),
-        messages: Vec::new(),
-    }
+    serde_json::from_value(serde_json::json!({
+        "stream": true, "model": "dummy-model", "max_tokens": 1, "messages": []
+    }))
+    .unwrap()
 }
 
 fn test_stats_events() -> broadcast::Sender<StatsStreamEvent> {
@@ -582,13 +580,80 @@ fn chat_stream_chunks_preserve_delta_and_finish_shapes() {
         ),
     ] {
         let value: serde_json::Value =
-            serde_json::from_str(&chat_chunk_json("id", "model", chunk)).unwrap();
+            serde_json::from_str(&chat_chunk_json("id", "model", chunk, false)).unwrap();
         assert_eq!(value["choices"][0]["delta"], delta);
         assert_eq!(value["choices"][0]["finish_reason"], finish_reason);
         assert!(value.get("usage").is_none());
     }
 }
 
+#[tokio::test]
+async fn streaming_chat_usage_reports_actual_output_only_when_requested() {
+    let state = AppState {
+        output_tokens: OutputTokenConfig {
+            min: 100,
+            max: 100,
+            distribution: OutputTokenDistribution::Uniform,
+        },
+        context_length_tokens: 5,
+        ..test_state()
+    };
+    let app = Router::new()
+        .route("/v1/chat/completions", post(chat_completions))
+        .with_state(state);
+    let (address, server) = spawn_test_app(app).await;
+    for include_usage in [None, Some(false), Some(true)] {
+        let mut body = serde_json::json!({
+            "model": "dummy-model", "messages": [], "stream": true, "max_tokens": 100
+        });
+        if let Some(include_usage) = include_usage {
+            body["stream_options"] = serde_json::json!({"include_usage": include_usage});
+        }
+        let response = json_response(
+            address,
+            "POST",
+            "/v1/chat/completions",
+            "connection: close\r\nx-input-tokens: 2\r\nx-output-tokens: 100",
+            &body.to_string(),
+        )
+        .await;
+        assert!(response.starts_with("HTTP/1.1 200 OK"));
+        let data: Vec<_> = response
+            .lines()
+            .filter_map(|line| line.strip_prefix("data: "))
+            .collect();
+        assert_eq!(data.last(), Some(&"[DONE]"));
+        let events: Vec<serde_json::Value> = data[..data.len() - 1]
+            .iter()
+            .map(|data| serde_json::from_str(data).unwrap())
+            .collect();
+        assert_eq!(
+            events
+                .iter()
+                .filter(|event| event["choices"][0]["delta"]["content"].is_string())
+                .count(),
+            3
+        );
+        if include_usage == Some(true) {
+            let (usage, output_events) = events.split_last().unwrap();
+            assert_eq!(usage["choices"], serde_json::json!([]));
+            assert_eq!(
+                usage["usage"],
+                serde_json::json!({"prompt_tokens": 2, "completion_tokens": 3, "total_tokens": 5})
+            );
+            assert!(
+                output_events
+                    .iter()
+                    .all(|event| event.get("usage") == Some(&serde_json::Value::Null))
+            );
+        } else {
+            assert!(events.iter().all(|event| event.get("usage").is_none()));
+        }
+    }
+    server.abort();
+    let _ = server.await;
+}
+
 #[tokio::test]
 async fn embeddings_endpoint_returns_json_without_stream() {
     let state = test_state();

src/libraries/rust/stargate/crates/stargate-bench/BUILD.bazel

@@ -4,6 +4,7 @@
 # Benchmark harness. Not packaged into an OCI image.
 
 load("@stargate_crates//:defs.bzl", "aliases", "all_crate_deps")
+load("@rules_python//python:defs.bzl", "py_test")
 load("@rules_rust//rust:defs.bzl", "rust_binary", "rust_test")
 
 # Workspace-local Cargo deps (own-workspace `workspace = true`).
@@ -39,3 +40,19 @@ rust_test(
     },
     deps = all_crate_deps(normal_dev = True),
 )
+
+py_test(
+    name = "benchmark_usage_test",
+    srcs = ["benchmark_usage_test.py"],
+    main = "benchmark_usage_test.py",
+    args = [
+        "$(rootpath //src/libraries/rust/stargate/crates/mock-dynamo:mock-dynamo)",
+        "$(rootpath :stargate-bench)",
+    ],
+    data = [
+        "//src/libraries/rust/stargate/crates/mock-dynamo:mock-dynamo",
+        ":stargate-bench",
+    ],
+    python_version = "3.11",
+    size = "small",
+)

src/libraries/rust/stargate/crates/stargate-bench/benchmark_usage_test.py

@@ -0,0 +1,90 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Exercise the real benchmark driver against its built-in mock backend."""
+
+import json
+import os
+from pathlib import Path
+import re
+import subprocess
+import sys
+import tempfile
+import time
+
+
+def test_usage(mock_binary, bench_binary):
+    with tempfile.TemporaryDirectory(dir=os.environ.get("TEST_TMPDIR")) as directory:
+        root = Path(directory)
+        server_log = root / "mock.log"
+        with server_log.open("w") as log:
+            server = subprocess.Popen(
+                [
+                    str(mock_binary), "--http-listen-addr", "127.0.0.1:0",
+                    "--model-name", "model-a", "--num-tokens", "100",
+                    "--context-length-tokens", "5", "--token-delay-ms", "1",
+                ],
+                env=dict(os.environ, RUST_LOG="info"),
+                stdout=log,
+                stderr=subprocess.STDOUT,
+            )
+            try:
+                deadline = time.monotonic() + 5
+                while True:
+                    output = server_log.read_text()
+                    address = re.search(r"http://127\.0\.0\.1:\d+/v1/chat/completions", output)
+                    if address:
+                        break
+                    if server.poll() is not None or time.monotonic() >= deadline:
+                        raise AssertionError(f"mock did not become ready:\n{output}")
+                    time.sleep(0.01)
+
+                manifest = root / "manifest.json"
+                manifest.write_text(json.dumps({
+                    "manifest_version": 1,
+                    "benchmark_name": "streamed-usage",
+                    "metadata": {},
+                    "model": "model-a",
+                    "seed": 1,
+                    "request_count": 1,
+                    "max_concurrency": 1,
+                    "stargate_count": 1,
+                    "backend_count": 1,
+                    "requests": [{
+                        "request_index": 0,
+                        "request_id": "usage-request",
+                        "scheduled_offset_ms": 0,
+                        "routing_key": None,
+                        "cache_affinity_key": None,
+                        "input_tokens": 2,
+                        "output_tokens": 100,
+                        "backend_behavior_class": "uniform",
+                    }],
+                }))
+                results = root / "results.jsonl"
+                driven = subprocess.run(
+                    [str(bench_binary), "drive", "--manifest", str(manifest),
+                     "--endpoint", address.group(), "--output", str(results)],
+                    capture_output=True, text=True, timeout=10,
+                )
+                assert driven.returncode == 0, driven.stdout + driven.stderr
+                result = json.loads(results.read_text())
+                assert result["ok"], result
+                assert result["output_tokens"] == 100, result
+                # The five-token context leaves three generated tokens after the prompt.
+                assert result["observed_output_tokens"] == 3, result
+                assert result["first_output_ms"] is not None, result
+            finally:
+                if server.poll() is None:
+                    server.terminate()
+                    try:
+                        server.wait(timeout=5)
+                    except subprocess.TimeoutExpired:
+                        server.kill()
+                        server.wait()
+
+
+if __name__ == "__main__":
+    if len(sys.argv) != 3:
+        raise SystemExit("usage: benchmark_usage_test.py MOCK_BINARY BENCH_BINARY")
+    test_usage(*(Path(argument).resolve() for argument in sys.argv[1:]))