Skip to content

fix(gateway): include image input tokens in accounting #1618

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
steebchen merged 4 commits into from
Feb 9, 2026
Merged

fix(gateway): include image input tokens in accounting #1618

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
e4148e4
fix(gateway): include image input tokens in prompt token accounting
steebchen Feb 9, 2026
f7447bb
fix: guard image token adjustment to Google-only
steebchen Feb 9, 2026
857cfd2
fix: use parsed model name for image count check
steebchen Feb 9, 2026
49c4981
fix: compute streaming total_tokens from parts
steebchen Feb 9, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
158 changes: 106 additions & 52 deletions apps/gateway/src/chat/chat.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -299,12 +299,6 @@ chat.openapi(completions, async (c) => {
});
}

// Count input images from messages for cost calculation (only for gemini-3-pro-image-preview)
const inputImageCount =
modelInput === "gemini-3-pro-image-preview"
? countInputImages(messages)
: 0;

// Extract reasoning_effort as mutable variable for auto-routing modification
let reasoning_effort = validationResult.data.reasoning_effort;

Expand Down Expand Up @@ -378,6 +372,12 @@ chat.openapi(completions, async (c) => {
const requestedModel = parseResult.requestedModel;
const customProviderName = parseResult.customProviderName;

// Count input images from messages for cost calculation
const inputImageCount =
requestedModel === "gemini-3-pro-image-preview"
? countInputImages(messages)
: 0;

// Resolve model info and filter deactivated providers
const modelInfoResult = resolveModelInfo(
requestedModel,
Expand Down Expand Up @@ -1646,9 +1646,16 @@ chat.openapi(completions, async (c) => {
content: fullContent || null,
reasoningContent: fullReasoningContent || null,
finishReason: cachedStreamingResponse.metadata.finishReason,
promptTokens: promptTokens?.toString() || null,
promptTokens:
(costs.promptTokens ?? promptTokens)?.toString() || null,
completionTokens: completionTokens?.toString() || null,
totalTokens: totalTokens?.toString() || null,
totalTokens: costs.imageInputTokens
? (
(costs.promptTokens || promptTokens || 0) +
(completionTokens || 0) +
(reasoningTokens || 0)
).toString()
: totalTokens?.toString() || null,
reasoningTokens: reasoningTokens?.toString() || null,
cachedTokens: cachedTokens?.toString() || null,
hasError: false,
Expand All @@ -1669,7 +1676,7 @@ chat.openapi(completions, async (c) => {
discount: costs.discount ?? null,
pricingTier: costs.pricingTier ?? null,
dataStorageCost: calculateDataStorageCost(
promptTokens,
costs.promptTokens ?? promptTokens,
cachedTokens,
completionTokens,
reasoningTokens,
Expand Down Expand Up @@ -1785,9 +1792,20 @@ chat.openapi(completions, async (c) => {
content: cachedContent || null,
reasoningContent: cachedReasoningContent || null,
finishReason: cachedResponse.choices?.[0]?.finish_reason || null,
promptTokens: cachedResponse.usage?.prompt_tokens || null,
promptTokens:
(
cachedCosts.promptTokens ?? cachedResponse.usage?.prompt_tokens
)?.toString() || null,
completionTokens: cachedResponse.usage?.completion_tokens || null,
totalTokens: cachedResponse.usage?.total_tokens || null,
totalTokens: cachedCosts.imageInputTokens
? (
(cachedCosts.promptTokens ||
cachedResponse.usage?.prompt_tokens ||
0) +
(cachedResponse.usage?.completion_tokens || 0) +
(cachedResponse.usage?.reasoning_tokens || 0)
).toString()
: cachedResponse.usage?.total_tokens || null,
Comment thread
steebchen marked this conversation as resolved.
reasoningTokens: cachedResponse.usage?.reasoning_tokens || null,
cachedTokens:
cachedResponse.usage?.prompt_tokens_details?.cached_tokens || null,
Expand All @@ -1809,7 +1827,7 @@ chat.openapi(completions, async (c) => {
discount: cachedCosts.discount ?? null,
pricingTier: cachedCosts.pricingTier ?? null,
dataStorageCost: calculateDataStorageCost(
cachedResponse.usage?.prompt_tokens,
cachedCosts.promptTokens ?? cachedResponse.usage?.prompt_tokens,
cachedResponse.usage?.prompt_tokens_details?.cached_tokens,
cachedResponse.usage?.completion_tokens,
cachedResponse.usage?.reasoning_tokens,
Expand Down Expand Up @@ -2307,11 +2325,15 @@ chat.openapi(completions, async (c) => {
reasoningContent: null,
finishReason: "canceled",
promptTokens: billCancelled
? estimatedPromptTokens?.toString()
? (
cancelledCosts?.promptTokens ?? estimatedPromptTokens
)?.toString()
: null,
completionTokens: billCancelled ? "0" : null,
totalTokens: billCancelled
? estimatedPromptTokens?.toString()
? (
cancelledCosts?.promptTokens ?? estimatedPromptTokens
)?.toString()
: null,
reasoningTokens: null,
cachedTokens: null,
Expand All @@ -2335,7 +2357,7 @@ chat.openapi(completions, async (c) => {
discount: cancelledCosts?.discount ?? null,
dataStorageCost: billCancelled
? calculateDataStorageCost(
estimatedPromptTokens,
cancelledCosts?.promptTokens ?? estimatedPromptTokens,
null,
0,
null,
Expand Down Expand Up @@ -3049,17 +3071,14 @@ chat.openapi(completions, async (c) => {
streamingCosts.completionTokens ||
finalCompletionTokens ||
0,
total_tokens: (() => {
const fallbackTotal =
(streamingCosts.promptTokens ||
finalPromptTokens ||
0) +
total_tokens: Math.max(
1,
(streamingCosts.promptTokens || finalPromptTokens || 0) +
(streamingCosts.completionTokens ||
finalCompletionTokens ||
0) +
(reasoningTokens || 0);
return Math.max(1, finalTotalTokens ?? fallbackTotal);
})(),
(reasoningTokens || 0),
),
...(shouldIncludeCosts && {
cost_usd_total: streamingCosts.totalCost,
cost_usd_input: streamingCosts.inputCost,
Expand Down Expand Up @@ -3775,36 +3794,41 @@ chat.openapi(completions, async (c) => {
finish_reason: null,
},
],
usage: {
prompt_tokens: Math.max(
usage: (() => {
// Only add image input tokens for providers that
// exclude them from upstream usage (Google)
const providerExcludesImageInput =
usedProvider === "google-ai-studio" ||
usedProvider === "google-vertex" ||
usedProvider === "obsidian";
const imageInputAdj = providerExcludesImageInput
? inputImageCount * 560
: 0;
const adjPrompt = Math.max(
1,
Math.round(
promptTokens && promptTokens > 0
? promptTokens + inputImageCount * 560
: (calculatedPromptTokens || 1) + inputImageCount * 560,
? promptTokens + imageInputAdj
: (calculatedPromptTokens || 1) + imageInputAdj,
),
),
completion_tokens: Math.round(
);
const adjCompletion = Math.round(
completionTokens || calculatedCompletionTokens || 0,
),
total_tokens: Math.round(
totalTokens ||
calculatedTotalTokens ||
Math.max(
1,
(promptTokens && promptTokens > 0
? promptTokens + inputImageCount * 560
: (calculatedPromptTokens || 1) +
inputImageCount * 560) +
(completionTokens || calculatedCompletionTokens || 0),
),
),
...(cachedTokens !== null && {
prompt_tokens_details: {
cached_tokens: cachedTokens,
},
}),
},
);
return {
prompt_tokens: adjPrompt,
completion_tokens: adjCompletion,
total_tokens: Math.max(
1,
Math.round(adjPrompt + adjCompletion),
),
...(cachedTokens !== null && {
prompt_tokens_details: {
cached_tokens: cachedTokens,
},
}),
};
})(),
};

await writeSSEAndCache({
Expand Down Expand Up @@ -3883,6 +3907,17 @@ chat.openapi(completions, async (c) => {
project.organizationId,
);

// Use costs.promptTokens as canonical value (includes image input
// tokens for providers that exclude them from upstream usage)
if (costs.promptTokens !== null && costs.promptTokens !== undefined) {
const promptDelta =
(costs.promptTokens || 0) - (calculatedPromptTokens || 0);
if (promptDelta > 0) {
calculatedPromptTokens = costs.promptTokens;
calculatedTotalTokens = (calculatedTotalTokens || 0) + promptDelta;
}
}

// Extract plugin IDs for logging (streaming - no healing applied)
const streamingPluginIds = plugins?.map((p) => p.id) || [];

Expand Down Expand Up @@ -4342,9 +4377,13 @@ chat.openapi(completions, async (c) => {
content: null,
reasoningContent: null,
finishReason: "canceled",
promptTokens: billCancelled ? estimatedPromptTokens?.toString() : null,
promptTokens: billCancelled
? (cancelledCosts?.promptTokens ?? estimatedPromptTokens)?.toString()
: null,
completionTokens: billCancelled ? "0" : null,
totalTokens: billCancelled ? estimatedPromptTokens?.toString() : null,
totalTokens: billCancelled
? (cancelledCosts?.promptTokens ?? estimatedPromptTokens)?.toString()
: null,
reasoningTokens: null,
cachedTokens: null,
hasError: false,
Expand All @@ -4365,7 +4404,7 @@ chat.openapi(completions, async (c) => {
discount: cancelledCosts?.discount ?? null,
dataStorageCost: billCancelled
? calculateDataStorageCost(
estimatedPromptTokens,
cancelledCosts?.promptTokens ?? estimatedPromptTokens,
null,
0,
null,
Expand Down Expand Up @@ -4640,7 +4679,7 @@ chat.openapi(completions, async (c) => {
}

// Estimate tokens if not provided by the API
const { calculatedPromptTokens, calculatedCompletionTokens } = estimateTokens(
let { calculatedPromptTokens, calculatedCompletionTokens } = estimateTokens(
usedProvider,
messages,
content,
Expand Down Expand Up @@ -4681,6 +4720,21 @@ chat.openapi(completions, async (c) => {
project.organizationId,
);

// Use costs.promptTokens as canonical value (includes image input
// tokens for providers that exclude them from upstream usage)
if (costs.promptTokens !== null && costs.promptTokens !== undefined) {
const promptDelta =
(costs.promptTokens || 0) - (calculatedPromptTokens || 0);
if (promptDelta > 0) {
calculatedPromptTokens = costs.promptTokens;
totalTokens = (
(calculatedPromptTokens || 0) +
(calculatedCompletionTokens || 0) +
(calculatedReasoningTokens || 0)
).toString();
}
}
Comment thread
coderabbitai[bot] marked this conversation as resolved.

// Transform response to OpenAI format for non-OpenAI providers
// Include costs in response for all users
const shouldIncludeCosts = true;
Expand Down
4 changes: 2 additions & 2 deletions apps/gateway/src/lib/costs.spec.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -311,8 +311,8 @@ describe("calculateCosts", () => {
// Each input image is 560 tokens at $2/1M = $0.00112 per image
expect(result.imageInputTokens).toBe(1120); // 2 * 560
expect(result.imageInputCost).toBeCloseTo(0.00224); // 1120 * 2e-6
// promptTokens should NOT include image tokens
expect(result.promptTokens).toBe(1000);
// promptTokens should include image input tokens
expect(result.promptTokens).toBe(2120); // 1000 text + 1120 image
expect(result.inputCost).toBeGreaterThan(0);
// totalCost should include image input cost
expect(result.totalCost).toBeGreaterThan(
Expand Down
11 changes: 10 additions & 1 deletion apps/gateway/src/lib/costs.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -368,7 +368,16 @@ export async function calculateCosts(
imageInputCost: imageInputCost?.toNumber() ?? null,
imageOutputCost: imageOutputCost?.toNumber() ?? null,
totalCost: totalCost.toNumber(),
promptTokens: calculatedPromptTokens,
// Only add image input tokens to promptTokens for providers whose upstream
// usage excludes them (Google). Other providers (OpenAI, xAI) already
// include image tokens in their reported prompt_tokens.
promptTokens:
imageInputTokens &&
(provider === "google-ai-studio" ||
provider === "google-vertex" ||
provider === "obsidian")
? (calculatedPromptTokens || 0) + imageInputTokens
: calculatedPromptTokens,
completionTokens: calculatedCompletionTokens,
cachedTokens,
estimatedCost: isEstimated,
Expand Down
Loading