From 13b905895ad952161d1d96afc2b9ca3442749141 Mon Sep 17 00:00:00 2001
From: Jeongkyu Shin <inureyes@gmail.com>
Date: Thu, 9 Jul 2026 14:52:00 +0900
Subject: [PATCH 1/3] fix(router): emit usage on the disaggregated
 /v1/chat/completions responses

PR #392 (issue #387) made the router carry the worker's authoritative generated-token count over the wire for the chat path, but only used it to fix finish_reason; the count was resolved and then discarded, so router chat clients got no token accounting at all, unlike the single-node chat route and the router's own /v1/completions path.

Non-streaming route_chat now threads the already-resolved prompt_tokens (previously discarded as _prompt_tokens) and completion_tokens (from the existing resolve_completion_tokens call) into chat_completion_json, which adds a usage object with prompt_tokens, completion_tokens, and total_tokens matching the single-node ChatCompletionResponse shape.

Streaming route_chat now reads stream_options.include_usage the same way the completions path does, resolves completion_tokens once for both finish_reason and the usage chunk, and emits a new chat_chunk_usage SSE chunk (empty choices, populated usage) after the finish chunk on success only; a handoff failure still emits no usage chunk, mirroring the completions path's existing behavior.

Added unit tests in router_front_tests.rs for chat_completion_json's usage field and chat_chunk_usage's shape, and a new E2E test in tests/disaggregated_router_e2e.rs (disaggregated_router_chat_usage_matches_single_node_byte_fallback) that compares router vs single-node chat usage for both non-streaming and streaming on a byte-fallback tokenizer (Gemma), gated behind the same checkpoint-presence guard as the other real-model tests in that file.

Refs #398
---
 src/server/router_front.rs        | 101 ++++++++--
 src/server/router_front_tests.rs  |  58 +++++-
 tests/disaggregated_router_e2e.rs | 320 ++++++++++++++++++++++++++++++
 3 files changed, 459 insertions(+), 20 deletions(-)

diff --git a/src/server/router_front.rs b/src/server/router_front.rs
index 79887199c..40415fd84 100644
--- a/src/server/router_front.rs
+++ b/src/server/router_front.rs
@@ -670,12 +670,19 @@ async fn route_chat(state: Arc<RouterState>, request: ChatCompletionRequest) ->
     // (`chatcmpl-<n>`) is preserved exactly.
     let request_id = state.next_id.fetch_add(1, Ordering::Relaxed);
     let request_id_str = format!("chatcmpl-{request_id}");
-    let (_prompt_tokens, rx) =
+    let (prompt_tokens, rx) =
         start_handoff(&state, &prompt, request_id, &request_id_str, &opts).await?;
 
     if request.stream {
         // Streaming: spawn a task that drives the handoff and feeds SSE events
-        // into an unbounded channel; return the SSE response immediately.
+        // into an unbounded channel; return the SSE response immediately. The
+        // optional trailing usage chunk mirrors the router's own
+        // `/v1/completions` streaming path and single-node chat (issue #398).
+        let include_usage = request
+            .stream_options
+            .as_ref()
+            .map(|o| o.include_usage)
+            .unwrap_or(false);
         let (chunk_tx, chunk_rx) =
             tokio::sync::mpsc::unbounded_channel::<Result<Event, Infallible>>();
         let state2 = state.clone();
@@ -731,18 +738,20 @@ async fn route_chat(state: Arc<RouterState>, request: ChatCompletionRequest) ->
             emit_filtered(filter.flush());
 
             let completed = result.is_ok();
+            // Resolve the authoritative completion-token count on success (issue
+            // #387); `None` on a handoff failure, which also means "stop" for
+            // `finish_reason` and no usage chunk below (issue #398), matching
+            // the router's own `/v1/completions` streaming path.
+            let completion_tokens = result
+                .as_ref()
+                .ok()
+                .map(|outcome| resolve_completion_tokens(outcome, frame_counted, max_tokens));
             // Mirror single-node chat: "length" when the whole token budget was
             // generated, else "stop". On a handoff failure finish as "stop"
             // alongside the visible error delta emitted just above.
-            let finish_reason = match &result {
-                Ok(outcome) => {
-                    if resolve_completion_tokens(outcome, frame_counted, max_tokens) >= max_tokens {
-                        "length"
-                    } else {
-                        "stop"
-                    }
-                }
-                Err(_) => "stop",
+            let finish_reason = match completion_tokens {
+                Some(n) if n >= max_tokens => "length",
+                _ => "stop",
             };
             if let Err(e) = result {
                 let _ = chunk_tx.send(Ok(sse_event(&chat_chunk_error(
@@ -756,6 +765,20 @@ async fn route_chat(state: Arc<RouterState>, request: ChatCompletionRequest) ->
                 &model,
                 finish_reason,
             ))));
+
+            // Emit the usage chunk only on success, matching single-node
+            // `stream_chat_completion` (`if include_usage && let Ok(ref r) =
+            // result`) and the router's own `/v1/completions` streaming path.
+            // On a handoff failure (`completion_tokens` is `None`) no usage
+            // chunk is sent.
+            if include_usage && let Some(completion_tokens) = completion_tokens {
+                let _ = chunk_tx.send(Ok(sse_event(&chat_chunk_usage(
+                    &request_id_str2,
+                    &model,
+                    prompt_tokens,
+                    completion_tokens,
+                ))));
+            }
             let _ = chunk_tx.send(Ok(Event::default().data("[DONE]")));
 
             state2.pending.lock().unwrap().remove(&request_id);
@@ -767,12 +790,15 @@ async fn route_chat(state: Arc<RouterState>, request: ChatCompletionRequest) ->
             .into_response())
     } else {
         // Non-streaming: collect all tokens (filtered) then return a single
-        // JSON object with `content` and, when present, `reasoning_content`.
+        // JSON object with `content`, `reasoning_content` when present, and a
+        // `usage` block matching single-node's `ChatCompletionResponse` shape
+        // (issue #398).
         let mut filter = stream_filter;
         let mut content = String::new();
         let mut reasoning = String::new();
         let mut rx = rx;
         let finish_reason;
+        let completion_tokens;
         {
             let mut frame_counted = 0usize;
             let mut absorb = |emit: FilterOutput| {
@@ -800,10 +826,10 @@ async fn route_chat(state: Arc<RouterState>, request: ChatCompletionRequest) ->
             let outcome = r?;
             // Mirror single-node chat finish_reason: "length" when the whole
             // token budget was generated, else "stop", using the worker's
-            // authoritative count when present (issue #387).
-            finish_reason = if resolve_completion_tokens(&outcome, frame_counted, opts.max_tokens)
-                >= opts.max_tokens
-            {
+            // authoritative count when present (issue #387). The same
+            // resolved count feeds `usage.completion_tokens` below (issue #398).
+            completion_tokens = resolve_completion_tokens(&outcome, frame_counted, opts.max_tokens);
+            finish_reason = if completion_tokens >= opts.max_tokens {
                 "length"
             } else {
                 "stop"
@@ -815,6 +841,8 @@ async fn route_chat(state: Arc<RouterState>, request: ChatCompletionRequest) ->
             &content,
             &reasoning,
             finish_reason,
+            prompt_tokens,
+            completion_tokens,
         ))
         .into_response())
     }
@@ -1462,16 +1490,48 @@ fn chat_chunk_error(id: &str, model: &str, msg: &str) -> serde_json::Value {
     })
 }
 
