[inureyes]

Summary

The disaggregated router's /v1/chat/completions path never emitted a usage object, diverging from the single-node chat route (which reports usage in both streaming and non-streaming responses) and from the router's own /v1/completions path. PR #392 (issue #387) had already made the router resolve the worker's authoritative generated-token count for the chat path, but only used it for finish_reason; this PR wires that already-resolved count, plus the previously-discarded prompt token count, into the chat usage block.

What changed

• src/server/router_front.rs: non-streaming route_chat now passes prompt_tokens (previously discarded as _prompt_tokens) and the resolved completion_tokens into chat_completion_json, which now serializes a usage object (prompt_tokens, completion_tokens, total_tokens) matching the single-node ChatCompletionResponse shape.
• src/server/router_front.rs: streaming route_chat now reads stream_options.include_usage the same way the router's /v1/completions streaming 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, only on a successful handoff — a handoff failure still emits no usage chunk, matching the existing completions-path behavior.
• src/server/router_front_tests.rs: unit tests for chat_completion_json's new usage field and chat_chunk_usage's chunk shape.
• tests/disaggregated_router_e2e.rs: a new E2E test, disaggregated_router_chat_usage_matches_single_node_byte_fallback, that spins up a three-process disaggregated stack plus a single-node reference and asserts usage.prompt_tokens / usage.completion_tokens / usage.total_tokens (and finish_reason) match between router and single-node, for both the non-streaming response and the streaming final usage chunk. Gated behind the same checkpoint-presence guard as the other real-model tests in that file.

E2E fixture note (issue #708)

The test was first written against the Gemma byte-fallback checkpoint (mirroring the issue #387 completions test) and its first-ever real run uncovered a latent main defect: 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), after which the router 500s every request. That defect is filed as issue #708 (with root cause: the vacuous .all() guard in CachePool::extract_paged_blocks plus the fatal serving-loop exit); the pre-existing #392 completions E2E is unrunnable for the same reason. The chat-usage E2E in this PR is therefore targeted at MiniCPM-2B (mlx-community/MiniCPM-2B-sft-4bit-llama-format-mlx), which keeps the byte-fallback tokenizer requirement (Llama-style SentencePiece, byte_fallback = true) while running as a plain dense Llama model — pool-backed Fp16, inside the handoff scope.

Test plan

• cargo check --lib --tests
• cargo test --lib server::router_front (12 passed, including the 3 new unit tests)
• cargo clippy --lib --tests -- -D warnings
• cargo fmt --all -- --check
• cargo test --test disaggregated_router_e2e --release --features metal,accelerate disaggregated_router_chat_usage_matches_single_node_byte_fallback -- --ignored — passed on the real three-node stack with minicpm-2b-4bit: non-stream and stream usage parity (prompt_tokens 20 / completion_tokens 16 / total_tokens 36) and finish_reason both match single-node; the usage chunk carries choices: [] and is emitted only when stream_options.include_usage is set.
• cargo test --test disaggregated_router_e2e --release --features metal,accelerate disaggregated_router_streams_correct_output -- --ignored — pre-existing qwen3 router E2E still green on this branch.

Closes #398

[inureyes]

Implementation Review Summary

Intent

Emit a usage object on the disaggregated router's /v1/chat/completions responses (streaming and non-streaming), matching single-node parity (issue #398).

Findings Addressed

• None. The implementation is correct, complete, and well-integrated; no CRITICAL/HIGH/MEDIUM findings were surfaced by the independent review or by a supplementary automated review pass.

Remaining Items

• None.

Acceptance criteria (issue #398)

• Non-streaming /v1/chat/completions carries usage with prompt_tokens, completion_tokens, total_tokens matching the single-node ChatCompletionResponse shape (chat_completion_json, router_front.rs:1527-1555). The router omits the optional prompt_tokens_details.cached_tokens breakdown, which is consistent with its own /v1/completions CompletionResponse (no prompt-prefix cache in the disaggregated path); the three required fields match.
• Streaming honors stream_options.include_usage and emits a final usage chunk only on a successful handoff. Gating if include_usage && let Some(completion_tokens) = completion_tokens (router_front.rs:774) is semantically identical to single-node if include_usage && let Ok(ref r) = result (chat.rs:910) and the router's own completions path (router_front.rs:1033). No usage chunk on handoff failure.
• usage.completion_tokens uses the authoritative resolve_completion_tokens value (worker wire count, byte-fallback correct, clamped to max_tokens), resolved once and shared with finish_reason.
• E2E parity test disaggregated_router_chat_usage_matches_single_node_byte_fallback added to tests/disaggregated_router_e2e.rs behind the checkpoint-presence guard, asserting non-streaming and streaming usage + finish_reason parity on the Gemma byte-fallback checkpoint.

Parity notes

• SSE ordering matches single-node and the router's completions path: content/reasoning chunks, then the finish chunk, then the usage chunk (empty choices: []), then [DONE].
• The usage chunk carries choices: [] per the OpenAI streaming-usage convention (unit-tested in chat_chunk_usage_shape).
• Non-usage chunks carry no usage field (equivalent to single-node's skip_serializing_if omission).
• The new chat chunk follows the router's existing hand-built-JSON chat chunk convention (sse_event), consistent with the sibling chat_chunk_* builders.

Verification

• All stated requirements implemented
• No placeholder/mock code remaining
• Integrated into router code flow (route_chat streaming + non-streaming branches wired; only callers are route_chat and the unit tests)
• Project conventions followed (reuses resolve_completion_tokens, shared chat-template + build_generate_options; no shared hotspot functions touched, so no // Used by: updates needed)
• Existing modules reused where applicable
• No unintended structural changes (additive only)
• Tests pass: cargo fmt --all -- --check, cargo check --lib --tests, cargo test --lib server::router_front (12 passed incl. 3 new), cargo clippy --lib --tests -- -D warnings, cross-repo ref check (exit 0). The real-model E2E test remains gated on the gemma-3-1b-it-4bit checkpoint (not run in this environment).