fix(console): bill google non-stream zen usage (#28829)
Some checks are pending
deploy / deploy (push) Waiting to run

This commit is contained in:
Jack 2026-05-26 01:07:50 +08:00 committed by GitHub
parent b5553839d0
commit e1406e05a3
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 76 additions and 2 deletions

View file

@ -249,8 +249,9 @@ export async function handler(
if (!isStream || [400, 404, 429].includes(res.status)) {
const json = await res.json()
await rateLimiter?.track()
if (json.usage) {
const usageInfo = providerInfo.normalizeUsage(json.usage)
const usage = providerInfo.extractUsage(json)
if (usage) {
const usageInfo = providerInfo.normalizeUsage(usage)
const costInfo = calculateCost(modelInfo, usageInfo)
await trialLimiter?.track(usageInfo)
await modelTpmLimiter?.track(providerInfo.id, providerInfo.model, usageInfo)

View file

@ -175,6 +175,7 @@ export const anthropicHelper: ProviderHelper = ({ reqModel, providerModel }) =>
retrieve: () => usage,
}
},
extractUsage: (response: any) => response.usage,
normalizeUsage: (usage: Usage) => ({
inputTokens: usage.input_tokens ?? 0,
outputTokens: usage.output_tokens ?? 0,

View file

@ -58,6 +58,7 @@ export const googleHelper: ProviderHelper = ({ providerModel }) => ({
retrieve: () => usage,
}
},
extractUsage: (response: any) => response.usageMetadata,
normalizeUsage: (usage: Usage) => {
const inputTokens = usage.promptTokenCount ?? 0
const outputTokens = usage.candidatesTokenCount ?? 0

View file

@ -58,6 +58,7 @@ export const oaCompatHelper: ProviderHelper = ({ adjustCacheUsage }) => ({
retrieve: () => usage,
}
},
extractUsage: (response: any) => response.usage,
normalizeUsage: (usage: Usage) => {
let inputTokens = usage.prompt_tokens ?? 0
const outputTokens = usage.completion_tokens ?? 0

View file

@ -43,6 +43,7 @@ export const openaiHelper: ProviderHelper = ({ workspaceID }) => ({
retrieve: () => usage,
}
},
extractUsage: (response: any) => response.usage ?? response.response?.usage,
normalizeUsage: (usage: Usage) => {
const inputTokens = usage.input_tokens ?? 0
const outputTokens = usage.output_tokens ?? 0

View file

@ -49,6 +49,7 @@ export type ProviderHelper = (input: {
parse: (chunk: string) => void
retrieve: () => any
}
extractUsage: (response: any) => any
normalizeUsage: (usage: any) => UsageInfo
}

View file

@ -0,0 +1,68 @@
import { describe, expect, test } from "bun:test"
import type { ZenData } from "@opencode-ai/console-core/model.js"
import type { ProviderHelper } from "../src/routes/zen/util/provider/provider"
import { anthropicHelper } from "../src/routes/zen/util/provider/anthropic"
import { googleHelper } from "../src/routes/zen/util/provider/google"
import { oaCompatHelper } from "../src/routes/zen/util/provider/openai-compatible"
import { openaiHelper } from "../src/routes/zen/util/provider/openai"
const providers = {
anthropic: anthropicHelper({ reqModel: "claude-haiku-4-5", providerModel: "claude-haiku-4-5" }),
google: googleHelper({ reqModel: "gemini-3-flash", providerModel: "gemini-3-flash" }),
openai: openaiHelper({ reqModel: "gpt-5", providerModel: "gpt-5" }),
"oa-compat": oaCompatHelper({ reqModel: "gpt-5-nano", providerModel: "gpt-5-nano" }),
} satisfies Record<ZenData.Format, ReturnType<ProviderHelper>>
describe("provider usage extraction", () => {
test("extracts Google non-stream usage metadata", () => {
const usage = providers.google.extractUsage({
usageMetadata: {
promptTokenCount: 10,
candidatesTokenCount: 3,
thoughtsTokenCount: 2,
cachedContentTokenCount: 4,
},
})
expect(providers.google.normalizeUsage(usage)).toEqual({
inputTokens: 6,
outputTokens: 3,
reasoningTokens: 2,
cacheReadTokens: 4,
cacheWrite5mTokens: undefined,
cacheWrite1hTokens: undefined,
})
})
test("parses Google stream usage metadata", () => {
const usageParser = providers.google.createUsageParser()
usageParser.parse(
'data: {"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":3,"thoughtsTokenCount":2,"cachedContentTokenCount":4}}',
)
expect(providers.google.normalizeUsage(usageParser.retrieve())).toEqual({
inputTokens: 6,
outputTokens: 3,
reasoningTokens: 2,
cacheReadTokens: 4,
cacheWrite5mTokens: undefined,
cacheWrite1hTokens: undefined,
})
})
test("extracts nested OpenAI Responses usage", () => {
expect(
providers.openai.extractUsage({
response: {
usage: {
input_tokens: 5,
output_tokens: 7,
},
},
}),
).toEqual({
input_tokens: 5,
output_tokens: 7,
})
})
})