fix(gateway): include image input tokens in accounting - #1618
Merged
Merged
Conversation
Include image input tokens in the returned promptTokens so that input + completion = total tokens. Billing remains unchanged as image tokens are still charged at their own rate separately. Co-Authored-By: Claude Haiku 4.5 <noreply@anthropic.com>
Contributor
WalkthroughChat gateway token accounting now prefers costs.promptTokens (when provided) and includes image input tokens for specific providers; logging, usage objects, and total token calculations were updated across streaming, cached, cancelled, and error paths to reflect costs-based accounting. Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant Gateway
participant Costs
participant Provider
participant Storage
Client->>Gateway: Send chat request (may include images)
Gateway->>Costs: calculateCosts(request, provider)
Costs-->>Gateway: {promptTokens, imageInputTokens, completionTokens, ...}
Gateway->>Provider: Request/stream model using adjusted accounting
Provider-->>Gateway: Streamed response / final usage info
Gateway->>Storage: Persist usage (uses costs.promptTokens + image adjustments)
Gateway-->>Client: Return response + canonicalized usage (total_tokens from costs)
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Possibly related PRs
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
🧪 Generate unit tests (beta)
No actionable comments were generated in the recent review. 🎉 🧹 Recent nitpick comments
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. Comment |
There was a problem hiding this comment.
Pull request overview
This PR updates gateway token accounting so that image input tokens are included in promptTokens, making prompt + completion (+ reasoning) = total consistent for analytics/reporting while keeping image billing at its own rate.
Changes:
- Update
calculateCosts()to returnpromptTokensinclusive of image input tokens. - Adjust chat logging paths to prefer the updated
costs.promptTokensand to recomputetotalTokenswhere needed. - Update unit test expectations for image token accounting.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 4 comments.
| File | Description |
|---|---|
| apps/gateway/src/lib/costs.ts | Adjusts returned promptTokens to include image input tokens. |
| apps/gateway/src/lib/costs.spec.ts | Updates test expectations to match new promptTokens semantics. |
| apps/gateway/src/chat/chat.ts | Updates DB logging to use the new prompt/total token accounting when images are present. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+3908
to
+3914
| // Include image input tokens in prompt/total for consistent accounting | ||
| if (costs.imageInputTokens) { | ||
| calculatedPromptTokens = | ||
| (calculatedPromptTokens || 0) + costs.imageInputTokens; | ||
| calculatedTotalTokens = | ||
| (calculatedTotalTokens || 0) + costs.imageInputTokens; | ||
| } |
There was a problem hiding this comment.
This block mutates calculatedPromptTokens/calculatedTotalTokens by adding costs.imageInputTokens. If calculatedPromptTokens already included image tokens from upstream usage for some providers/models, this will double-count in logs and any later calculations that rely on these variables. Prefer setting these values from costs.promptTokens (which is already the source of truth) or guard the addition behind an explicit “image tokens not included in upstream usage” condition.
Suggested change
| // Include image input tokens in prompt/total for consistent accounting | |
| if (costs.imageInputTokens) { | |
| calculatedPromptTokens = | |
| (calculatedPromptTokens || 0) + costs.imageInputTokens; | |
| calculatedTotalTokens = | |
| (calculatedTotalTokens || 0) + costs.imageInputTokens; | |
| } | |
| // Sync token accounting with costs to avoid double-counting image tokens | |
| if (typeof costs.promptTokens === "number") { | |
| calculatedPromptTokens = costs.promptTokens; | |
| } | |
| if (typeof costs.completionTokens === "number") { | |
| calculatedCompletionTokens = costs.completionTokens; | |
| } | |
| if (typeof costs.cachedTokens === "number") { | |
| cachedTokens = costs.cachedTokens; | |
| } | |
| calculatedTotalTokens = | |
| (calculatedPromptTokens || 0) + | |
| (calculatedCompletionTokens || 0) + | |
| (cachedTokens || 0); |
Comment on lines
+4718
to
+4728
| // Include image input tokens in prompt/total for consistent accounting | ||
| if (costs.imageInputTokens) { | ||
| calculatedPromptTokens = | ||
| (calculatedPromptTokens || 0) + costs.imageInputTokens; | ||
| totalTokens = ( | ||
| (calculatedPromptTokens || 0) + | ||
| (calculatedCompletionTokens || 0) + | ||
| (calculatedReasoningTokens || 0) | ||
| ).toString(); | ||
| } | ||
|
|
There was a problem hiding this comment.
Recomputing calculatedPromptTokens/totalTokens by adding costs.imageInputTokens has the same double-counting risk as the streaming path if upstream promptTokens already include image tokens for a given provider/model. Consider using costs.promptTokens directly as the canonical prompt token count (and only applying a manual image-token adjustment when you know the provider usage excludes image tokens).
Suggested change
| // Include image input tokens in prompt/total for consistent accounting | |
| if (costs.imageInputTokens) { | |
| calculatedPromptTokens = | |
| (calculatedPromptTokens || 0) + costs.imageInputTokens; | |
| totalTokens = ( | |
| (calculatedPromptTokens || 0) + | |
| (calculatedCompletionTokens || 0) + | |
| (calculatedReasoningTokens || 0) | |
| ).toString(); | |
| } | |
| // Use costs.promptTokens as the canonical prompt token count to avoid | |
| // double-counting image tokens that may already be included upstream. |
Only add image input tokens to promptTokens for providers that exclude them from upstream usage (Google). Other providers like OpenAI and xAI already include image tokens in their reported prompt_tokens, so adding them again would double-count. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
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)
apps/gateway/src/chat/chat.ts (1)
3066-3083:⚠️ Potential issue | 🟠 MajorStreaming [DONE] usage chunk still has
prompt + completion ≠ totalfor Google providers with image inputs.The PR updates
prompt_tokens(line 3067) to usestreamingCosts.promptTokenswhich now includes image input tokens for Google. However,total_tokens(line 3083) still prefersfinalTotalTokens— the upstream value that does not include image input tokens. This means the client-facing SSE usage in the common [DONE] path still has the exact mismatch the PR aims to fix.When
finalTotalTokensis non-null (which is the typical case for Google streaming), the fallback tofallbackTotalis never reached.Proposed fix
total_tokens: (() => { const fallbackTotal = (streamingCosts.promptTokens || finalPromptTokens || 0) + (streamingCosts.completionTokens || finalCompletionTokens || 0) + (reasoningTokens || 0); - return Math.max(1, finalTotalTokens ?? fallbackTotal); + // If costs include image input tokens, recompute total + // to ensure prompt + completion = total + if (streamingCosts.imageInputTokens) { + return Math.max(1, fallbackTotal); + } + return Math.max(1, finalTotalTokens ?? fallbackTotal); })(),
🤖 Fix all issues with AI agents
In `@apps/gateway/src/chat/chat.ts`:
- Around line 4727-4740: totalTokens is being reassigned from number | null to a
string when you do .toString() inside the costs.promptTokens branch, causing a
type mismatch for downstream code (originating from parseProviderResponse). Fix
by keeping totalTokens as a number: when costs.promptTokens >
calculatedPromptTokens set calculatedPromptTokens = costs.promptTokens and
recompute totalTokens as a numeric sum of calculatedPromptTokens,
calculatedCompletionTokens, and calculatedReasoningTokens (do not call
.toString()); only call .toString() where totalTokens is formatted for logging
or output (e.g., the log/print site around line ~4879). Ensure references:
totalTokens, costs.promptTokens, calculatedPromptTokens,
calculatedCompletionTokens, calculatedReasoningTokens, and
parseProviderResponse.
🧹 Nitpick comments (1)
apps/gateway/src/chat/chat.ts (1)
3800-3835: Extract the560tokens-per-image constant to avoid duplication and future divergence.The
inputImageCount * 560calculation at line 3808 duplicates the logic defined inapps/gateway/src/lib/costs.ts:286asTOKENS_PER_INPUT_IMAGE = 560. If the per-image token rate changes in the costs module, this client-facing usage chunk won't be updated, causing the client to see different token counts than what's billed. Extract this as a shared constant (or reference the one from costs.ts) to ensure both code paths always use the same value.
The inputImageCount check compared against modelInput which includes the provider prefix (e.g. "google-ai-studio/gemini-3-pro-image-preview") so it never matched. Move the check after parseModelInput and compare against requestedModel instead. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Ensures total_tokens in client-facing SSE responses is always computed from prompt + completion + reasoning tokens rather than preferring upstream values that exclude image input tokens. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Contributor
|
Caution Failed to replace (edit) comment. This is likely due to insufficient permissions or the comment being deleted. Error details |
3 similar comments
Contributor
|
Caution Failed to replace (edit) comment. This is likely due to insufficient permissions or the comment being deleted. Error details |
Contributor
|
Caution Failed to replace (edit) comment. This is likely due to insufficient permissions or the comment being deleted. Error details |
Contributor
|
Caution Failed to replace (edit) comment. This is likely due to insufficient permissions or the comment being deleted. Error details |
Contributor
|
Caution Failed to replace (edit) comment. This is likely due to insufficient permissions or the comment being deleted. Error details |
3 similar comments
Contributor
|
Caution Failed to replace (edit) comment. This is likely due to insufficient permissions or the comment being deleted. Error details |
Contributor
|
Caution Failed to replace (edit) comment. This is likely due to insufficient permissions or the comment being deleted. Error details |
Contributor
|
Caution Failed to replace (edit) comment. This is likely due to insufficient permissions or the comment being deleted. Error details |
3 tasks
2 tasks
4 tasks
This was referenced Jun 30, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Fixed image token accounting so that
input + completion = total. Image input tokens are now included in the returnedpromptTokensfield rather than being only in the separateimageInputTokensmetadata field.What changed
calculateCostsnow returnspromptTokensthat includes image input tokensImpact
🤖 Generated with Claude Code
Summary by CodeRabbit