Compare commits

...

1 Commits

Author SHA1 Message Date
Kit Langton
5e238ee8b9 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.
2026-03-18 15:28:50 -04:00
7 changed files with 116 additions and 12 deletions

View File

@@ -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,
})
}
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,
cost: calculateOccuredCost(billingSource, costInfo),
}),
)
const body = JSON.stringify(responseConverter(json))
logger.metric({ response_length: body.length })
logger.debug("RESPONSE: " + body)
dataDumper?.provideResponse(body)

View File

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

View File

@@ -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: () => {

View File

@@ -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: () => {

View File

@@ -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: () => {

View File

@@ -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: () => {

View File

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