fix(stargate-bench): measure validated stream output

NOTICE

@@ -484,4 +484,5 @@ The following third-party licenses are included in this repository:
     src/libraries/go/lib/vendor/sigs.k8s.io/structured-merge-diff/v6/LICENSE
     src/libraries/go/lib/vendor/sigs.k8s.io/yaml/LICENSE
     src/libraries/java/nv-boot-parent/NOTICE
+    src/libraries/rust/stargate/crates/stargate-bench/NOTICE
     src/uis/nvcf-ui/NOTICE

dependencies.md

@@ -1942,6 +1942,7 @@ Generated by `go run -C ./tools/collect-dependencies .`. Refresh: `go run -C ./t
 - `Rust`: `serde` (`serde 1`, `serde 1.0`, `serde =1.0.228`)
 - `Rust`: `serde-json` (`serde_json 1`, `serde_json 1.0`, `serde_json =1.0.150`)
 - `Rust`: `serde-with` (`serde_with 3`, `serde_with 3.0`)
+- `Rust`: `sse-core =0.2.3`
 - `Rust`: `tempfile` (`tempfile 3`, `tempfile 3.20.0`, `tempfile =3.27.0`)
 - `Rust`: `thiserror` (`thiserror 2`, `thiserror =2.0.18`)
 - `Rust`: `url` (`url 2`, `url =2.5.8`)

src/libraries/rust/stargate/Cargo.toml

@@ -62,6 +62,7 @@ serde = { version = "=1.0.228", features = ["derive"] }
 serde_json = "=1.0.150"
 serde_yaml_ng = "=0.10.0"
 sonic-rs = "=0.5.8"
+sse-core = { version = "=0.2.3", default-features = false, features = ["std"] }
 thiserror = "=2.0.18"
 tokio = { version = "=1.52.3", features = ["full"] }
 tokio-stream = "=0.1.18"

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,7 +4,8 @@
 # Benchmark harness. Not packaged into an OCI image.
 
 load("@stargate_crates//:defs.bzl", "aliases", "all_crate_deps")
-load("@rules_rust//rust:defs.bzl", "rust_binary")
+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`).
 _WORKSPACE_DEPS = [
@@ -27,3 +28,31 @@ rust_binary(
     visibility = ["//visibility:public"],
     deps = _WORKSPACE_DEPS + all_crate_deps(normal = True),
 )
+
+rust_test(
+    name = "stargate-bench_test",
+    crate = ":stargate-bench",
+    data = ["//src/libraries/rust/stargate:benches"],
+    # Scenario lookup uses the manifest's ancestors. Resolve it within the
+    # test's runfiles tree instead of the build sandbox's source directory.
+    rustc_env = {
+        "CARGO_MANIFEST_DIR": "src/libraries/rust/stargate/crates/stargate-bench",
+    },
+    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/Cargo.toml

@@ -45,6 +45,7 @@ rustls = { workspace = true }
 serde = { workspace = true }
 serde_json = { workspace = true }
 serde_yaml_ng = { workspace = true }
+sse-core = { workspace = true }
 stargate = { workspace = true }
 stargate-proto = { workspace = true }
 stargate-protocol = { workspace = true }

src/libraries/rust/stargate/crates/stargate-bench/NOTICE

@@ -0,0 +1,6 @@
+Stargate benchmark third-party software
+
+sse-core 0.2.3
+Copyright (c) 2026 Max Shteimberg
+https://github.com/PizzasBear/sse-rs
+Used under the Apache License 2.0: https://www.apache.org/licenses/LICENSE-2.0