+/// Final usage chunk, sent only when the client requested
+/// `stream_options.include_usage` and the handoff succeeded (issue #398).
+/// `choices` is empty per the OpenAI streaming-usage-chunk convention.
+/// `completion_tokens` is the caller's already-resolved authoritative count
+/// (issue #387); mirrors the router's own `/v1/completions` streaming usage
+/// chunk (`CompletionChunk::usage`) and single-node chat's
+/// `ChatCompletionChunk::usage_with_cache` (without the prompt-cache
+/// breakdown, which the router does not track).
+fn chat_chunk_usage(
+    id: &str,
+    model: &str,
+    prompt_tokens: usize,
+    completion_tokens: usize,
+) -> serde_json::Value {
+    serde_json::json!({
+        "id": id,
+        "object": "chat.completion.chunk",
+        "model": model,
+        "choices": [],
+        "usage": {
+            "prompt_tokens": prompt_tokens,
+            "completion_tokens": completion_tokens,
+            "total_tokens": prompt_tokens + completion_tokens
+        }
+    })
+}
+
 /// Non-streaming response body. `reasoning_content` is included only when
 /// the filter routed any thinking text. `finish_reason` ("stop" or "length")
 /// is derived from the worker's authoritative token count (issue #387) so it
-/// matches the single-node chat route.
+/// matches the single-node chat route. `usage` mirrors single-node's
+/// `ChatCompletionResponse` shape (`prompt_tokens`, `completion_tokens`,
+/// `total_tokens`); `completion_tokens` is the caller's already-resolved
+/// authoritative count (issue #398).
 fn chat_completion_json(
     id: &str,
     model: &str,
     content: &str,
     reasoning: &str,
     finish_reason: &str,
+    prompt_tokens: usize,
+    completion_tokens: usize,
 ) -> serde_json::Value {
     let mut message = serde_json::json!({"role": "assistant", "content": content});
     if !reasoning.is_empty() {
@@ -1485,7 +1545,12 @@ fn chat_completion_json(
             "index": 0,
             "message": message,
             "finish_reason": finish_reason
-        }]
+        }],
+        "usage": {
+            "prompt_tokens": prompt_tokens,
+            "completion_tokens": completion_tokens,
+            "total_tokens": prompt_tokens + completion_tokens
+        }
     })
 }
 
diff --git a/src/server/router_front_tests.rs b/src/server/router_front_tests.rs
index b842760ba..073dce392 100644
--- a/src/server/router_front_tests.rs
+++ b/src/server/router_front_tests.rs
@@ -13,11 +13,13 @@
 // limitations under the License.
 
 //! Unit tests for the disaggregated router front-end's usage accounting
-//! (issue #387).
+//! (issues #387, #398).
 //!
 //! Model-free: exercise `resolve_completion_tokens`, the pure function that
 //! turns the worker's authoritative wire-carried token count (or the emitted-
-//! piece fallback) into the reported `usage.completion_tokens`.
+//! piece fallback) into the reported `usage.completion_tokens`, and the chat
+//! response/chunk builders (`chat_completion_json`, `chat_chunk_usage`) that
+//! surface it on `/v1/chat/completions` (issue #398).
 
 use super::*;
 
@@ -109,6 +111,58 @@ fn authoritative_count_fixes_finish_reason_flip() {
     );
 }
 
