fix(ai): preserve cache_write_tokens in completions stream usage closes #2802

This commit is contained in:
Mario Zechner
2026-04-04 21:39:11 +02:00
parent a7acef92a7
commit 6044cabb15
4 changed files with 185 additions and 8 deletions

View File

@@ -734,24 +734,34 @@ function parseChunkUsage(
rawUsage: {
prompt_tokens?: number;
completion_tokens?: number;
prompt_tokens_details?: { cached_tokens?: number };
prompt_tokens_details?: { cached_tokens?: number; cache_write_tokens?: number };
completion_tokens_details?: { reasoning_tokens?: number };
},
model: Model<"openai-completions">,
): AssistantMessage["usage"] {
const cachedTokens = rawUsage.prompt_tokens_details?.cached_tokens || 0;
const promptTokens = rawUsage.prompt_tokens || 0;
const reportedCachedTokens = rawUsage.prompt_tokens_details?.cached_tokens || 0;
const cacheWriteTokens = rawUsage.prompt_tokens_details?.cache_write_tokens || 0;
const reasoningTokens = rawUsage.completion_tokens_details?.reasoning_tokens || 0;
// OpenAI includes cached tokens in prompt_tokens, so subtract to get non-cached input
const input = (rawUsage.prompt_tokens || 0) - cachedTokens;
// Normalize to pi-ai semantics:
// - cacheRead: hits from cache created by previous requests only
// - cacheWrite: tokens written to cache in this request
// Some OpenAI-compatible providers (observed on OpenRouter) report cached_tokens
// as (previous hits + current writes). In that case, remove cacheWrite from cacheRead.
const cacheReadTokens =
cacheWriteTokens > 0 ? Math.max(0, reportedCachedTokens - cacheWriteTokens) : reportedCachedTokens;
const input = Math.max(0, promptTokens - cacheReadTokens - cacheWriteTokens);
// Compute totalTokens ourselves since we add reasoning_tokens to output
// and some providers (e.g., Groq) don't include them in total_tokens
const outputTokens = (rawUsage.completion_tokens || 0) + reasoningTokens;
const usage: AssistantMessage["usage"] = {
input,
output: outputTokens,
cacheRead: cachedTokens,
cacheWrite: 0,
totalTokens: input + outputTokens + cachedTokens,
cacheRead: cacheReadTokens,
cacheWrite: cacheWriteTokens,
totalTokens: input + outputTokens + cacheReadTokens + cacheWriteTokens,
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
};
calculateCost(model, usage);