Skip to content

Commit e1406e0

Browse files
fix(console): bill google non-stream zen usage (anomalyco#28829)
1 parent b555383 commit e1406e0

7 files changed

Lines changed: 76 additions & 2 deletions

File tree

packages/console/app/src/routes/zen/util/handler.ts

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -249,8 +249,9 @@ export async function handler(
249249
if (!isStream || [400, 404, 429].includes(res.status)) {
250250
const json = await res.json()
251251
await rateLimiter?.track()
252-
if (json.usage) {
253-
const usageInfo = providerInfo.normalizeUsage(json.usage)
252+
const usage = providerInfo.extractUsage(json)
253+
if (usage) {
254+
const usageInfo = providerInfo.normalizeUsage(usage)
254255
const costInfo = calculateCost(modelInfo, usageInfo)
255256
await trialLimiter?.track(usageInfo)
256257
await modelTpmLimiter?.track(providerInfo.id, providerInfo.model, usageInfo)

packages/console/app/src/routes/zen/util/provider/anthropic.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -175,6 +175,7 @@ export const anthropicHelper: ProviderHelper = ({ reqModel, providerModel }) =>
175175
retrieve: () => usage,
176176
}
177177
},
178+
extractUsage: (response: any) => response.usage,
178179
normalizeUsage: (usage: Usage) => ({
179180
inputTokens: usage.input_tokens ?? 0,
180181
outputTokens: usage.output_tokens ?? 0,

packages/console/app/src/routes/zen/util/provider/google.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,7 @@ export const googleHelper: ProviderHelper = ({ providerModel }) => ({
5858
retrieve: () => usage,
5959
}
6060
},
61+
extractUsage: (response: any) => response.usageMetadata,
6162
normalizeUsage: (usage: Usage) => {
6263
const inputTokens = usage.promptTokenCount ?? 0
6364
const outputTokens = usage.candidatesTokenCount ?? 0

packages/console/app/src/routes/zen/util/provider/openai-compatible.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,7 @@ export const oaCompatHelper: ProviderHelper = ({ adjustCacheUsage }) => ({
5858
retrieve: () => usage,
5959
}
6060
},
61+
extractUsage: (response: any) => response.usage,
6162
normalizeUsage: (usage: Usage) => {
6263
let inputTokens = usage.prompt_tokens ?? 0
6364
const outputTokens = usage.completion_tokens ?? 0

packages/console/app/src/routes/zen/util/provider/openai.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@ export const openaiHelper: ProviderHelper = ({ workspaceID }) => ({
4343
retrieve: () => usage,
4444
}
4545
},
46+
extractUsage: (response: any) => response.usage ?? response.response?.usage,
4647
normalizeUsage: (usage: Usage) => {
4748
const inputTokens = usage.input_tokens ?? 0
4849
const outputTokens = usage.output_tokens ?? 0

packages/console/app/src/routes/zen/util/provider/provider.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,7 @@ export type ProviderHelper = (input: {
4949
parse: (chunk: string) => void
5050
retrieve: () => any
5151
}
52+
extractUsage: (response: any) => any
5253
normalizeUsage: (usage: any) => UsageInfo
5354
}
5455

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,68 @@
1+
import { describe, expect, test } from "bun:test"
2+
import type { ZenData } from "@opencode-ai/console-core/model.js"
3+
import type { ProviderHelper } from "../src/routes/zen/util/provider/provider"
4+
import { anthropicHelper } from "../src/routes/zen/util/provider/anthropic"
5+
import { googleHelper } from "../src/routes/zen/util/provider/google"
6+
import { oaCompatHelper } from "../src/routes/zen/util/provider/openai-compatible"
7+
import { openaiHelper } from "../src/routes/zen/util/provider/openai"
8+
9+
const providers = {
10+
anthropic: anthropicHelper({ reqModel: "claude-haiku-4-5", providerModel: "claude-haiku-4-5" }),
11+
google: googleHelper({ reqModel: "gemini-3-flash", providerModel: "gemini-3-flash" }),
12+
openai: openaiHelper({ reqModel: "gpt-5", providerModel: "gpt-5" }),
13+
"oa-compat": oaCompatHelper({ reqModel: "gpt-5-nano", providerModel: "gpt-5-nano" }),
14+
} satisfies Record<ZenData.Format, ReturnType<ProviderHelper>>
15+
16+
describe("provider usage extraction", () => {
17+
test("extracts Google non-stream usage metadata", () => {
18+
const usage = providers.google.extractUsage({
19+
usageMetadata: {
20+
promptTokenCount: 10,
21+
candidatesTokenCount: 3,
22+
thoughtsTokenCount: 2,
23+
cachedContentTokenCount: 4,
24+
},
25+
})
26+
27+
expect(providers.google.normalizeUsage(usage)).toEqual({
28+
inputTokens: 6,
29+
outputTokens: 3,
30+
reasoningTokens: 2,
31+
cacheReadTokens: 4,
32+
cacheWrite5mTokens: undefined,
33+
cacheWrite1hTokens: undefined,
34+
})
35+
})
36+
37+
test("parses Google stream usage metadata", () => {
38+
const usageParser = providers.google.createUsageParser()
39+
usageParser.parse(
40+
'data: {"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":3,"thoughtsTokenCount":2,"cachedContentTokenCount":4}}',
41+
)
42+
43+
expect(providers.google.normalizeUsage(usageParser.retrieve())).toEqual({
44+
inputTokens: 6,
45+
outputTokens: 3,
46+
reasoningTokens: 2,
47+
cacheReadTokens: 4,
48+
cacheWrite5mTokens: undefined,
49+
cacheWrite1hTokens: undefined,
50+
})
51+
})
52+
53+
test("extracts nested OpenAI Responses usage", () => {
54+
expect(
55+
providers.openai.extractUsage({
56+
response: {
57+
usage: {
58+
input_tokens: 5,
59+
output_tokens: 7,
60+
},
61+
},
62+
}),
63+
).toEqual({
64+
input_tokens: 5,
65+
output_tokens: 7,
66+
})
67+
})
68+
})

0 commit comments

Comments
 (0)