+// ── Chat `usage` parity (issue #398) ─────────────────────────────────────────
+
+/// The non-streaming chat JSON body carries a `usage` object with
+/// `prompt_tokens`, `completion_tokens`, and `total_tokens`, matching the
+/// single-node `ChatCompletionResponse` shape, using the caller-supplied
+/// (already-resolved) counts verbatim.
+#[test]
+fn chat_completion_json_includes_usage() {
+    let body = chat_completion_json("chatcmpl-1", "qwen3", "hello", "", "stop", 12, 5);
+    assert_eq!(body["usage"]["prompt_tokens"], serde_json::json!(12));
+    assert_eq!(body["usage"]["completion_tokens"], serde_json::json!(5));
+    assert_eq!(body["usage"]["total_tokens"], serde_json::json!(17));
+    // The rest of the shape is untouched by the usage addition.
+    assert_eq!(body["id"], serde_json::json!("chatcmpl-1"));
+    assert_eq!(body["object"], serde_json::json!("chat.completion"));
+    assert_eq!(
+        body["choices"][0]["finish_reason"],
+        serde_json::json!("stop")
+    );
+    assert_eq!(
+        body["choices"][0]["message"]["content"],
+        serde_json::json!("hello")
+    );
+}
+
+/// `reasoning_content` is included on the message only when non-empty; the
+/// `usage` object is present either way.
+#[test]
+fn chat_completion_json_usage_present_with_reasoning() {
+    let body = chat_completion_json("chatcmpl-2", "qwen3", "answer", "thinking...", "stop", 8, 3);
+    assert_eq!(
+        body["choices"][0]["message"]["reasoning_content"],
+        serde_json::json!("thinking...")
+    );
+    assert_eq!(body["usage"]["completion_tokens"], serde_json::json!(3));
+    assert_eq!(body["usage"]["total_tokens"], serde_json::json!(11));
+}
+
+/// The streaming usage chunk carries an empty `choices` array (per the OpenAI
+/// streaming-usage-chunk convention, matching `CompletionChunk::usage` and
+/// `ChatCompletionChunk::usage_with_cache`) and the caller-supplied counts.
+#[test]
+fn chat_chunk_usage_shape() {
+    let chunk = chat_chunk_usage("chatcmpl-3", "qwen3", 20, 7);
+    assert_eq!(chunk["id"], serde_json::json!("chatcmpl-3"));
+    assert_eq!(chunk["object"], serde_json::json!("chat.completion.chunk"));
+    assert_eq!(chunk["choices"], serde_json::json!([]));
+    assert_eq!(chunk["usage"]["prompt_tokens"], serde_json::json!(20));
+    assert_eq!(chunk["usage"]["completion_tokens"], serde_json::json!(7));
+    assert_eq!(chunk["usage"]["total_tokens"], serde_json::json!(27));
+}
+
 // ── /router/stats topology-disclosure redaction (issue #389) ────────────────
 
 /// Build a small `StatsNode` fixture for the redaction tests.
diff --git a/tests/disaggregated_router_e2e.rs b/tests/disaggregated_router_e2e.rs
index d0e00b07c..889172ba4 100644
--- a/tests/disaggregated_router_e2e.rs
+++ b/tests/disaggregated_router_e2e.rs
@@ -58,6 +58,12 @@ const GEMMA_DIR: &str = "gemma-3-1b-it-4bit";
 /// echoed `model`).
 const GEMMA_MODEL_ALIAS: &str = "gemma-completions";
 
+/// Model alias for the Gemma byte-fallback CHAT usage parity test (issue #398).
+/// A distinct alias from [`GEMMA_MODEL_ALIAS`] keeps the completions-path and
+/// chat-path fixtures independent even though each test spawns its own
+/// isolated set of processes.
+const GEMMA_CHAT_MODEL_ALIAS: &str = "gemma-chat";
+
 /// Expected concatenated SSE content from the router for the test prompt.
 ///
 /// This is the single-node greedy reference: a hybrid `mlxcel-server` answers
@@ -1070,6 +1076,320 @@ async fn disaggregated_router_completions_match_single_node_byte_fallback() {
     );
 }
 
