fix(stream): preserve usage chunk before non-standard upstream frames

[Ankairis]

Problem

When streaming with upstreams that send non-standard chunks after the usage-bearing finish_reason chunk (e.g. OpenCode.ai sends x-opencode-type and extra cost frames), OaiStreamHandler overwrites lastStreamData on every frame. handleLastResponse then sees a chunk without usage, containStreamUsage stays false, and the code falls back to ResponseText2Usage (local token counting). This causes cache_tokens to always be 0 regardless of actual upstream cache hits.

Root Cause

relay/channel/openai/relay-openai.go — OaiStreamHandler:

lastStreamData = data   // ← overwritten by non-standard chunk
processTokenData(data)
// ...
handleLastResponse(lastStreamData)  // ← wrong chunk, no usage

Upstream stream order:

... finish_reason chunk (with usage)   ← has cache_tokens
... x-opencode-type chunk (no usage)   ← overwrites lastStreamData!
... [DONE]

Fix

1. Introduce usageStreamData — stores the last chunk that contained valid Usage
2. In the stream callback, parse each chunk and save it to usageStreamData if it has valid usage
3. Before falling back to ResponseText2Usage, try usageStreamData first
4. Pass usageStreamData to applyUsagePostProcessing so the original upstream body is still used for cache-token extraction

Testing

Verified with upstream curl test that OpenCode.ai returns prompt_cache_hit_tokens and prompt_tokens_details.cached_tokens correctly in the finish_reason chunk. After this fix, those values are preserved instead of being lost to the subsequent non-standard frame.

Closes #(issue if applicable)

Summary by CodeRabbit

• Bug Fixes
  • Improved usage reporting for streamed responses by recovering usage data from earlier valid stream updates.
  • Added fallback estimation when usage details are unavailable.
  • Preserved streamed content and trailing response frames while ensuring usage and cached token counts remain accurate in non-standard streaming scenarios.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 26c0b777-649c-4dbb-b73d-d6463175a862

📥 Commits

Reviewing files that changed from the base of the PR and between dfb5f87 and ce2960f.

📒 Files selected for processing (2)

• relay/channel/openai/relay-openai.go
• relay/channel/openai/relay_openai_stream_usage_test.go

🚧 Files skipped from review as they are similar to previous changes (1)

• relay/channel/openai/relay-openai.go

---

Walkthrough

OaiStreamHandler preserves a valid usage-bearing SSE chunk, uses it for final usage resolution and post-processing, and verifies recovery when a trailing non-standard frame follows it.

Changes

Stream usage handling

Layer / File(s)	Summary
Capture and apply usage-bearing chunks relay/channel/openai/relay-openai.go, relay/channel/openai/relay_openai_stream_usage_test.go	The handler stores valid usage-bearing chunks, checks them before estimating usage, and prefers them for post-processing. The regression test verifies token counts and relayed stream content after a trailing non-standard frame.

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

Poem

I’m a rabbit saving usage bytes,
Checking chunks at stream-time gates.
A trailing frame cannot erase
The token counts from their proper place.
Hop, hop—clean streams leave a trace.

🚥 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
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly describes preserving the usage chunk before non-standard upstream frames, which is the main change.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ 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.}

[coderabbitai]

🧹 Nitpick comments (1)

relay/channel/openai/relay-openai.go (1)

140-149: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick win

Add a strings.Contains pre-filter before the per-chunk unmarshal.

Every incoming chunk is unmarshaled into a full dto.ChatCompletionsStreamResponse just to check for a usage field. Most chunks are content-only and won't contain usage, so this adds unnecessary overhead on the streaming hot path. A lightweight substring check before the unmarshal skips the majority of parses.

♻️ Proposed optimization

 			// 检测当前 chunk 是否含有 usage，保存备用
 			// (OpenCode.ai 等上游在 usage chunk 后还会发非标准块，会覆盖 lastStreamData)
-			var chunkWithUsage dto.ChatCompletionsStreamResponse
-			if err := common.UnmarshalJsonStr(data, &chunkWithUsage); err == nil && chunkWithUsage.Usage != nil {
-				if service.ValidUsage(chunkWithUsage.Usage) {
-					usageStreamData = data
-				}
+			if strings.Contains(data, `"usage"`) {
+				var chunkWithUsage dto.ChatCompletionsStreamResponse
+				if err := common.UnmarshalJsonStr(data, &chunkWithUsage); err == nil && chunkWithUsage.Usage != nil {
+					if service.ValidUsage(chunkWithUsage.Usage) {
+						usageStreamData = data
+					}
+				}
 			}

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@relay/channel/openai/relay-openai.go` around lines 140 - 149, In the
usage-detection block near the per-chunk unmarshal, add a strings.Contains
pre-filter for the usage field before calling common.UnmarshalJsonStr. Only
unmarshal into dto.ChatCompletionsStreamResponse and run service.ValidUsage when
the chunk contains that substring, preserving the existing assignment to
usageStreamData.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@relay/channel/openai/relay-openai.go`:
- Around line 140-149: In the usage-detection block near the per-chunk
unmarshal, add a strings.Contains pre-filter for the usage field before calling
common.UnmarshalJsonStr. Only unmarshal into dto.ChatCompletionsStreamResponse
and run service.ValidUsage when the chunk contains that substring, preserving
the existing assignment to usageStreamData.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: f6d38d5f-2263-4512-a5ea-62b9f3522e67

📥 Commits

Reviewing files that changed from the base of the PR and between 4e57038 and 5b50a2f.

📒 Files selected for processing (1)

• relay/channel/openai/relay-openai.go

[huanggefan]

This issue exists in the opencode interface.
When stream=True
Finally, an x-opencode-type chunk will be returned.

This results in the cache hits recorded by new-api always being 0

[Ankairis]

是的，根因就是 OpenCode.ai 在 finish_reason（含 usage）之后发了 x-opencode-type 非标准块，lastStreamData 被覆盖导致 usage 丢失。

这个 PR 的改法是在每个 chunk 处理时独立保存最后一个含有效 usage 的 chunk，回退计费前优先用它，同时对非 usage chunk 增加 strings.Contains("\"usage\"") 预检避免多余反序列化。

目前已在我这边生产环境验证通过，cache_tokens 从恒为 0 正确恢复为实际值。

[Ankairis]

Rebased onto the latest main — the conflict in relay/channel/openai/relay-openai.go is resolved and the change is gofmt-clean.

I also added a regression test (relay/channel/openai/relay_openai_stream_usage_test.go) that replays the exact failure shape: a usage-bearing finish_reason chunk followed by a non-standard x-opencode-type frame. It asserts the real usage and cached_tokens are recovered, and I verified it fails on the previous code and passes with this fix.

This has been running in production on my side since the original submission — OpenCode streams now record real cache_tokens instead of 0. @seefs001 I'd appreciate a review when you have a moment.