Browse Source

fix(console): bill google non-stream zen usage (#28829)

Jack 2 months ago
parent
commit
e1406e05a3

+ 3 - 2
packages/console/app/src/routes/zen/util/handler.ts

@@ -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)

+ 1 - 0
packages/console/app/src/routes/zen/util/provider/anthropic.ts

@@ -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,

+ 1 - 0
packages/console/app/src/routes/zen/util/provider/google.ts

@@ -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

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

@@ -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

+ 1 - 0
packages/console/app/src/routes/zen/util/provider/openai.ts

@@ -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

+ 1 - 0
packages/console/app/src/routes/zen/util/provider/provider.ts

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

+ 68 - 0
packages/console/app/test/providerUsage.test.ts

@@ -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,
+    })
+  })
+})