+/// Issue #398: `POST /v1/chat/completions` through the disaggregated router must
+/// report a `usage` object (`prompt_tokens`, `completion_tokens`, `total_tokens`)
+/// identical to single-node, in both the non-streaming response and the
+/// streaming final usage chunk (gated on `stream_options.include_usage`), for a
+/// BYTE-FALLBACK tokenizer (Gemma). Before this fix the router chat path never
+/// emitted `usage` at all.
+///
+/// Reuses the same byte-fallback prompt idea as
+/// `disaggregated_router_completions_match_single_node_byte_fallback` (issue
+/// #387) so the authoritative wire-carried token count is actually exercised:
+/// Gemma's SentencePiece tokenizer emits a multi-byte character (the `é` in
+/// "café") as several `<0xXX>` byte-fallback model tokens that surface as a
+/// single detokenized text piece, so a naive emitted-piece count would
+/// under-count `completion_tokens`.
+///
+/// Gated like the other real-model tests: `#[ignore]` plus a checkpoint-presence
+/// guard that skips cleanly when the Gemma checkpoint is absent.
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+#[ignore = "spawns four real mlxcel-server processes loading gemma-3-1b-it-4bit; run with --ignored"]
+async fn disaggregated_router_chat_usage_matches_single_node_byte_fallback() {
+    let model_dir = repo_model_dir(GEMMA_DIR);
+    if !model_dir.exists() {
+        eprintln!(
+            "Skipping {GEMMA_DIR}: model directory not found at {}.\n\
+             Fetch with: ./target/release/mlxcel download mlx-community/gemma-3-1b-it-4bit",
+            model_dir.display()
+        );
+        return;
+    }
+    let binary = repo_binary_path("mlxcel-server");
+    if !binary.exists() {
+        eprintln!("Skipping: mlxcel-server binary not found");
+        return;
+    }
+    let model_arg = model_dir.to_string_lossy().to_string();
+    // Same byte-fallback-forcing prompt as the completions parity test above,
+    // phrased as a chat turn.
+    let messages = serde_json::json!([{
+        "role": "user",
+        "content": "Spell the French word for coffee. It is café. Repeat it:"
+    }]);
+    let nonstream_body = serde_json::json!({
+        "model": GEMMA_CHAT_MODEL_ALIAS,
+        "messages": messages,
+        "max_tokens": 16,
+        "temperature": 0.0
+    });
+    let stream_body = serde_json::json!({
+        "model": GEMMA_CHAT_MODEL_ALIAS,
+        "messages": messages,
+        "max_tokens": 16,
+        "temperature": 0.0,
+        "stream": true,
+        "stream_options": {"include_usage": true}
+    });
+    let client = reqwest::Client::new();
+
+    // ---- Single-node reference (started with --alias for model parity) ----
+    let (ref_nonstream, ref_usage_chunk) = {
+        let ports = reserve_ports(1);
+        let http = ports[0].to_string();
+        let _single = spawn_role_server(&[
+            "-m",
+            &model_arg,
+            "--host",
+            "127.0.0.1",
+            "--port",
+            &http,
+            "--alias",
+            GEMMA_CHAT_MODEL_ALIAS,
+            "--no-warmup",
+        ]);
+        let deadline = Instant::now() + Duration::from_secs(240);
+        let health_url = format!("http://127.0.0.1:{http}/health");
+        assert!(
+            wait_for_http_health(&health_url, deadline).await,
+            "single-node Gemma reference never became healthy at {health_url}"
+        );
+        let chat_url = format!("http://127.0.0.1:{http}/v1/chat/completions");
+
+        let ns_resp = client
+            .post(&chat_url)
+            .json(&nonstream_body)
+            .send()
+            .await
+            .expect("POST single-node /v1/chat/completions (non-stream)");
+        assert!(
+            ns_resp.status().is_success(),
+            "single-node non-stream chat completion returned HTTP {}",
+            ns_resp.status()
+        );
+        let ns_json = ns_resp
+            .json::<serde_json::Value>()
+            .await
+            .expect("parse single-node non-stream chat completion JSON");
+
+        let st_resp = client
+            .post(&chat_url)
+            .json(&stream_body)
+            .send()
+            .await
+            .expect("POST single-node /v1/chat/completions (stream)");
+        assert!(
+            st_resp.status().is_success(),
+            "single-node stream chat completion returned HTTP {}",
+            st_resp.status()
+        );
+        let st_body = st_resp.text().await.expect("read single-node SSE body");
+        let st_chunks = parse_completion_chunks(&st_body);
+        let usage_chunk = st_chunks
+            .iter()
+            .find(|c| !c["usage"].is_null())
+            .cloned()
+            .expect("single-node chat stream must carry a final usage chunk");
+
+        (ns_json, usage_chunk)
+        // _single drops here, killing the reference server.
+    };
+    eprintln!("single-node Gemma chat non-stream reference: {ref_nonstream}");
+    let ref_text = ref_nonstream["choices"][0]["message"]["content"]
+        .as_str()
+        .unwrap_or("");
+    assert!(
+        !ref_text.is_empty(),
+        "single-node Gemma chat reference produced empty content; parity check would be vacuous"
+    );
+    // Guard the test's premise: the byte-fallback path is only exercised if the
+    // greedy output actually contains a multi-byte character.
+    assert!(
+        !ref_text.is_ascii(),
+        "single-node Gemma chat output {ref_text:?} is pure ASCII; the byte-fallback \
+         count path is not exercised. Adjust the prompt so the output contains a \
+         multi-byte character."
+    );
+    assert!(
+        !ref_nonstream["usage"].is_null(),
+        "single-node chat non-stream response must carry a usage object; \
+         parity check would be vacuous"
+    );
+
+    // ---- Three-process disaggregated router run ----
+    let ports = reserve_ports(6);
+    let prefill_http = ports[0].to_string();
+    let decode_http = ports[1].to_string();
+    let router_http = ports[2].to_string();
+    let prefill_serving_addr = format!("127.0.0.1:{}", ports[3]);
+    let decode_serving_addr = format!("127.0.0.1:{}", ports[4]);
+    let router_serving_addr = format!("127.0.0.1:{}", ports[5]);
+
+    let _decode = spawn_role_server(&[
+        "-m",
+        &model_arg,
+        "--host",
+        "127.0.0.1",
+        "--port",
+        &decode_http,
+        "--parallel",
+        "2",
+        "--max-batch-size",
+        "2",
+        "--decode-storage-backend",
+        "paged",
+        "--node-role",
+        "decode",
+        "--serving-bind",
+        &decode_serving_addr,
+        "--no-warmup",
+    ]);
+    let _prefill = spawn_role_server(&[
+        "-m",
+        &model_arg,
+        "--host",
+        "127.0.0.1",
+        "--port",
+        &prefill_http,
+        "--parallel",
+        "2",
+        "--max-batch-size",
+        "2",
+        "--decode-storage-backend",
+        "paged",
+        "--node-role",
+        "prefill",
+        "--serving-bind",
+        &prefill_serving_addr,
+        "--decode-peers",
+        &decode_serving_addr,
+        "--no-warmup",
+    ]);
+    let _router = spawn_role_server(&[
+        "-m",
+        &model_arg,
+        "--host",
+        "127.0.0.1",
+        "--port",
+        &router_http,
+        "--node-role",
+        "router",
+        "--serving-bind",
+        &router_serving_addr,
+        "--prefill-peers",
+        &prefill_serving_addr,
+        "--decode-peers",
+        &decode_serving_addr,
+        "--no-warmup",
+    ]);
+
+    let deadline = Instant::now() + Duration::from_secs(240);
+    assert!(
+        wait_for_tcp(&decode_serving_addr, deadline).await,
+        "decode node serving transport never came up"
+    );
+    assert!(
+        wait_for_tcp(&prefill_serving_addr, deadline).await,
+        "prefill node serving transport never came up"
+    );
+    let health_url = format!("http://127.0.0.1:{router_http}/health");
+    assert!(
+        wait_for_http_health(&health_url, deadline).await,
+        "router HTTP /health never returned 200"
+    );
+    let router_chat_url = format!("http://127.0.0.1:{router_http}/v1/chat/completions");
+
+    // Non-streaming parity: the router must now emit a `usage` object at all,
+    // and it must match single-node exactly for the byte-fallback tokenizer.
+    let router_ns_resp = client
+        .post(&router_chat_url)
+        .json(&nonstream_body)
+        .send()
+        .await
+        .expect("POST router /v1/chat/completions (non-stream)");
+    assert!(
+        router_ns_resp.status().is_success(),
+        "router non-stream chat completion returned HTTP {}",
+        router_ns_resp.status()
+    );
+    let router_ns = router_ns_resp
+        .json::<serde_json::Value>()
+        .await
+        .expect("parse router non-stream chat completion JSON");
+    eprintln!("router Gemma chat non-stream: {router_ns}");
+    assert!(
+        !router_ns["usage"].is_null(),
+        "router non-stream /v1/chat/completions must carry a usage object (issue #398)"
+    );
+    assert_eq!(
+        router_ns["usage"]["prompt_tokens"], ref_nonstream["usage"]["prompt_tokens"],
+        "router usage.prompt_tokens diverges from single-node"
+    );
+    assert_eq!(
+        router_ns["usage"]["completion_tokens"], ref_nonstream["usage"]["completion_tokens"],
+        "router usage.completion_tokens diverges from single-node for a byte-fallback tokenizer"
+    );
+    assert_eq!(
+        router_ns["usage"]["total_tokens"], ref_nonstream["usage"]["total_tokens"],
+        "router usage.total_tokens diverges from single-node"
+    );
+    assert_eq!(
+        router_ns["choices"][0]["finish_reason"], ref_nonstream["choices"][0]["finish_reason"],
+        "router finish_reason diverges from single-node for a byte-fallback tokenizer"
+    );
+
+    // Streaming parity: the final usage chunk (stream_options.include_usage)
+    // must also match single-node exactly.
+    let router_st_resp = client
+        .post(&router_chat_url)
+        .json(&stream_body)
+        .send()
+        .await
+        .expect("POST router /v1/chat/completions (stream)");
+    assert!(
+        router_st_resp.status().is_success(),
+        "router stream chat completion returned HTTP {}",
+        router_st_resp.status()
+    );
+    let router_st_body = router_st_resp.text().await.expect("read router SSE body");
+    let router_st_chunks = parse_completion_chunks(&router_st_body);
+    let router_usage_chunk = router_st_chunks
+        .iter()
+        .find(|c| !c["usage"].is_null())
+        .cloned()
+        .unwrap_or_else(|| {
+            panic!(
+                "router chat stream carried no usage chunk despite \
+                 stream_options.include_usage=true (issue #398); chunks: {router_st_chunks:?}"
+            )
+        });
+    assert_eq!(
+        router_usage_chunk["usage"]["prompt_tokens"], ref_usage_chunk["usage"]["prompt_tokens"],
+        "router stream usage.prompt_tokens diverges from single-node"
+    );
+    assert_eq!(
+        router_usage_chunk["usage"]["completion_tokens"],
+        ref_usage_chunk["usage"]["completion_tokens"],
+        "router stream usage.completion_tokens diverges from single-node for a \
+         byte-fallback tokenizer"
+    );
+    assert_eq!(
+        router_usage_chunk["usage"]["total_tokens"], ref_usage_chunk["usage"]["total_tokens"],
+        "router stream usage.total_tokens diverges from single-node"
+    );
+    // The usage chunk itself carries no choices, per the OpenAI streaming-usage
+    // convention (matches the router's own /v1/completions usage chunk).
+    assert_eq!(
+        router_usage_chunk["choices"],
+        serde_json::json!([]),
+        "router chat usage chunk must carry an empty choices array"
+    );
+    eprintln!(
+        "OK: router /v1/chat/completions matches single-node usage for a byte-fallback \
+         tokenizer (non-stream + stream)."
+    );
+}
+
 // ── Multi-node routing, balancing, and failover (issue #201) ─────────────
 
 /// Spawn a `--node-role decode` server wired to its own serving address.

