fix(xai): preserve prompt cache tokens in streaming responses

[stofancy]

Summary

Fixes #6144

The xAI channel's streaming handler split the upstream usage across two objects: it forwarded the complete upstream usage to the client, but fed a stripped usage (only 3 scalar fields) into billing. The net effect: every streaming request's cached_tokens was silently zeroed in logs/billing, while the client still received the correct value — and when cache_ratio < 1, this caused real overcharging.

Root cause — dual-path divergence

xAIStreamHandler in relay/channel/xai/text.go keeps two usage objects:

usage := &dto.Usage{}                       // object A — used for BILLING (returned upstream)
...
var xAIResp *dto.ChatCompletionsStreamResponse
common.UnmarshalJsonStr(data, &xAIResp)     // xAIResp.Usage = object B — upstream's full usage

if xAIResp.Usage != nil {
    // object A: only 3 scalars copied → CachedTokens dropped
    usage.PromptTokens = xAIResp.Usage.PromptTokens
    usage.TotalTokens = xAIResp.Usage.TotalTokens
    usage.CompletionTokens = usage.TotalTokens - usage.PromptTokens
}

openaiResponse := streamResponseXAI2OpenAI(xAIResp, usage)  // forwarded to client
helper.ObjectData(c, openaiResponse)                         // → uses object B (full)
...
return usage, nil   // → object A (stripped) goes to PostTextConsumeQuota

streamResponseXAI2OpenAI forwards the original xAIResp.Usage (object B):

openAIResp := &dto.ChatCompletionsStreamResponse{
    ...
    Usage: xAIResp.Usage,   // full upstream usage, incl. cached_tokens
}

	forwarded to client	internal billing
object used	xAIResp.Usage (upstream original, full)	usage (rebuilt, stripped)
cached_tokens	preserved	lost (→ 0)

This is why a downstream client (e.g. LiteLLM) sees the correct cached_tokens while new-api's own console logs show 0 for the same request.

Proof via packet capture

Streaming request sent directly to new-api (stream_options.include_usage=true), the final chunk new-api returns to the client:

{"prompt_tokens": 1816, "completion_tokens": 64, "total_tokens": 1880,
 "prompt_tokens_details": {"cached_tokens": 1792, ...}}

Client gets cached_tokens = 1792 (correct), but the same request is logged in-console with cache_tokens = 0.

The non-streaming xAIHandler is unaffected because it returns the full xaiResponse.Usage. The OpenAI channel is also unaffected — handleLastResponse (relay/channel/openai/helper.go) does *usage = lastStreamResponse.Usage, copying the whole struct. The xAI streaming handler is the only place that rebuilds usage by hand.

Why it matters — billing impact

The billing formula (service/text_quota.go, OpenAI semantic):

promptQuota = (PromptTokens - CacheTokens) + CacheTokens * CacheRatio

Cached tokens are subtracted from the base, then re-added at CacheRatio (the discount). With the bug zeroing CacheTokens, cached tokens get billed at full price instead of the discounted rate — an overcharge of 1 - cache_ratio per cached token.

Verified with a 6-turn incremental conversation (same prompt, streaming vs non-streaming). Actual console logs: non-streaming recorded cache growth correctly (7424→8064→8704→9344→10112); streaming logged 0 every turn and charged 2–2.6× the correct quota. Full table in #6144.

No impact when cache_ratio == 1 (default for unlisted models) or when there's no cache hit.

Fix

Mirror the OpenAI channel's approach: copy the full upstream usage so billing and forwarding share the same source of truth.

*usage = *xAIResp.Usage
if usage.CompletionTokens == 0 && usage.TotalTokens > usage.PromptTokens {
    usage.CompletionTokens = usage.TotalTokens - usage.PromptTokens
}

The total - prompt fallback is preserved for upstreams that omit CompletionTokens. No change needed to applyUsagePostProcessing: xAI reports cache hits via the standard prompt_tokens_details.cached_tokens path, preserved by the struct copy.

Verification

• go build ./... ✅
• go vet ./relay/channel/xai/ ./relay/channel/openai/ ✅

Reproduction / expected vs actual

See #6144 for the full reproduction steps, the dual-path source analysis, and the 6-turn console-log table.

[coderabbitai]

Walkthrough

The xAI streaming handler now copies the complete upstream usage object, preserving nested fields such as cached-token details. It retains a fallback calculation for missing completion tokens before downstream response conversion.

Changes

xAI streaming usage

Layer / File(s)	Summary
Preserve upstream usage fields relay/channel/xai/text.go	The stream handler copies the full upstream usage structure and recalculates completion tokens only when the upstream value is zero and total tokens exceed prompt tokens.

Estimated code review effort: 2 (Simple) | ~10 minutes

Suggested reviewers: seefs001

Poem

I’m a rabbit guarding tokens bright,
Cached fields now hop into sight.
Prompt and total keep their tune,
Completion fills in when it’s due.
Stream on, little bytes—boing!

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	The fix copies the full upstream xAI usage payload, preserving cached token details and matching the issue's streaming billing requirement.
Out of Scope Changes check	✅ Passed	The diff stays limited to xAI streaming usage handling and doesn't introduce unrelated changes.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly matches the main change: preserving xAI prompt cache tokens in streaming responses.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[stofancy]

@seefs001