diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index 9d0f38ed..834ad5db 100644 --- a/packages/ai/CHANGELOG.md +++ b/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 diff --git a/packages/ai/src/providers/openai-completions.ts b/packages/ai/src/providers/openai-completions.ts index 79446578..ef4968dd 100644 --- a/packages/ai/src/providers/openai-completions.ts +++ b/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); diff --git a/packages/ai/test/openai-completions-tool-choice.test.ts b/packages/ai/test/openai-completions-tool-choice.test.ts index cd174897..62185d3e 100644 --- a/packages/ai/test/openai-completions-tool-choice.test.ts +++ b/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; diff --git a/packages/ai/test/openrouter-cache-write-repro.test.ts b/packages/ai/test/openrouter-cache-write-repro.test.ts new file mode 100644 index 00000000..ca492967 --- /dev/null +++ b/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); + }, + ); +});