fix(console): bill google non-stream zen usage (#28829)
This commit is contained in:
@@ -249,8 +249,9 @@ export async function handler(
|
|||||||
if (!isStream || [400, 404, 429].includes(res.status)) {
|
if (!isStream || [400, 404, 429].includes(res.status)) {
|
||||||
const json = await res.json()
|
const json = await res.json()
|
||||||
await rateLimiter?.track()
|
await rateLimiter?.track()
|
||||||
if (json.usage) {
|
const usage = providerInfo.extractUsage(json)
|
||||||
const usageInfo = providerInfo.normalizeUsage(json.usage)
|
if (usage) {
|
||||||
|
const usageInfo = providerInfo.normalizeUsage(usage)
|
||||||
const costInfo = calculateCost(modelInfo, usageInfo)
|
const costInfo = calculateCost(modelInfo, usageInfo)
|
||||||
await trialLimiter?.track(usageInfo)
|
await trialLimiter?.track(usageInfo)
|
||||||
await modelTpmLimiter?.track(providerInfo.id, providerInfo.model, usageInfo)
|
await modelTpmLimiter?.track(providerInfo.id, providerInfo.model, usageInfo)
|
||||||
|
|||||||
@@ -175,6 +175,7 @@ export const anthropicHelper: ProviderHelper = ({ reqModel, providerModel }) =>
|
|||||||
retrieve: () => usage,
|
retrieve: () => usage,
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
|
extractUsage: (response: any) => response.usage,
|
||||||
normalizeUsage: (usage: Usage) => ({
|
normalizeUsage: (usage: Usage) => ({
|
||||||
inputTokens: usage.input_tokens ?? 0,
|
inputTokens: usage.input_tokens ?? 0,
|
||||||
outputTokens: usage.output_tokens ?? 0,
|
outputTokens: usage.output_tokens ?? 0,
|
||||||
|
|||||||
@@ -58,6 +58,7 @@ export const googleHelper: ProviderHelper = ({ providerModel }) => ({
|
|||||||
retrieve: () => usage,
|
retrieve: () => usage,
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
|
extractUsage: (response: any) => response.usageMetadata,
|
||||||
normalizeUsage: (usage: Usage) => {
|
normalizeUsage: (usage: Usage) => {
|
||||||
const inputTokens = usage.promptTokenCount ?? 0
|
const inputTokens = usage.promptTokenCount ?? 0
|
||||||
const outputTokens = usage.candidatesTokenCount ?? 0
|
const outputTokens = usage.candidatesTokenCount ?? 0
|
||||||
|
|||||||
@@ -58,6 +58,7 @@ export const oaCompatHelper: ProviderHelper = ({ adjustCacheUsage }) => ({
|
|||||||
retrieve: () => usage,
|
retrieve: () => usage,
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
|
extractUsage: (response: any) => response.usage,
|
||||||
normalizeUsage: (usage: Usage) => {
|
normalizeUsage: (usage: Usage) => {
|
||||||
let inputTokens = usage.prompt_tokens ?? 0
|
let inputTokens = usage.prompt_tokens ?? 0
|
||||||
const outputTokens = usage.completion_tokens ?? 0
|
const outputTokens = usage.completion_tokens ?? 0
|
||||||
|
|||||||
@@ -43,6 +43,7 @@ export const openaiHelper: ProviderHelper = ({ workspaceID }) => ({
|
|||||||
retrieve: () => usage,
|
retrieve: () => usage,
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
|
extractUsage: (response: any) => response.usage ?? response.response?.usage,
|
||||||
normalizeUsage: (usage: Usage) => {
|
normalizeUsage: (usage: Usage) => {
|
||||||
const inputTokens = usage.input_tokens ?? 0
|
const inputTokens = usage.input_tokens ?? 0
|
||||||
const outputTokens = usage.output_tokens ?? 0
|
const outputTokens = usage.output_tokens ?? 0
|
||||||
|
|||||||
@@ -49,6 +49,7 @@ export type ProviderHelper = (input: {
|
|||||||
parse: (chunk: string) => void
|
parse: (chunk: string) => void
|
||||||
retrieve: () => any
|
retrieve: () => any
|
||||||
}
|
}
|
||||||
|
extractUsage: (response: any) => any
|
||||||
normalizeUsage: (usage: any) => UsageInfo
|
normalizeUsage: (usage: any) => UsageInfo
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -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,
|
||||||
|
})
|
||||||
|
})
|
||||||
|
})
|
||||||
Reference in New Issue
Block a user