From a249a7e12aeccaaf31442c2f0816af010b08c775 Mon Sep 17 00:00:00 2001
From: Jeongkyu Shin <inureyes@gmail.com>
Date: Thu, 9 Jul 2026 15:27:14 +0900
Subject: [PATCH 2/3] test(e2e): retarget the router chat-usage parity E2E from
 Gemma to MiniCPM-2B

The gemma-3-1b-it-4bit fixture can never pass this test: gemma3 is a model-owned paged family, and a disaggregated handoff request for it crashes the prefill node's serving loop (PagedBlockPool::read_block_contents: layer 0 has no pool tensors to read from) before any usage can be emitted - discovered on the first-ever real run of the Gemma byte-fallback E2Es and filed as issue #708. MiniCPM-2B (llama-format export, mlx-community/MiniCPM-2B-sft-4bit-llama-format-mlx) keeps the byte-fallback tokenizer requirement (Llama-style SentencePiece with byte_fallback = true) while running as a plain dense Llama model, which is pool-backed Fp16 and inside the handoff scope, so the test now actually runs: non-stream and stream usage parity (prompt 20 / completion 16 / total 36) and finish_reason both match single-node on the three-node harness.
---
 tests/disaggregated_router_e2e.rs | 73 ++++++++++++++++++++-----------
 1 file changed, 48 insertions(+), 25 deletions(-)

diff --git a/tests/disaggregated_router_e2e.rs b/tests/disaggregated_router_e2e.rs
index 889172ba4..7910aff32 100644
--- a/tests/disaggregated_router_e2e.rs
+++ b/tests/disaggregated_router_e2e.rs
@@ -58,11 +58,30 @@ const GEMMA_DIR: &str = "gemma-3-1b-it-4bit";
 /// echoed `model`).
 const GEMMA_MODEL_ALIAS: &str = "gemma-completions";
 
