fix(router): emit usage on the disaggregated /v1/chat/completions responses

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

src/server/router_front.rs

@@ -670,12 +670,15 @@ 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 = wants_stream_usage(&request);
         let (chunk_tx, chunk_rx) =
             tokio::sync::mpsc::unbounded_channel::<Result<Event, Infallible>>();
         let state2 = state.clone();
@@ -731,18 +734,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 +761,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 +786,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 +822,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 +837,8 @@ async fn route_chat(state: Arc<RouterState>, request: ChatCompletionRequest) ->
             &content,
             &reasoning,
             finish_reason,
+            prompt_tokens,
+            completion_tokens,
         ))
         .into_response())
     }
@@ -1257,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
@@ -1462,16 +1499,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 +1554,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
+        }
     })
 }
 

src/server/router_front_tests.rs

@@ -13,11 +13,15 @@
 // 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`; the chat
+//! response/chunk builders (`chat_completion_json`, `chat_chunk_usage`) that
+//! 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::*;
 
@@ -109,6 +113,97 @@ 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));
+}
+
+/// 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.