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

packages/ai/CHANGELOG.md

@@ -2,6 +2,10 @@
 
 ## [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
 
 ### Added

packages/ai/src/providers/openai-completions.ts

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

packages/ai/test/openai-completions-tool-choice.test.ts

@@ -13,7 +13,7 @@ const mockState = vi.hoisted(() => ({
 				usage?: {
 					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 };
 				};
 		  }>
@@ -430,6 +430,91 @@ describe("openai-completions tool_choice", () => {
 		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 () => {
 		const model = getModel("openrouter", "deepseek/deepseek-r1")!;
 		let payload: unknown;

packages/ai/test/openrouter-cache-write-repro.test.ts

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