fix(ai): preserve cache_write_tokens in completions stream usage closes #2802
This commit is contained in:
@@ -2,6 +2,10 @@
|
|||||||
|
|
||||||
## [Unreleased]
|
## [Unreleased]
|
||||||
|
|
||||||
|
### Fixed
|
||||||
|
|
||||||
|
- Fixed OpenAI-compatible completions streaming usage to preserve `prompt_tokens_details.cache_write_tokens` and normalize OpenRouter `cached_tokens` to previous-request cache hits only, preventing cache read/write double counting in `usage` and cost calculation ([#2802](https://github.com/badlogic/pi-mono/issues/2802))
|
||||||
|
|
||||||
## [0.65.0] - 2026-04-03
|
## [0.65.0] - 2026-04-03
|
||||||
|
|
||||||
### Added
|
### Added
|
||||||
|
|||||||
@@ -734,24 +734,34 @@ function parseChunkUsage(
|
|||||||
rawUsage: {
|
rawUsage: {
|
||||||
prompt_tokens?: number;
|
prompt_tokens?: number;
|
||||||
completion_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 };
|
completion_tokens_details?: { reasoning_tokens?: number };
|
||||||
},
|
},
|
||||||
model: Model<"openai-completions">,
|
model: Model<"openai-completions">,
|
||||||
): AssistantMessage["usage"] {
|
): 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;
|
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
|
// Compute totalTokens ourselves since we add reasoning_tokens to output
|
||||||
// and some providers (e.g., Groq) don't include them in total_tokens
|
// and some providers (e.g., Groq) don't include them in total_tokens
|
||||||
const outputTokens = (rawUsage.completion_tokens || 0) + reasoningTokens;
|
const outputTokens = (rawUsage.completion_tokens || 0) + reasoningTokens;
|
||||||
const usage: AssistantMessage["usage"] = {
|
const usage: AssistantMessage["usage"] = {
|
||||||
input,
|
input,
|
||||||
output: outputTokens,
|
output: outputTokens,
|
||||||
cacheRead: cachedTokens,
|
cacheRead: cacheReadTokens,
|
||||||
cacheWrite: 0,
|
cacheWrite: cacheWriteTokens,
|
||||||
totalTokens: input + outputTokens + cachedTokens,
|
totalTokens: input + outputTokens + cacheReadTokens + cacheWriteTokens,
|
||||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||||
};
|
};
|
||||||
calculateCost(model, usage);
|
calculateCost(model, usage);
|
||||||
|
|||||||
@@ -13,7 +13,7 @@ const mockState = vi.hoisted(() => ({
|
|||||||
usage?: {
|
usage?: {
|
||||||
prompt_tokens: number;
|
prompt_tokens: number;
|
||||||
completion_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 };
|
completion_tokens_details: { reasoning_tokens: number };
|
||||||
};
|
};
|
||||||
}>
|
}>
|
||||||
@@ -430,6 +430,91 @@ describe("openai-completions tool_choice", () => {
|
|||||||
expect(response.content).toEqual([{ type: "text", text: "OK" }]);
|
expect(response.content).toEqual([{ type: "text", text: "OK" }]);
|
||||||
});
|
});
|
||||||
|
|
||||||
|
it("preserves prompt_tokens_details.cache_write_tokens from chunk usage", async () => {
|
||||||
|
mockState.chunks = [
|
||||||
|
{
|
||||||
|
id: "chatcmpl-cache-write",
|
||||||
|
choices: [{ delta: { content: "OK" }, finish_reason: null }],
|
||||||
|
},
|
||||||
|
{
|
||||||
|
id: "chatcmpl-cache-write",
|
||||||
|
choices: [{ delta: {}, finish_reason: "stop" }],
|
||||||
|
usage: {
|
||||||
|
prompt_tokens: 100,
|
||||||
|
completion_tokens: 5,
|
||||||
|
prompt_tokens_details: { cached_tokens: 50, cache_write_tokens: 30 },
|
||||||
|
completion_tokens_details: { reasoning_tokens: 0 },
|
||||||
|
},
|
||||||
|
},
|
||||||
|
];
|
||||||
|
|
||||||
|
const { compat: _compat, ...baseModel } = getModel("openai", "gpt-4o-mini")!;
|
||||||
|
const model = { ...baseModel, api: "openai-completions" } as const;
|
||||||
|
const response = await streamSimple(
|
||||||
|
model,
|
||||||
|
{
|
||||||
|
messages: [
|
||||||
|
{
|
||||||
|
role: "user",
|
||||||
|
content: "Reply with exactly OK",
|
||||||
|
timestamp: Date.now(),
|
||||||
|
},
|
||||||
|
],
|
||||||
|
},
|
||||||
|
{ apiKey: "test" },
|
||||||
|
).result();
|
||||||
|
|
||||||
|
expect(response.usage.input).toBe(50);
|
||||||
|
expect(response.usage.cacheRead).toBe(20);
|
||||||
|
expect(response.usage.cacheWrite).toBe(30);
|
||||||
|
expect(response.usage.totalTokens).toBe(105);
|
||||||
|
});
|
||||||
|
|
||||||
|
it("preserves prompt_tokens_details.cache_write_tokens from choice usage fallback", async () => {
|
||||||
|
mockState.chunks = [
|
||||||
|
{
|
||||||
|
id: "chatcmpl-cache-write-choice",
|
||||||
|
choices: [{ delta: { content: "OK" }, finish_reason: null }],
|
||||||
|
},
|
||||||
|
{
|
||||||
|
id: "chatcmpl-cache-write-choice",
|
||||||
|
choices: [
|
||||||
|
{
|
||||||
|
delta: {},
|
||||||
|
finish_reason: "stop",
|
||||||
|
usage: {
|
||||||
|
prompt_tokens: 100,
|
||||||
|
completion_tokens: 5,
|
||||||
|
prompt_tokens_details: { cached_tokens: 50, cache_write_tokens: 30 },
|
||||||
|
completion_tokens_details: { reasoning_tokens: 0 },
|
||||||
|
},
|
||||||
|
},
|
||||||
|
],
|
||||||
|
},
|
||||||
|
];
|
||||||
|
|
||||||
|
const { compat: _compat, ...baseModel } = getModel("openai", "gpt-4o-mini")!;
|
||||||
|
const model = { ...baseModel, api: "openai-completions" } as const;
|
||||||
|
const response = await streamSimple(
|
||||||
|
model,
|
||||||
|
{
|
||||||
|
messages: [
|
||||||
|
{
|
||||||
|
role: "user",
|
||||||
|
content: "Reply with exactly OK",
|
||||||
|
timestamp: Date.now(),
|
||||||
|
},
|
||||||
|
],
|
||||||
|
},
|
||||||
|
{ apiKey: "test" },
|
||||||
|
).result();
|
||||||
|
|
||||||
|
expect(response.usage.input).toBe(50);
|
||||||
|
expect(response.usage.cacheRead).toBe(20);
|
||||||
|
expect(response.usage.cacheWrite).toBe(30);
|
||||||
|
expect(response.usage.totalTokens).toBe(105);
|
||||||
|
});
|
||||||
|
|
||||||
it("uses OpenRouter reasoning object instead of reasoning_effort", async () => {
|
it("uses OpenRouter reasoning object instead of reasoning_effort", async () => {
|
||||||
const model = getModel("openrouter", "deepseek/deepseek-r1")!;
|
const model = getModel("openrouter", "deepseek/deepseek-r1")!;
|
||||||
let payload: unknown;
|
let payload: unknown;
|
||||||
|
|||||||
78
packages/ai/test/openrouter-cache-write-repro.test.ts
Normal file
78
packages/ai/test/openrouter-cache-write-repro.test.ts
Normal file
@@ -0,0 +1,78 @@
|
|||||||
|
import { describe, expect, it } from "vitest";
|
||||||
|
import { getModel } from "../src/models.js";
|
||||||
|
import { completeSimple } from "../src/stream.js";
|
||||||
|
|
||||||
|
function createLongSystemPrompt(): string {
|
||||||
|
const nonce = `${Date.now()}-${Math.random()}`;
|
||||||
|
return `You are a concise assistant.\nCache nonce: ${nonce}\n\n${Array(80)
|
||||||
|
.fill(
|
||||||
|
"Prompt-caching probe content. Keep this exact text stable across requests so the provider can reuse prefix tokens and report cache read and cache write usage.",
|
||||||
|
)
|
||||||
|
.join("\n\n")}`;
|
||||||
|
}
|
||||||
|
|
||||||
|
describe.skipIf(!process.env.OPENROUTER_API_KEY)("OpenRouter cache_write repro E2E", () => {
|
||||||
|
it(
|
||||||
|
"regression: preserves cache_write_tokens on openai-completions stream path",
|
||||||
|
{ retry: 2, timeout: 90000 },
|
||||||
|
async () => {
|
||||||
|
const model = getModel("openrouter", "google/gemini-2.5-flash");
|
||||||
|
const context = {
|
||||||
|
systemPrompt: createLongSystemPrompt(),
|
||||||
|
messages: [
|
||||||
|
{
|
||||||
|
role: "user" as const,
|
||||||
|
content: "Reply with exactly: OK",
|
||||||
|
timestamp: Date.now(),
|
||||||
|
},
|
||||||
|
],
|
||||||
|
};
|
||||||
|
|
||||||
|
const options = {
|
||||||
|
apiKey: process.env.OPENROUTER_API_KEY!,
|
||||||
|
maxTokens: 32,
|
||||||
|
temperature: 0,
|
||||||
|
onPayload: (payload: unknown) => {
|
||||||
|
const params = payload as {
|
||||||
|
messages?: Array<{
|
||||||
|
role?: string;
|
||||||
|
content?: string | Array<{ type?: string; text?: string; cache_control?: { type: string } }>;
|
||||||
|
}>;
|
||||||
|
};
|
||||||
|
const messages = params.messages;
|
||||||
|
if (!Array.isArray(messages)) return payload;
|
||||||
|
|
||||||
|
for (let i = messages.length - 1; i >= 0; i--) {
|
||||||
|
const msg = messages[i];
|
||||||
|
if (msg.role !== "user") continue;
|
||||||
|
if (typeof msg.content === "string") {
|
||||||
|
msg.content = [{ type: "text", text: msg.content, cache_control: { type: "ephemeral" } }];
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
if (!Array.isArray(msg.content)) continue;
|
||||||
|
for (let j = msg.content.length - 1; j >= 0; j--) {
|
||||||
|
const part = msg.content[j];
|
||||||
|
if (part.type === "text") {
|
||||||
|
part.cache_control = { type: "ephemeral" };
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
return payload;
|
||||||
|
},
|
||||||
|
};
|
||||||
|
|
||||||
|
const first = await completeSimple(model, context, options);
|
||||||
|
expect(first.stopReason, first.errorMessage).toBe("stop");
|
||||||
|
|
||||||
|
const second = await completeSimple(model, context, options);
|
||||||
|
expect(second.stopReason, second.errorMessage).toBe("stop");
|
||||||
|
|
||||||
|
// Regression expectation: cache_write_tokens from provider usage must be preserved.
|
||||||
|
// With the cache_control marker above, at least one of the two calls should create cache.
|
||||||
|
const hasCacheWrite = first.usage.cacheWrite > 0 || second.usage.cacheWrite > 0;
|
||||||
|
expect(hasCacheWrite).toBe(true);
|
||||||
|
},
|
||||||
|
);
|
||||||
|
});
|
||||||
Reference in New Issue
Block a user