-/// Model alias for the Gemma byte-fallback CHAT usage parity test (issue #398).
-/// A distinct alias from [`GEMMA_MODEL_ALIAS`] keeps the completions-path and
-/// chat-path fixtures independent even though each test spawns its own
-/// isolated set of processes.
-const GEMMA_CHAT_MODEL_ALIAS: &str = "gemma-chat";
+/// MiniCPM-2B (llama-format export) checkpoint directory name (issue #398).
+/// Two properties make it the chat-usage parity fixture:
+///
+/// * its Llama-style SentencePiece tokenizer has `byte_fallback = true`
+///   (`<0xXX>` byte pieces in-vocab), satisfying the byte-fallback tokenizer
+///   requirement of issue #398, and
+/// * the llama-format export runs as a plain dense Llama model, so it is
+///   pool-backed Fp16 — inside the disaggregated handoff scope (#125).
+///
+/// Gemma, the natural byte-fallback reference used by [`GEMMA_DIR`], cannot be
+/// used here: gemma3 is a model-owned paged family and the prefill handoff
+/// does not support it — the request crashes the prefill node's serving loop
+/// (`PagedBlockPool::read_block_contents: layer 0 has no pool tensors`,
+/// issue #708).
+/// Fetch with:
+/// `./target/release/mlxcel download mlx-community/MiniCPM-2B-sft-4bit-llama-format-mlx`
+/// and place/symlink it at `models/minicpm-2b-4bit`.
+const MINICPM_DIR: &str = "minicpm-2b-4bit";
+
+/// Model alias for the MiniCPM byte-fallback CHAT usage parity test (issue
+/// #398). A distinct alias from [`GEMMA_MODEL_ALIAS`] keeps the
+/// completions-path and chat-path fixtures independent even though each test
+/// spawns its own isolated set of processes.
+const MINICPM_CHAT_MODEL_ALIAS: &str = "minicpm-chat";
 
 /// Expected concatenated SSE content from the router for the test prompt.
 ///
