diff --git a/packages/console/app/src/routes/zen/util/handler.ts b/packages/console/app/src/routes/zen/util/handler.ts index 1996d48c567..d8bfc8ce338 100644 --- a/packages/console/app/src/routes/zen/util/handler.ts +++ b/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) diff --git a/packages/console/app/src/routes/zen/util/provider/anthropic.ts b/packages/console/app/src/routes/zen/util/provider/anthropic.ts index 15fe75b8481..993eebfd498 100644 --- a/packages/console/app/src/routes/zen/util/provider/anthropic.ts +++ b/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 diff --git a/packages/console/app/src/routes/zen/util/provider/google.ts b/packages/console/app/src/routes/zen/util/provider/google.ts index ecf3b2d4d4d..4d5446a9ee7 100644 --- a/packages/console/app/src/routes/zen/util/provider/google.ts +++ b/packages/console/app/src/routes/zen/util/provider/google.ts @@ -36,6 +36,7 @@ export const googleHelper: ProviderHelper = ({ providerModel }) => ({ modifyBody: (body: Record) => { return body }, + extractBodyUsage: (body: any) => body?.usageMetadata, createBinaryStreamDecoder: () => undefined, streamSeparator: "\r\n\r\n", createUsageParser: () => { diff --git a/packages/console/app/src/routes/zen/util/provider/openai-compatible.ts b/packages/console/app/src/routes/zen/util/provider/openai-compatible.ts index 046bf8f0c62..d826cf84715 100644 --- a/packages/console/app/src/routes/zen/util/provider/openai-compatible.ts +++ b/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: () => { diff --git a/packages/console/app/src/routes/zen/util/provider/openai.ts b/packages/console/app/src/routes/zen/util/provider/openai.ts index 596b38cc5a4..4a3a228992b 100644 --- a/packages/console/app/src/routes/zen/util/provider/openai.ts +++ b/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: () => { diff --git a/packages/console/app/src/routes/zen/util/provider/provider.ts b/packages/console/app/src/routes/zen/util/provider/provider.ts index 1f9492845f8..3f7eee80aba 100644 --- a/packages/console/app/src/routes/zen/util/provider/provider.ts +++ b/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, apiKey: string) => void modifyBody: (body: Record, workspaceID?: string) => Record + extractBodyUsage: (body: any) => any createBinaryStreamDecoder: () => ((chunk: Uint8Array) => Uint8Array | undefined) | undefined streamSeparator: string createUsageParser: () => { diff --git a/packages/console/app/test/provider-usage.test.ts b/packages/console/app/test/provider-usage.test.ts new file mode 100644 index 00000000000..a11d212110a --- /dev/null +++ b/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, + }) + }) +})