fixes costing and token counts for streaming requests - #528
Conversation
|
Caution Review failedThe pull request is closed. 📝 WalkthroughSummary by CodeRabbit
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
📜 Recent review detailsConfiguration used: CodeRabbit UI Review profile: CHILL Plan: Pro 📒 Files selected for processing (30)
Comment |
This stack of pull requests is managed by Graphite. Learn more about stacking. |
c79c40c to
f21909b
Compare
af2710e to
f880604
Compare
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
core/providers/openai.go(1 hunks)tests/core-providers/openrouter_test.go(1 hunks)tests/core-providers/scenarios/chat_completion_stream.go(7 hunks)
🧰 Additional context used
🧬 Code graph analysis (2)
tests/core-providers/openrouter_test.go (3)
tests/core-providers/config/setup.go (1)
SetupTest(51-60)tests/core-providers/config/account.go (2)
ComprehensiveTestConfig(42-53)TestScenarios(20-39)core/schemas/bifrost.go (1)
OpenRouter(54-54)
tests/core-providers/scenarios/chat_completion_stream.go (1)
core/schemas/bifrost.go (2)
BifrostResponse(515-528)BifrostStream(827-830)
⏰ Context from checks skipped due to timeout of 900000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (6)
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
tests/core-providers/openrouter_test.go (1)
44-110: Stabilize new OpenRouter Anthropic/Llama suites to reduce CI flake.These add real E2E paths against specific models. To keep CI stable across rate limits/model availability, consider gating with an env/flag (e.g.,
RUN_OPENROUTER_E2E=1) or honoringSkipReasonwhen credentials/quotas aren’t present. Optionally addt.Parallel()only where rate limits permit.If desired, I can add a small helper that checks for credentials/quotas and calls
t.SkipfwithSkipReasonwhen unavailable.tests/core-providers/scenarios/chat_completion_stream.go (1)
49-49: Make usage assertions tolerant and tied to the final metadata chunk.Some providers/proxies only emit usage in the final metadata chunk (choices=[]), while others may include zero/placeholder usage in interim chunks. Restrict per-chunk validations to metadata-only chunks to avoid false failures; keep stronger checks on the last chunk.
Apply diffs to guard by
len(response.Choices) == 0:- if response.Usage != nil { + if response.Usage != nil && len(response.Choices) == 0 { hasReceivedUsage = true ... }and in the tool-call stream:
- if response.Usage != nil { + if response.Usage != nil && len(response.Choices) == 0 { hasReceivedUsageWithTools = true ... }This preserves the stronger final-usage checks below while avoiding brittle assumptions mid‑stream. Based on docs. (community.openai.com)
Also applies to: 101-115, 130-148, 155-160, 217-219, 236-250, 285-292
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
core/providers/openai.go(1 hunks)tests/core-providers/openrouter_test.go(1 hunks)tests/core-providers/scenarios/chat_completion_stream.go(7 hunks)
🧰 Additional context used
🧬 Code graph analysis (2)
tests/core-providers/openrouter_test.go (3)
tests/core-providers/config/setup.go (1)
SetupTest(51-60)tests/core-providers/config/account.go (2)
ComprehensiveTestConfig(42-53)TestScenarios(20-39)core/schemas/bifrost.go (1)
OpenRouter(54-54)
tests/core-providers/scenarios/chat_completion_stream.go (1)
core/schemas/bifrost.go (2)
BifrostResponse(515-528)BifrostStream(827-830)
⏰ Context from checks skipped due to timeout of 900000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (11)
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
🔇 Additional comments (1)
core/providers/openai.go (1)
502-515: Refactor usage accumulator and expand scanner buffer
- Change the accumulator declaration before the stream loop from
tousage := &schemas.LLMUsage{}var usage *schemas.LLMUsage- Replace the in-chunk summing logic with a nil-check + latest-first (fallback to deltas) block:
- if usage != nil { - usage.PromptTokens += response.Usage.PromptTokens - usage.CompletionTokens += response.Usage.CompletionTokens - usage.TotalTokens += response.Usage.TotalTokens - } else { - usage = response.Usage - } + if usage == nil { + usage = &schemas.LLMUsage{} + } + // Use the highest TotalTokens seen to avoid double-counting; sum only if deltas + if response.Usage.TotalTokens >= usage.TotalTokens { + *usage = *response.Usage + } else { + usage.PromptTokens += response.Usage.PromptTokens + usage.CompletionTokens += response.Usage.CompletionTokens + usage.TotalTokens += response.Usage.TotalTokens + } + response.Usage = nil- Immediately after
addscanner := bufio.NewScanner(resp.Body)scanner.Buffer(make([]byte, 0, 64*1024), 2*1024*1024)- Manually verify the provider’s streaming behavior (OpenAI vs. proxies) to confirm whether usage only appears on the final chunk (so “latest” is sufficient) or if you must sum deltas.
8aa3421 to
598ed26
Compare
598ed26 to
62870f9
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
framework/logstore/tables.go (1)
235-242: Guard against nil inner error before clearing it
l.ErrorDetailsParsed.Errorcan legitimately benil(e.g., when upstream error metadata was absent), so dereferencing it here will panic and break request logging. Please guard the assignment so we only clear the nested error when the parent struct exists.if l.ErrorDetailsParsed != nil { - l.ErrorDetailsParsed.Error.Error = nil + if l.ErrorDetailsParsed.Error != nil { + l.ErrorDetailsParsed.Error.Error = nil + } if data, err := json.Marshal(l.ErrorDetailsParsed); err != nil {
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (30)
core/bifrost.go(1 hunks)core/changelog.md(1 hunks)core/providers/openai.go(1 hunks)core/version(1 hunks)framework/changelog.md(1 hunks)framework/logstore/sqlite.go(1 hunks)framework/logstore/tables.go(1 hunks)framework/version(1 hunks)plugins/governance/changelog.md(1 hunks)plugins/governance/version(1 hunks)plugins/jsonparser/changelog.md(1 hunks)plugins/jsonparser/version(1 hunks)plugins/logging/changelog.md(1 hunks)plugins/logging/main.go(1 hunks)plugins/logging/operations.go(4 hunks)plugins/logging/streaming.go(6 hunks)plugins/maxim/changelog.md(1 hunks)plugins/maxim/version(1 hunks)plugins/mocker/changelog.md(1 hunks)plugins/mocker/version(1 hunks)plugins/semanticcache/changelog.md(1 hunks)plugins/semanticcache/version(1 hunks)plugins/telemetry/changelog.md(1 hunks)plugins/telemetry/version(1 hunks)tests/core-providers/openrouter_test.go(1 hunks)tests/core-providers/scenarios/chat_completion_stream.go(7 hunks)transports/bifrost-http/handlers/logging.go(1 hunks)transports/changelog.md(1 hunks)transports/version(1 hunks)ui/app/logs/views/logDetailsSheet.tsx(1 hunks)
✅ Files skipped from review due to trivial changes (16)
- plugins/maxim/version
- core/changelog.md
- framework/version
- plugins/semanticcache/version
- framework/logstore/sqlite.go
- core/version
- plugins/logging/main.go
- framework/changelog.md
- transports/version
- plugins/mocker/changelog.md
- plugins/logging/changelog.md
- plugins/jsonparser/version
- plugins/mocker/version
- plugins/semanticcache/changelog.md
- plugins/jsonparser/changelog.md
- transports/bifrost-http/handlers/logging.go
🚧 Files skipped from review as they are similar to previous changes (2)
- core/providers/openai.go
- tests/core-providers/openrouter_test.go
🧰 Additional context used
🧬 Code graph analysis (4)
plugins/logging/operations.go (2)
plugins/logging/main.go (2)
LoggerPlugin(117-132)CreatedTimestampKey(39-39)framework/logstore/tables.go (2)
Log(70-116)Log(119-121)
tests/core-providers/scenarios/chat_completion_stream.go (1)
core/schemas/bifrost.go (2)
BifrostResponse(515-528)BifrostStream(827-830)
core/bifrost.go (1)
core/schemas/bifrost.go (1)
BifrostContextKeyStreamEndIndicator(104-104)
plugins/logging/streaming.go (2)
plugins/logging/main.go (1)
LoggerPlugin(117-132)core/schemas/bifrost.go (1)
BifrostError(838-847)
⏰ Context from checks skipped due to timeout of 900000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (15)
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
🔇 Additional comments (11)
plugins/governance/changelog.md (1)
4-5: Changelog update looks good. Entries follow the established format and clearly capture the two version bumps.plugins/logging/operations.go (4)
40-40: LGTM! Context propagation improvement.The function call correctly uses the new signature with context as the first parameter, enabling better context-aware latency calculations.
145-145: LGTM! Consistent context propagation.The error handling path correctly uses the updated function signature with context as the first parameter.
185-185: LGTM! Context propagation maintained.The final chunk processing correctly uses the updated function signature with context as the first parameter.
270-285: Approve calculateLatency fallback and retry logic.
The retryOnNotFound implementation in plugins/logging/main.go correctly retries on logstore.ErrNotFound up to three times with 1s delays and respects context cancellation, so the fallback in calculateLatency is safe as implemented.plugins/logging/streaming.go (6)
9-9: LGTM! Appropriate JSON library choice.Adding sonic for JSON marshaling is a good choice for performance-critical serialization operations.
129-129: LGTM! Enhanced function signature for better error handling.The addition of context and error parameters enables proper context propagation and streaming error handling in accumulated chunks processing.
142-142: LGTM! Consistent context propagation.The calculateLatency call correctly uses the updated signature with context as the first parameter, maintaining consistency across the codebase.
151-153: LGTM! Proper error status handling.The conditional logic correctly sets the status to "error" when a response error is provided, ensuring accurate logging of streaming response states.
168-173: LGTM! Proper error vs success case handling.The conditional logic correctly separates error serialization from normal output message handling. The use of sonic.Marshal for error serialization is appropriate.
Consider handling the sonic.Marshal error explicitly, though the current implementation using
_is acceptable for this use case.
412-412: LGTM! Call site properly updated.The function call correctly uses the new signature, properly dereferencing the context and passing the error parameter for comprehensive streaming error handling.
62870f9 to
568f1bb
Compare
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
plugins/logging/operations.go (1)
40-45: Don’t fail the whole update on latency lookup errors.updateLogEntry currently returns on calculateLatency error, unlike the streaming path which degrades gracefully. Prefer logging a warning and proceeding without latency to avoid dropping updates when created_at lookup transiently fails.
- latency, err := p.calculateLatency(ctx, requestID, timestamp) - if err != nil { - return err - } - updates["latency"] = latency + if latency, err := p.calculateLatency(ctx, requestID, timestamp); err == nil { + updates["latency"] = latency + } else { + p.logger.Warn("calculateLatency failed for %s: %v", requestID, err) + }
🧹 Nitpick comments (4)
plugins/logging/operations.go (2)
185-190: Final-chunk latency: degrade gracefully instead of erroring.If created_at fetch fails here, the entire final update is aborted. Log and continue; omit latency rather than failing the write.
- latency, err = p.calculateLatency(ctx, requestID, timestamp) - if err != nil { - return fmt.Errorf("failed to get created_at for latency calculation: %w", err) - } - needsLatency = true + if l, err := p.calculateLatency(ctx, requestID, timestamp); err == nil { + latency = l + needsLatency = true + } else { + p.logger.Warn("calculateLatency failed for %s: %v", requestID, err) + }
270-285: Clamp negative latency to 0 to handle clock skew.If currentTime < created_at (clock drift), latency becomes negative. Clamp to 0ms.
- if ctxTimestamp, ok := ctx.Value(CreatedTimestampKey).(time.Time); ok { - return float64(currentTime.Sub(ctxTimestamp).Nanoseconds()) / 1e6, nil - } + if ctxTimestamp, ok := ctx.Value(CreatedTimestampKey).(time.Time); ok { + d := currentTime.Sub(ctxTimestamp) + if d < 0 { + d = 0 + } + return float64(d.Nanoseconds()) / 1e6, nil + } @@ - return float64(currentTime.Sub(originalEntry.CreatedAt).Nanoseconds()) / 1e6, nil + d := currentTime.Sub(originalEntry.CreatedAt) + if d < 0 { + d = 0 + } + return float64(d.Nanoseconds()) / 1e6, nilplugins/logging/streaming.go (2)
141-146: Guard latency calc when FinalTimestamp is zero.If FinalTimestamp wasn’t set, latency uses zero time. Default to now.
- latency, err := p.calculateLatency(ctx, requestID, accumulator.FinalTimestamp) + ts := accumulator.FinalTimestamp + if ts.IsZero() { + ts = time.Now() + } + latency, err := p.calculateLatency(ctx, requestID, ts)
129-137: Defer order is subtle; add intent comment.cleanupStreamAccumulator is deferred after Unlock defer, so it executes first (LIFO) while the mutex is still held. It’s correct but non‑obvious—add a comment to prevent accidental reordering.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (30)
core/bifrost.go(1 hunks)core/changelog.md(1 hunks)core/providers/openai.go(1 hunks)core/version(1 hunks)framework/changelog.md(1 hunks)framework/logstore/sqlite.go(1 hunks)framework/logstore/tables.go(1 hunks)framework/version(1 hunks)plugins/governance/changelog.md(1 hunks)plugins/governance/version(1 hunks)plugins/jsonparser/changelog.md(1 hunks)plugins/jsonparser/version(1 hunks)plugins/logging/changelog.md(1 hunks)plugins/logging/main.go(1 hunks)plugins/logging/operations.go(4 hunks)plugins/logging/streaming.go(6 hunks)plugins/maxim/changelog.md(1 hunks)plugins/maxim/version(1 hunks)plugins/mocker/changelog.md(1 hunks)plugins/mocker/version(1 hunks)plugins/semanticcache/changelog.md(1 hunks)plugins/semanticcache/version(1 hunks)plugins/telemetry/changelog.md(1 hunks)plugins/telemetry/version(1 hunks)tests/core-providers/openrouter_test.go(1 hunks)tests/core-providers/scenarios/chat_completion_stream.go(7 hunks)transports/bifrost-http/handlers/logging.go(1 hunks)transports/changelog.md(1 hunks)transports/version(1 hunks)ui/app/logs/views/logDetailsSheet.tsx(1 hunks)
✅ Files skipped from review due to trivial changes (4)
- plugins/governance/version
- framework/logstore/sqlite.go
- plugins/logging/main.go
- plugins/semanticcache/version
🚧 Files skipped from review as they are similar to previous changes (13)
- tests/core-providers/openrouter_test.go
- plugins/maxim/version
- tests/core-providers/scenarios/chat_completion_stream.go
- framework/logstore/tables.go
- plugins/semanticcache/changelog.md
- framework/changelog.md
- plugins/telemetry/version
- transports/bifrost-http/handlers/logging.go
- ui/app/logs/views/logDetailsSheet.tsx
- transports/version
- framework/version
- plugins/telemetry/changelog.md
- transports/changelog.md
🧰 Additional context used
🧬 Code graph analysis (3)
plugins/logging/streaming.go (2)
plugins/logging/main.go (1)
LoggerPlugin(117-132)core/schemas/bifrost.go (1)
BifrostError(838-847)
plugins/logging/operations.go (2)
plugins/logging/main.go (2)
LoggerPlugin(117-132)CreatedTimestampKey(39-39)framework/logstore/tables.go (2)
Log(70-116)Log(119-121)
core/bifrost.go (1)
core/schemas/bifrost.go (1)
BifrostContextKeyStreamEndIndicator(104-104)
⏰ Context from checks skipped due to timeout of 900000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (15)
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
- GitHub Check: Graphite / mergeability_check
🔇 Additional comments (7)
plugins/governance/changelog.md (1)
4-5: Changelog entries aligned with version bumps.Documenting the core 1.1.38 and framework 1.0.24 upgrades here keeps governance release notes in sync with the rest of the stack.
plugins/logging/changelog.md (1)
4-6: New changelog entries capture the logging fix.Noting the streaming logging fix alongside the coordinated core/framework bumps gives downstream consumers clear release context.
plugins/maxim/changelog.md (1)
4-5: Maxim changelog stays consistent with global upgrades.Listing both core and framework version bumps keeps this plugin’s release notes aligned with the wider rollout.
plugins/mocker/version (1)
1-1: Plugin version bump looks consistent.The 1.2.15 bump matches the coordinated plugin release cadence reflected elsewhere in the stack.
core/bifrost.go (1)
1138-1149: Thanks for restoring post-hook recovery handling.This mirrors the non-stream path: we now respect post-hook overrides (both recovered errors and synthetic responses) while still marking the stream end and releasing the channel. Looks solid.
plugins/logging/streaming.go (2)
412-413: LGTM on context/error propagation and single-run gating.processAccumulatedChunks now receives ctx and respErr, and IsComplete prevents double processing. Looks good.
If you can, run the OpenRouter streaming tests to ensure error-finalization still fires the callback once.
58-66: No nil‐check needed for Function—FunctionCall is a non‐pointer struct
TheFunctionfield onToolCallis declared as aFunctionCall(not*FunctionCall), so it can never be nil and does not require a guard.Likely an incorrect or invalid review comment.
568f1bb to
7496a33
Compare
Merge activity
|
7496a33 to
44990dc
Compare
## Summary Improves token usage tracking in OpenAI streaming responses and expands OpenRouter test coverage with additional model tests. ## Changes - Fixed token usage calculation in OpenAI streaming responses by properly accumulating usage data when it arrives in multiple chunks - Enhanced streaming tests to validate token usage information throughout the stream and in final responses - Added comprehensive tests for Anthropic Claude and Meta Llama models via OpenRouter - Removed the requirement for an OpenRouter API key in tests to allow CI to run these tests ## Type of change - [x] Bug fix - [x] Feature - [ ] Refactor - [ ] Documentation - [ ] Chore/CI ## Affected areas - [x] Core (Go) - [ ] Transports (HTTP) - [x] Providers/Integrations - [ ] Plugins - [ ] UI (Next.js) - [ ] Docs ## How to test Run the OpenRouter tests to verify proper token usage tracking and model support: ```sh # Run all OpenRouter tests go test -v ./tests/core-providers -run TestOpenRouter # Run specific model tests go test -v ./tests/core-providers -run TestOpenRouterAnthropic go test -v ./tests/core-providers -run TestOpenRouterMetaLlama ``` ## Breaking changes - [ ] Yes - [x] No ## Related issues Fixes issues with token usage tracking in streaming responses where usage information was being overwritten instead of accumulated. ## Security considerations No security implications. ## Checklist - [x] I added/updated tests where appropriate - [x] I verified builds succeed (Go and UI) - [x] I verified the CI pipeline passes locally if applicable
2 participants
Summary
Improves token usage tracking in OpenAI streaming responses and expands OpenRouter test coverage with additional model tests.
Changes
Type of change
Affected areas
How to test
Run the OpenRouter tests to verify proper token usage tracking and model support:
Breaking changes
Related issues
Fixes issues with token usage tracking in streaming responses where usage information was being overwritten instead of accumulated.
Security considerations
No security implications.
Checklist