@@ -1080,27 +1099,31 @@ async fn disaggregated_router_completions_match_single_node_byte_fallback() {
 /// report a `usage` object (`prompt_tokens`, `completion_tokens`, `total_tokens`)
 /// identical to single-node, in both the non-streaming response and the
 /// streaming final usage chunk (gated on `stream_options.include_usage`), for a
-/// BYTE-FALLBACK tokenizer (Gemma). Before this fix the router chat path never
-/// emitted `usage` at all.
+/// BYTE-FALLBACK tokenizer (MiniCPM-2B, llama-format export). Before this fix
+/// the router chat path never emitted `usage` at all.
 ///
 /// Reuses the same byte-fallback prompt idea as
 /// `disaggregated_router_completions_match_single_node_byte_fallback` (issue
-/// #387) so the authoritative wire-carried token count is actually exercised:
-/// Gemma's SentencePiece tokenizer emits a multi-byte character (the `é` in
-/// "café") as several `<0xXX>` byte-fallback model tokens that surface as a
-/// single detokenized text piece, so a naive emitted-piece count would
-/// under-count `completion_tokens`.
+/// #387) so the authoritative wire-carried token count is exercised end to
+/// end: with a byte-fallback SentencePiece tokenizer, a multi-byte character
+/// missing from the vocab is emitted as several `<0xXX>` byte-fallback model
+/// tokens that surface as a single detokenized text piece, so a naive
+/// emitted-piece count would under-count `completion_tokens` (the divergence
+/// itself is pinned by the `resolve_completion_tokens` unit tests). Gemma
+/// cannot serve as the fixture — see [`MINICPM_DIR`].
 ///
 /// Gated like the other real-model tests: `#[ignore]` plus a checkpoint-presence
-/// guard that skips cleanly when the Gemma checkpoint is absent.
+/// guard that skips cleanly when the MiniCPM checkpoint is absent.
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
-#[ignore = "spawns four real mlxcel-server processes loading gemma-3-1b-it-4bit; run with --ignored"]
+#[ignore = "spawns four real mlxcel-server processes loading minicpm-2b-4bit; run with --ignored"]
 async fn disaggregated_router_chat_usage_matches_single_node_byte_fallback() {
-    let model_dir = repo_model_dir(GEMMA_DIR);
+    let model_dir = repo_model_dir(MINICPM_DIR);
     if !model_dir.exists() {
         eprintln!(
-            "Skipping {GEMMA_DIR}: model directory not found at {}.\n\
-             Fetch with: ./target/release/mlxcel download mlx-community/gemma-3-1b-it-4bit",
+            "Skipping {MINICPM_DIR}: model directory not found at {}.\n\
+             Fetch with: ./target/release/mlxcel download \
+             mlx-community/MiniCPM-2B-sft-4bit-llama-format-mlx (place it at \
+             models/{MINICPM_DIR})",
             model_dir.display()
         );
         return;
@@ -1118,13 +1141,13 @@ async fn disaggregated_router_chat_usage_matches_single_node_byte_fallback() {
         "content": "Spell the French word for coffee. It is café. Repeat it:"
     }]);
     let nonstream_body = serde_json::json!({
-        "model": GEMMA_CHAT_MODEL_ALIAS,
+        "model": MINICPM_CHAT_MODEL_ALIAS,
         "messages": messages,
         "max_tokens": 16,
         "temperature": 0.0
     });
     let stream_body = serde_json::json!({
-        "model": GEMMA_CHAT_MODEL_ALIAS,
+        "model": MINICPM_CHAT_MODEL_ALIAS,
         "messages": messages,
         "max_tokens": 16,
         "temperature": 0.0,
@@ -1145,14 +1168,14 @@ async fn disaggregated_router_chat_usage_matches_single_node_byte_fallback() {
             "--port",
             &http,
             "--alias",
-            GEMMA_CHAT_MODEL_ALIAS,
+            MINICPM_CHAT_MODEL_ALIAS,
             "--no-warmup",
         ]);
         let deadline = Instant::now() + Duration::from_secs(240);
         let health_url = format!("http://127.0.0.1:{http}/health");
         assert!(
             wait_for_http_health(&health_url, deadline).await,
-            "single-node Gemma reference never became healthy at {health_url}"
+            "single-node MiniCPM reference never became healthy at {health_url}"
         );
         let chat_url = format!("http://127.0.0.1:{http}/v1/chat/completions");
 
@@ -1194,19 +1217,19 @@ async fn disaggregated_router_chat_usage_matches_single_node_byte_fallback() {
         (ns_json, usage_chunk)
         // _single drops here, killing the reference server.
     };
-    eprintln!("single-node Gemma chat non-stream reference: {ref_nonstream}");
+    eprintln!("single-node MiniCPM chat non-stream reference: {ref_nonstream}");
     let ref_text = ref_nonstream["choices"][0]["message"]["content"]
         .as_str()
         .unwrap_or("");
     assert!(
         !ref_text.is_empty(),
-        "single-node Gemma chat reference produced empty content; parity check would be vacuous"
+        "single-node MiniCPM chat reference produced empty content; parity check would be vacuous"
     );
     // Guard the test's premise: the byte-fallback path is only exercised if the
     // greedy output actually contains a multi-byte character.
     assert!(
         !ref_text.is_ascii(),
-        "single-node Gemma chat output {ref_text:?} is pure ASCII; the byte-fallback \
+        "single-node MiniCPM chat output {ref_text:?} is pure ASCII; the byte-fallback \
          count path is not exercised. Adjust the prompt so the output contains a \
          multi-byte character."
     );
@@ -1316,7 +1339,7 @@ async fn disaggregated_router_chat_usage_matches_single_node_byte_fallback() {
         .json::<serde_json::Value>()
         .await
         .expect("parse router non-stream chat completion JSON");
-    eprintln!("router Gemma chat non-stream: {router_ns}");
+    eprintln!("router MiniCPM chat non-stream: {router_ns}");
     assert!(
         !router_ns["usage"].is_null(),
         "router non-stream /v1/chat/completions must carry a usage object (issue #398)"

From cf0d2471e7bfdecc1fe528fde0073ac647809992 Mon Sep 17 00:00:00 2001
From: Jeongkyu Shin <inureyes@gmail.com>
Date: Thu, 9 Jul 2026 15:46:27 +0900
Subject: [PATCH 3/3] test(router): cover the chat streaming include_usage
 default gate

The chat streaming usage chunk added in 13b905895 is gated on `stream_options.include_usage`, but nothing exercised that gate directly: the two builder-function unit tests only check chunk shape, and the E2E test only exercises the opted-in (true) path, so a regression that always (or never) sent the usage chunk regardless of the request would go undetected.

Extract the gate's boolean resolution into `wants_stream_usage` (pure function, behavior-preserving) and add three unit tests covering the omitted-`stream_options`, explicit-`false`, and explicit-`true` cases.

Also update docs/CONTINUOUS_BATCHING.md's usage/finish_reason limitation bullet, which was stale: it still described the pre-#387 emitted-text-piece counting as the primary path, when #387 already made the worker's authoritative wire-carried count (issue #387) the primary source and text-piece counting only the mixed-version-cluster fallback. Added a bullet noting that /v1/chat/completions now reports the same usage shape as /v1/completions in both modes (issue #398).

Validation:
- cargo test --lib server::router_front (15 passed, including the 3 new tests)
- cargo fmt --all -- --check
- cargo clippy --lib --tests --features metal,accelerate -- -D warnings
- python3 scripts/ci/check_cross_repo_refs.py (advisory pass; all flagged refs verified as real lablup/mlxcel issues)

Refs #398
---
 docs/CONTINUOUS_BATCHING.md      |  3 ++-
 src/server/router_front.rs       | 19 ++++++++++----
 src/server/router_front_tests.rs | 45 ++++++++++++++++++++++++++++++--
 3 files changed, 59 insertions(+), 8 deletions(-)

diff --git a/docs/CONTINUOUS_BATCHING.md b/docs/CONTINUOUS_BATCHING.md
index 2c080e706..809a8ea2e 100644
--- a/docs/CONTINUOUS_BATCHING.md
+++ b/docs/CONTINUOUS_BATCHING.md
@@ -168,7 +168,8 @@ curl http://127.0.0.1:8300/v1/chat/completions \
 - Text-only. The router serves `POST /v1/chat/completions` and `POST /v1/completions`; multimodal requests are rejected.
 - On `/v1/completions`, three option groups are rejected with HTTP 400 because the disaggregated wire protocol does not carry the data they require: `logprobs`, `response_format` (structured output), and explicit reasoning/thinking budgets. Requests without those fields work normally.
 - `stop` sequences are not forwarded to the worker nodes by the router. This is a pre-existing limitation shared with the chat path.
-- `completion_tokens` in the usage block is counted from emitted detokenized text pieces. For byte-level BPE tokenizers (Qwen, Llama) this equals the worker's token count exactly. For byte-fallback tokenizers (Gemma `<0xXX>` byte sequences) it can under-count, which may flip `finish_reason` between `"length"` and `"stop"`. A precise fix requires the disaggregated wire protocol to carry the worker's token count; a follow-up issue tracks the protocol extension.
+- `completion_tokens` in the usage block uses the worker's authoritative generated-token count carried over the disaggregated wire protocol (issue #387), so it is exact for both byte-level BPE tokenizers (Qwen, Llama) and byte-fallback tokenizers (Gemma `<0xXX>` byte sequences), and `finish_reason` matches single-node. A node that predates the wire field (a mixed-version cluster) falls back to counting emitted detokenized text pieces, which can under-count for byte-fallback tokenizers and flip `finish_reason` between `"length"` and `"stop"`.
+- `/v1/chat/completions` reports the same `usage` shape as `/v1/completions` (`prompt_tokens`, `completion_tokens`, `total_tokens`): the non-streaming response always carries it, and the streaming response sends it as a trailing chunk with empty `choices` when the request sets `stream_options.include_usage`, matching single-node (issue #398).
 - The router stream filter suppresses model-specific structural markers (`<think>`, tool-call delimiters) and routes thinking content to `reasoning_content`. Tool-call parsing is not yet supported on the router path: only `content` and `reasoning_content` are emitted.
 - Co-locating the roles on one machine adds transport hops without the scaling
   benefit of separate prefill and decode hardware, so a single-machine setup is
diff --git a/src/server/router_front.rs b/src/server/router_front.rs
index 40415fd84..17ed52c15 100644
--- a/src/server/router_front.rs
+++ b/src/server/router_front.rs
@@ -678,11 +678,7 @@ async fn route_chat(state: Arc<RouterState>, request: ChatCompletionRequest) ->
         // into an unbounded channel; return the SSE response immediately. The
         // optional trailing usage chunk mirrors the router's own
         // `/v1/completions` streaming path and single-node chat (issue #398).
-        let include_usage = request
-            .stream_options
-            .as_ref()
-            .map(|o| o.include_usage)
-            .unwrap_or(false);
+        let include_usage = wants_stream_usage(&request);
         let (chunk_tx, chunk_rx) =
             tokio::sync::mpsc::unbounded_channel::<Result<Event, Infallible>>();
         let state2 = state.clone();
@@ -1285,6 +1281,19 @@ fn resolve_completion_tokens(
         .unwrap_or(frame_counted.min(max_tokens))
 }
 
+/// Whether the client requested the trailing streaming usage chunk via
+/// `stream_options.include_usage`. Defaults to `false` when the request omits
+/// `stream_options` entirely, or when the field is present but `false`,
+/// matching single-node chat/completions (`chat.rs`, `completions.rs`) and
+/// the router's own `/v1/completions` streaming path (issue #398).
+fn wants_stream_usage(request: &ChatCompletionRequest) -> bool {
+    request
+        .stream_options
+        .as_ref()
+        .map(|o| o.include_usage)
+        .unwrap_or(false)
+}
+
 /// Consume the two-part disaggregated result (prefill first token + decode
 /// continuation), call `on_token` for every text piece in order, and return the
 /// authoritative generation outcome once the stream is finalized (or an error
diff --git a/src/server/router_front_tests.rs b/src/server/router_front_tests.rs
index 073dce392..da5458520 100644
--- a/src/server/router_front_tests.rs
+++ b/src/server/router_front_tests.rs
@@ -17,9 +17,11 @@
 //!
 //! Model-free: exercise `resolve_completion_tokens`, the pure function that
 //! turns the worker's authoritative wire-carried token count (or the emitted-
-//! piece fallback) into the reported `usage.completion_tokens`, and the chat
+//! piece fallback) into the reported `usage.completion_tokens`; the chat
 //! response/chunk builders (`chat_completion_json`, `chat_chunk_usage`) that
-//! surface it on `/v1/chat/completions` (issue #398).
+//! surface it on `/v1/chat/completions`; and `wants_stream_usage`, the
+//! `stream_options.include_usage` gate that decides whether the streaming
+//! usage chunk is sent at all (issue #398).
 
 use super::*;
 
@@ -163,6 +165,45 @@ fn chat_chunk_usage_shape() {
     assert_eq!(chunk["usage"]["total_tokens"], serde_json::json!(27));
 }
 
+/// A request that omits `stream_options` entirely (the common case: most
+/// clients never set it) must not opt into the trailing usage chunk. This is
+/// the default this PR must not flip, since nothing else guards the streaming
+/// gate itself (`chat_completion_json` and `chat_chunk_usage` above only
+/// cover the chunk *shape*, not *whether* it is sent).
+#[test]
+fn wants_stream_usage_defaults_to_false_without_stream_options() {
+    let json = r#"{"model": "qwen3", "messages": [{"role": "user", "content": "hi"}]}"#;
+    let request: ChatCompletionRequest = serde_json::from_str(json).unwrap();
+    assert!(request.stream_options.is_none());
+    assert!(!wants_stream_usage(&request));
+}
+
+/// `stream_options.include_usage: false` (explicit, not just absent) is also
+/// a no.
+#[test]
+fn wants_stream_usage_respects_explicit_false() {
+    let json = r#"{
+        "model": "qwen3",
+        "messages": [{"role": "user", "content": "hi"}],
+        "stream_options": {"include_usage": false}
+    }"#;
+    let request: ChatCompletionRequest = serde_json::from_str(json).unwrap();
+    assert!(!wants_stream_usage(&request));
+}
+
+/// `stream_options.include_usage: true` opts in, matching single-node
+/// chat/completions and the router's own `/v1/completions` streaming path.
+#[test]
+fn wants_stream_usage_respects_explicit_true() {
+    let json = r#"{
+        "model": "qwen3",
+        "messages": [{"role": "user", "content": "hi"}],
+        "stream_options": {"include_usage": true}
+    }"#;
+    let request: ChatCompletionRequest = serde_json::from_str(json).unwrap();
+    assert!(wants_stream_usage(&request));
+}
+
 // ── /router/stats topology-disclosure redaction (issue #389) ────────────────
 
 /// Build a small `StatsNode` fixture for the redaction tests.