Compare commits

...

1 Commits

Author SHA1 Message Date
Aiden Cline e7b56e8122 fix(ai): normalize loose provider usage 2026-08-05 18:43:02 +00:00
11 changed files with 214 additions and 66 deletions
@@ -237,7 +237,7 @@ export type AnthropicMessagesBody = Schema.Schema.Type<typeof AnthropicMessagesB
const AnthropicUsage = Schema.StructWithRest(
Schema.Struct({
input_tokens: Schema.optional(Schema.Number),
input_tokens: optionalNull(Schema.Number),
output_tokens: Schema.optional(Schema.Number),
cache_creation_input_tokens: optionalNull(Schema.Number),
cache_read_input_tokens: optionalNull(Schema.Number),
@@ -684,7 +684,7 @@ const mapFinishReason = (reason: string | null | undefined): FinishReason => {
// expose that subset through `output_tokens_details.thinking_tokens`.
const mapUsage = (usage: AnthropicUsage | undefined): Usage | undefined => {
if (!usage) return undefined
const nonCached = usage.input_tokens
const nonCached = usage.input_tokens ?? undefined
const cacheRead = usage.cache_read_input_tokens ?? undefined
const cacheWrite = usage.cache_creation_input_tokens ?? undefined
const inputTokens = ProviderShared.sumTokens(nonCached, cacheRead, cacheWrite)
+12 -10
View File
@@ -20,7 +20,7 @@ import {
} from "../schema"
import { BedrockEventStream } from "./bedrock-event-stream"
import { classifyProviderFailure } from "../provider-error"
import { JsonObject, optionalArray, ProviderShared } from "./shared"
import { JsonObject, optionalArray, optionalNull, ProviderShared } from "./shared"
import { BedrockAuth } from "./utils/bedrock-auth"
import { BedrockCache } from "./utils/bedrock-cache"
import { BedrockMedia } from "./utils/bedrock-media"
@@ -150,8 +150,8 @@ const BedrockUsageSchema = Schema.Struct({
inputTokens: Schema.optional(Schema.Number),
outputTokens: Schema.optional(Schema.Number),
totalTokens: Schema.optional(Schema.Number),
cacheReadInputTokens: Schema.optional(Schema.Number),
cacheWriteInputTokens: Schema.optional(Schema.Number),
cacheReadInputTokens: optionalNull(Schema.Number),
cacheWriteInputTokens: optionalNull(Schema.Number),
})
type BedrockUsageSchema = Schema.Schema.Type<typeof BedrockUsageSchema>
@@ -206,9 +206,9 @@ const BedrockEvent = Schema.Struct({
additionalModelResponseFields: Schema.optional(Schema.Unknown),
}),
),
metadata: Schema.optional(
metadata: optionalNull(
Schema.Struct({
usage: Schema.optional(BedrockUsageSchema),
usage: optionalNull(BedrockUsageSchema),
metrics: Schema.optional(Schema.Unknown),
}),
),
@@ -464,19 +464,21 @@ const mapFinishReason = (reason: string): FinishReason => {
// AWS reports inputTokens separately from cache reads and writes.
// Bedrock does not break reasoning out of outputTokens for current models.
const mapUsage = (usage: BedrockUsageSchema | undefined): Usage | undefined => {
const mapUsage = (usage: BedrockUsageSchema | null | undefined): Usage | undefined => {
if (!usage) return undefined
const cacheRead = usage.cacheReadInputTokens ?? undefined
const cacheWrite = usage.cacheWriteInputTokens ?? undefined
const inputTokens = ProviderShared.sumTokens(
usage.inputTokens,
usage.cacheReadInputTokens,
usage.cacheWriteInputTokens,
cacheRead,
cacheWrite,
)
return new Usage({
inputTokens,
outputTokens: usage.outputTokens,
nonCachedInputTokens: usage.inputTokens,
cacheReadInputTokens: usage.cacheReadInputTokens,
cacheWriteInputTokens: usage.cacheWriteInputTokens,
cacheReadInputTokens: cacheRead,
cacheWriteInputTokens: cacheWrite,
totalTokens: ProviderShared.totalTokens(inputTokens, usage.outputTokens, usage.totalTokens),
providerMetadata: { bedrock: usage },
})
@@ -76,12 +76,12 @@ const consumeFrames = (route: string) => (state: FrameBufferState, chunk: Uint8A
// before handing the object to the chunk schema. JSON decode goes
// through the shared Schema-driven codec to satisfy the package rule
// against ad-hoc `JSON.parse` calls.
const parsed = (yield* ProviderShared.parseJson(
const parsed = yield* ProviderShared.parseJson(
route,
payload,
"Failed to parse Bedrock Converse event-stream payload",
)) as Record<string, unknown>
delete parsed.p
)
if (ProviderShared.isRecord(parsed)) delete parsed.p
out.push({ [eventType]: parsed })
}
return [cursor, out] as const
+22 -17
View File
@@ -18,7 +18,7 @@ import {
type ToolCallPart,
type ToolDefinition,
} from "../schema"
import { JsonObject, optionalArray, ProviderShared } from "./shared"
import { JsonObject, optionalArray, optionalNull, ProviderShared } from "./shared"
import { GeminiToolSchema } from "./utils/gemini-tool-schema"
import { Lifecycle } from "./utils/lifecycle"
import { ToolSchemaProjection } from "./utils/tool-schema"
@@ -162,13 +162,16 @@ const GeminiBodyFields = {
const GeminiBody = Schema.Struct(GeminiBodyFields)
export type GeminiBody = Schema.Schema.Type<typeof GeminiBody>
const GeminiUsage = Schema.Struct({
cachedContentTokenCount: Schema.optional(Schema.Number),
thoughtsTokenCount: Schema.optional(Schema.Number),
promptTokenCount: Schema.optional(Schema.Number),
candidatesTokenCount: Schema.optional(Schema.Number),
totalTokenCount: Schema.optional(Schema.Number),
})
const GeminiUsage = Schema.StructWithRest(
Schema.Struct({
cachedContentTokenCount: optionalNull(Schema.Number),
thoughtsTokenCount: optionalNull(Schema.Number),
promptTokenCount: optionalNull(Schema.Number),
candidatesTokenCount: optionalNull(Schema.Number),
totalTokenCount: optionalNull(Schema.Number),
}),
[Schema.Record(Schema.String, Schema.Unknown)],
)
type GeminiUsage = Schema.Schema.Type<typeof GeminiUsage>
const GeminiCandidate = Schema.Struct({
@@ -178,7 +181,7 @@ const GeminiCandidate = Schema.Struct({
const GeminiEvent = Schema.Struct({
candidates: optionalArray(GeminiCandidate),
usageMetadata: Schema.optional(GeminiUsage),
usageMetadata: optionalNull(GeminiUsage),
})
type GeminiEvent = Schema.Schema.Type<typeof GeminiEvent>
@@ -422,23 +425,25 @@ const fromRequest = Effect.fn("Gemini.fromRequest")(function* (request: LLMReque
// `cachedContentTokenCount` subset. `candidatesTokenCount` is *exclusive*
// of `thoughtsTokenCount` — visible-only, not a total — so we sum the two
// to produce the inclusive `outputTokens` the rest of the contract expects.
const mapUsage = (usage: GeminiUsage | undefined) => {
const mapUsage = (usage: GeminiUsage | null | undefined) => {
if (!usage) return undefined
const cached = usage.cachedContentTokenCount
const nonCached = ProviderShared.subtractTokens(usage.promptTokenCount, cached)
const input = usage.promptTokenCount ?? undefined
const cached = input === undefined ? undefined : (usage.cachedContentTokenCount ?? undefined)
const visible = usage.candidatesTokenCount ?? undefined
const thoughts = visible === undefined ? undefined : (usage.thoughtsTokenCount ?? undefined)
const nonCached = ProviderShared.subtractTokens(input, cached)
// `candidatesTokenCount` is visible-only; sum with thoughts to produce the
// inclusive `outputTokens` the contract expects. Only compute the total
// when the visible component is reported — otherwise we'd fabricate an
// inclusive number from a partial breakdown.
const outputTokens =
usage.candidatesTokenCount !== undefined ? usage.candidatesTokenCount + (usage.thoughtsTokenCount ?? 0) : undefined
const outputTokens = visible === undefined ? undefined : visible + (thoughts ?? 0)
return new Usage({
inputTokens: usage.promptTokenCount,
inputTokens: input,
outputTokens,
nonCachedInputTokens: nonCached,
cacheReadInputTokens: cached,
reasoningTokens: usage.thoughtsTokenCount,
totalTokens: ProviderShared.totalTokens(usage.promptTokenCount, outputTokens, usage.totalTokenCount),
reasoningTokens: thoughts,
totalTokens: ProviderShared.totalTokens(input, outputTokens, usage.totalTokenCount ?? undefined),
providerMetadata: { google: usage },
})
}
+8 -6
View File
@@ -183,12 +183,12 @@ const OpenResponsesUsage = Schema.Struct({
input_tokens: Schema.optional(Schema.Number),
input_tokens_details: optionalNull(
Schema.Struct({
cached_tokens: Schema.optional(Schema.Number),
cache_write_tokens: Schema.optional(Schema.Number),
cached_tokens: optionalNull(Schema.Number),
cache_write_tokens: optionalNull(Schema.Number),
}),
),
output_tokens: Schema.optional(Schema.Number),
output_tokens_details: optionalNull(Schema.Struct({ reasoning_tokens: Schema.optional(Schema.Number) })),
output_tokens_details: optionalNull(Schema.Struct({ reasoning_tokens: optionalNull(Schema.Number) })),
total_tokens: Schema.optional(Schema.Number),
})
type OpenResponsesUsage = Schema.Schema.Type<typeof OpenResponsesUsage>
@@ -592,9 +592,11 @@ export const fromRequest = Effect.fn("OpenResponses.fromRequest")(function* (req
// non-cached breakdown.
const mapUsage = (usage: OpenResponsesUsage | null | undefined, providerMetadataKey: string) => {
if (!usage) return undefined
const cached = usage.input_tokens_details?.cached_tokens
const cacheWrite = usage.input_tokens_details?.cache_write_tokens
const reasoning = usage.output_tokens_details?.reasoning_tokens
const cached = usage.input_tokens === undefined ? undefined : (usage.input_tokens_details?.cached_tokens ?? undefined)
const cacheWrite =
usage.input_tokens === undefined ? undefined : (usage.input_tokens_details?.cache_write_tokens ?? undefined)
const reasoning =
usage.output_tokens === undefined ? undefined : (usage.output_tokens_details?.reasoning_tokens ?? undefined)
const nonCached = ProviderShared.subtractTokens(usage.input_tokens, ProviderShared.sumTokens(cached, cacheWrite))
return new Usage({
inputTokens: usage.input_tokens,
+23 -20
View File
@@ -146,18 +146,18 @@ export type OpenAIChatBody = Schema.Schema.Type<typeof OpenAIChatBody>
// byte stream into strings, then `Protocol.jsonEvent` decodes each string into
// this provider-native event shape.
const OpenAIChatUsage = Schema.Struct({
prompt_tokens: Schema.optional(Schema.Number),
completion_tokens: Schema.optional(Schema.Number),
total_tokens: Schema.optional(Schema.Number),
prompt_tokens: optionalNull(Schema.Number),
completion_tokens: optionalNull(Schema.Number),
total_tokens: optionalNull(Schema.Number),
prompt_tokens_details: optionalNull(
Schema.Struct({
cached_tokens: Schema.optional(Schema.Number),
cache_write_tokens: Schema.optional(Schema.Number),
cached_tokens: optionalNull(Schema.Number),
cache_write_tokens: optionalNull(Schema.Number),
}),
),
completion_tokens_details: optionalNull(
Schema.Struct({
reasoning_tokens: Schema.optional(Schema.Number),
reasoning_tokens: optionalNull(Schema.Number),
}),
),
})
@@ -168,7 +168,7 @@ const OpenAIChatToolCallDeltaFunction = Schema.Struct({
})
const OpenAIChatToolCallDelta = Schema.Struct({
index: Schema.Number,
index: optionalNull(Schema.Number),
id: optionalNull(Schema.String),
function: optionalNull(OpenAIChatToolCallDeltaFunction),
})
@@ -559,18 +559,20 @@ const mapFinishReason = (reason: string | null | undefined): FinishReason => {
// satisfied on both sides.
const mapUsage = (usage: OpenAIChatEvent["usage"]): Usage | undefined => {
if (!usage) return undefined
const cached = usage.prompt_tokens_details?.cached_tokens
const cacheWrite = usage.prompt_tokens_details?.cache_write_tokens
const reasoning = usage.completion_tokens_details?.reasoning_tokens
const nonCached = ProviderShared.subtractTokens(usage.prompt_tokens, ProviderShared.sumTokens(cached, cacheWrite))
const input = usage.prompt_tokens ?? undefined
const output = usage.completion_tokens ?? undefined
const cached = input === undefined ? undefined : (usage.prompt_tokens_details?.cached_tokens ?? undefined)
const cacheWrite = input === undefined ? undefined : (usage.prompt_tokens_details?.cache_write_tokens ?? undefined)
const reasoning = output === undefined ? undefined : (usage.completion_tokens_details?.reasoning_tokens ?? undefined)
const nonCached = ProviderShared.subtractTokens(input, ProviderShared.sumTokens(cached, cacheWrite))
return new Usage({
inputTokens: usage.prompt_tokens,
outputTokens: usage.completion_tokens,
inputTokens: input,
outputTokens: output,
nonCachedInputTokens: nonCached,
cacheReadInputTokens: cached,
cacheWriteInputTokens: cacheWrite,
reasoningTokens: reasoning,
totalTokens: ProviderShared.totalTokens(usage.prompt_tokens, usage.completion_tokens, usage.total_tokens),
totalTokens: ProviderShared.totalTokens(input, output, usage.total_tokens ?? undefined),
providerMetadata: { openai: usage },
})
}
@@ -694,24 +696,25 @@ const step = (state: ParserState, event: OpenAIChatEvent) =>
lifecycle = Lifecycle.textDelta(lifecycle, events, "text-0", delta.content)
}
for (const tool of toolDeltas) {
const current = tools[tool.index]
const pending = pendingTools[tool.index]
for (const [position, tool] of toolDeltas.entries()) {
const index = tool.index ?? position
const current = tools[index]
const pending = pendingTools[index]
const id = current?.id ?? pending?.id ?? (tool.id || undefined)
const name = current?.name ?? pending?.name ?? (tool.function?.name || undefined)
const text = `${pending?.input ?? ""}${tool.function?.arguments ?? ""}`
if (!current && (!id || !name)) {
pendingTools = { ...pendingTools, [tool.index]: { id: id || undefined, name: name || undefined, input: text } }
pendingTools = { ...pendingTools, [index]: { id: id || undefined, name: name || undefined, input: text } }
continue
}
if (pending) {
pendingTools = { ...pendingTools }
delete pendingTools[tool.index]
delete pendingTools[index]
}
const result = ToolStream.appendOrStart(
ADAPTER,
tools,
tool.index,
index,
{ id: id || undefined, name: name || undefined, text },
"OpenAI Chat tool call delta is missing id or name",
)
@@ -488,7 +488,7 @@ describe("Anthropic Messages route", () => {
{
type: "message_delta",
delta: { stop_reason: "end_turn", stop_sequence: "\n\nHuman:" },
usage: { output_tokens: 2 },
usage: { input_tokens: null, output_tokens: 2 },
},
{ type: "message_stop" },
)
@@ -388,12 +388,30 @@ describe("Bedrock Converse route", () => {
Effect.gen(function* () {
const body = eventStreamBody(
["messageStop", { stopReason: "end_turn" }],
["metadata", { usage: { inputTokens: 5, outputTokens: 2, totalTokens: 7 } }],
["metadata", { metrics: { latencyMs: 100 } }],
[
"metadata",
{
usage: {
inputTokens: 5,
outputTokens: 2,
totalTokens: 7,
cacheReadInputTokens: null,
cacheWriteInputTokens: null,
},
},
],
["metadata", { usage: null }],
["metadata", null],
)
const response = yield* LLMClient.generate(baseRequest).pipe(Effect.provide(fixedBytes(body)))
expect(response.usage).toMatchObject({ inputTokens: 5, outputTokens: 2, totalTokens: 7 })
expect(response.usage).toMatchObject({
inputTokens: 5,
outputTokens: 2,
totalTokens: 7,
cacheReadInputTokens: undefined,
cacheWriteInputTokens: undefined,
})
}),
)
+27 -4
View File
@@ -722,14 +722,37 @@ describe("Gemini route", () => {
}),
)
it.effect("leaves total usage undefined when component counts are missing", () =>
it.effect("keeps partial usage only in provider metadata", () =>
Effect.gen(function* () {
const response = yield* LLMClient.generate(request).pipe(
Effect.provide(fixedResponse(sseEvents({ usageMetadata: { thoughtsTokenCount: 1 } }))),
Effect.provide(
fixedResponse(
sseEvents({
usageMetadata: {
promptTokenCount: null,
candidatesTokenCount: null,
totalTokenCount: null,
thoughtsTokenCount: null,
cachedContentTokenCount: 1,
promptTokensDetails: [{ modality: "TEXT", tokenCount: 5 }],
candidatesTokensDetails: [{ modality: "TEXT", tokenCount: 2 }],
},
}),
),
),
)
expect(response.usage).toMatchObject({ reasoningTokens: 1 })
expect(response.usage?.totalTokens).toBeUndefined()
expect(response.usage).toMatchObject({
inputTokens: undefined,
outputTokens: undefined,
cacheReadInputTokens: undefined,
providerMetadata: {
google: {
promptTokensDetails: [{ modality: "TEXT", tokenCount: 5 }],
candidatesTokensDetails: [{ modality: "TEXT", tokenCount: 2 }],
},
},
})
}),
)
@@ -596,6 +596,37 @@ describe("OpenAI Chat route", () => {
}),
)
it.effect("accepts nullable usage counters", () =>
Effect.gen(function* () {
const body = sseEvents(
deltaChunk({ role: "assistant", content: "Hello" }),
deltaChunk({}, "stop"),
usageChunk({
prompt_tokens: null,
completion_tokens: null,
total_tokens: null,
prompt_tokens_details: { cached_tokens: 1, cache_write_tokens: null },
completion_tokens_details: { reasoning_tokens: null },
}),
)
const response = yield* LLMClient.generate(request).pipe(Effect.provide(fixedResponse(body)))
expect(response.usage).toEqual(
new Usage({
providerMetadata: {
openai: {
prompt_tokens: null,
completion_tokens: null,
total_tokens: null,
prompt_tokens_details: { cached_tokens: 1, cache_write_tokens: null },
completion_tokens_details: { reasoning_tokens: null },
},
},
}),
)
}),
)
it.effect("parses and replays OpenAI-compatible reasoning fields", () =>
Effect.gen(function* () {
const fields = ["reasoning_content", "reasoning", "reasoning_text"] as const
@@ -1048,6 +1079,36 @@ describe("OpenAI Chat route", () => {
}),
)
it.effect("assembles indexless streamed tool calls", () =>
Effect.gen(function* () {
const body = sseEvents(
deltaChunk({
tool_calls: [
{ id: "call_1", function: { name: "lookup", arguments: '{"query"' } },
{ index: null, id: "call_2", function: { name: "lookup", arguments: '{"query"' } },
],
}),
deltaChunk({
tool_calls: [
{ function: { arguments: ':"weather"}' } },
{ index: null, function: { arguments: ':"time"}' } },
],
}),
deltaChunk({}, "tool_calls"),
)
const response = yield* LLMClient.generate(
LLMRequest.update(request, {
tools: [ToolDefinition.make({ name: "lookup", description: "Lookup data", inputSchema: { type: "object" } })],
}),
).pipe(Effect.provide(fixedResponse(body)))
expect(response.toolCalls).toMatchObject([
{ id: "call_1", name: "lookup", input: { query: "weather" } },
{ id: "call_2", name: "lookup", input: { query: "time" } },
])
}),
)
it.effect("ignores empty identity fields on later tool call deltas", () =>
Effect.gen(function* () {
const body = sseEvents(
@@ -885,6 +885,40 @@ describe("OpenAI Responses route", () => {
}),
)
it.effect("accepts nullable token details", () =>
Effect.gen(function* () {
const response = yield* LLMClient.generate(request).pipe(
Effect.provide(
fixedResponse(
sseEvents({
type: "response.completed",
response: {
id: "resp_1",
usage: {
input_tokens: 5,
output_tokens: 2,
total_tokens: 7,
input_tokens_details: { cached_tokens: null, cache_write_tokens: null },
output_tokens_details: { reasoning_tokens: null },
},
},
}),
),
),
)
expect(response.usage).toMatchObject({
inputTokens: 5,
outputTokens: 2,
nonCachedInputTokens: 5,
totalTokens: 7,
cacheReadInputTokens: undefined,
cacheWriteInputTokens: undefined,
reasoningTokens: undefined,
})
}),
)
it.effect("preserves and replays assistant message phases", () =>
Effect.gen(function* () {
const response = yield* LLMClient.generate(request).pipe(