ソースを参照

fix(zen): read nested usage in non-stream responses

Extract usage from each provider's actual response shape so OpenAI Responses payloads nested under response.usage do not crash Zen's non-stream tracking path.
Kit Langton 4 週間 前
コミット
5e238ee8b9

+ 32 - 12
packages/console/app/src/routes/zen/util/handler.ts

@@ -219,20 +219,40 @@ export async function handler(
     // Handle non-streaming response
     if (!isStream) {
       const json = await res.json()
-      const usageInfo = providerInfo.normalizeUsage(json.usage)
-      const costInfo = calculateCost(modelInfo, usageInfo)
-      await trialLimiter?.track(usageInfo)
-      await rateLimiter?.track()
-      await trackUsage(sessionId, billingSource, authInfo, modelInfo, providerInfo, usageInfo, costInfo)
-      await reload(billingSource, authInfo, costInfo)
+      const usage = providerInfo.extractBodyUsage(json)
+
+      if (usage) {
+        const usageInfo = providerInfo.normalizeUsage(usage)
+        const costInfo = calculateCost(modelInfo, usageInfo)
+        await trialLimiter?.track(usageInfo)
+        await rateLimiter?.track()
+        await trackUsage(sessionId, billingSource, authInfo, modelInfo, providerInfo, usageInfo, costInfo)
+        await reload(billingSource, authInfo, costInfo)
+
+        const responseConverter = createResponseConverter(providerInfo.format, opts.format)
+        const body = JSON.stringify(
+          responseConverter({
+            ...json,
+            cost: calculateOccuredCost(billingSource, costInfo),
+          }),
+        )
+        logger.metric({ response_length: body.length })
+        logger.debug("RESPONSE: " + body)
+        dataDumper?.provideResponse(body)
+        dataDumper?.flush()
+        return new Response(body, {
+          status: resStatus,
+          statusText: res.statusText,
+          headers: resHeaders,
+        })
+      }
 
-      const responseConverter = createResponseConverter(providerInfo.format, opts.format)
-      const body = JSON.stringify(
-        responseConverter({
-          ...json,
-          cost: calculateOccuredCost(billingSource, costInfo),
-        }),
+      logger.debug(
+        "RESPONSE missing usage payload: " + JSON.stringify({ format: providerInfo.format, keys: Object.keys(json ?? {}) }),
       )
+
+      const responseConverter = createResponseConverter(providerInfo.format, opts.format)
+      const body = JSON.stringify(responseConverter(json))
       logger.metric({ response_length: body.length })
       logger.debug("RESPONSE: " + body)
       dataDumper?.provideResponse(body)

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

@@ -51,6 +51,7 @@ export const anthropicHelper: ProviderHelper = ({ reqModel, providerModel }) =>
             service_tier: "standard_only",
           }),
     }),
+    extractBodyUsage: (body: any) => body?.usage ?? body?.message?.usage,
     createBinaryStreamDecoder: () => {
       if (!isBedrock) return undefined
 

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

@@ -36,6 +36,7 @@ export const googleHelper: ProviderHelper = ({ providerModel }) => ({
   modifyBody: (body: Record<string, any>) => {
     return body
   },
+  extractBodyUsage: (body: any) => body?.usageMetadata,
   createBinaryStreamDecoder: () => undefined,
   streamSeparator: "\r\n\r\n",
   createUsageParser: () => {

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

@@ -34,6 +34,7 @@ export const oaCompatHelper: ProviderHelper = () => ({
       ...(body.stream ? { stream_options: { include_usage: true } } : {}),
     }
   },
+  extractBodyUsage: (body: any) => body?.usage,
   createBinaryStreamDecoder: () => undefined,
   streamSeparator: "\n\n",
   createUsageParser: () => {

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

@@ -22,6 +22,7 @@ export const openaiHelper: ProviderHelper = () => ({
     ...body,
     ...(workspaceID ? { safety_identifier: workspaceID } : {}),
   }),
+  extractBodyUsage: (body: any) => body?.usage ?? body?.response?.usage,
   createBinaryStreamDecoder: () => undefined,
   streamSeparator: "\n\n",
   createUsageParser: () => {

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

@@ -38,6 +38,7 @@ export type ProviderHelper = (input: { reqModel: string; providerModel: string }
   modifyUrl: (providerApi: string, isStream?: boolean) => string
   modifyHeaders: (headers: Headers, body: Record<string, any>, apiKey: string) => void
   modifyBody: (body: Record<string, any>, workspaceID?: string) => Record<string, any>
+  extractBodyUsage: (body: any) => any
   createBinaryStreamDecoder: () => ((chunk: Uint8Array) => Uint8Array | undefined) | undefined
   streamSeparator: string
   createUsageParser: () => {

+ 79 - 0
packages/console/app/test/provider-usage.test.ts

@@ -0,0 +1,79 @@
+import { describe, expect, test } from "bun:test"
+import { anthropicHelper } from "../src/routes/zen/util/provider/anthropic"
+import { googleHelper } from "../src/routes/zen/util/provider/google"
+import { openaiHelper } from "../src/routes/zen/util/provider/openai"
+import { oaCompatHelper } from "../src/routes/zen/util/provider/openai-compatible"
+
+describe("provider usage extraction", () => {
+  test("reads OpenAI Responses usage from response.usage", () => {
+    const helper = openaiHelper({ reqModel: "gpt-5.4", providerModel: "gpt-5.4" })
+
+    expect(
+      helper.extractBodyUsage({
+        response: {
+          usage: {
+            input_tokens: 13,
+            input_tokens_details: { cached_tokens: 3 },
+            output_tokens: 5,
+            output_tokens_details: { reasoning_tokens: 1 },
+          },
+        },
+      }),
+    ).toEqual({
+      input_tokens: 13,
+      input_tokens_details: { cached_tokens: 3 },
+      output_tokens: 5,
+      output_tokens_details: { reasoning_tokens: 1 },
+    })
+  })
+
+  test("reads Anthropic usage from message.usage", () => {
+    const helper = anthropicHelper({ reqModel: "claude-sonnet", providerModel: "claude-sonnet-4-5" })
+
+    expect(
+      helper.extractBodyUsage({
+        message: {
+          usage: {
+            input_tokens: 10,
+            output_tokens: 4,
+          },
+        },
+      }),
+    ).toEqual({
+      input_tokens: 10,
+      output_tokens: 4,
+    })
+  })
+
+  test("reads OA-compatible usage from usage", () => {
+    const helper = oaCompatHelper({ reqModel: "gpt-4o-mini", providerModel: "gpt-4o-mini" })
+
+    expect(
+      helper.extractBodyUsage({
+        usage: {
+          prompt_tokens: 8,
+          completion_tokens: 2,
+        },
+      }),
+    ).toEqual({
+      prompt_tokens: 8,
+      completion_tokens: 2,
+    })
+  })
+
+  test("reads Google usage from usageMetadata", () => {
+    const helper = googleHelper({ reqModel: "gemini-2.5-flash", providerModel: "gemini-2.5-flash" })
+
+    expect(
+      helper.extractBodyUsage({
+        usageMetadata: {
+          promptTokenCount: 11,
+          candidatesTokenCount: 3,
+        },
+      }),
+    ).toEqual({
+      promptTokenCount: 11,
+      candidatesTokenCount: 3,
+    })
+  })
+})