diff --git a/src/api/providers/__tests__/openai-usage-tracking.spec.ts b/src/api/providers/__tests__/openai-usage-tracking.spec.ts index e9c2f5e4fd..22ab2daf97 100644 --- a/src/api/providers/__tests__/openai-usage-tracking.spec.ts +++ b/src/api/providers/__tests__/openai-usage-tracking.spec.ts @@ -181,6 +181,31 @@ describe("OpenAiHandler with usage tracking fix", () => { }) }) + it("should report OpenAI-compatible cached prompt tokens", async () => { + mockCreate.mockImplementationOnce(async () => + asyncStreamFrom([ + { + choices: [{ delta: { content: "Cached response" }, index: 0 }], + usage: { + prompt_tokens: 5_053, + completion_tokens: 16, + total_tokens: 5_069, + prompt_tokens_details: { cached_tokens: 4_864 }, + }, + }, + ]), + ) + + const chunks = await collectStream(handler.createMessage(systemPrompt, messages)) + + expect(chunks).toContainEqual({ + type: "usage", + inputTokens: 5_053, + outputTokens: 16, + cacheReadTokens: 4_864, + }) + }) + it("should handle case where no usage is provided", async () => { // Override the mock for this specific test mockCreate.mockImplementationOnce(async (options) => { @@ -212,4 +237,160 @@ describe("OpenAiHandler with usage tracking fix", () => { expect(usageChunks).toHaveLength(0) }) }) + + it("should report cached prompt tokens from a non-streaming response", async () => { + const nonStreamingHandler = new OpenAiHandler({ ...mockOptions, openAiStreamingEnabled: false }) + mockCreate.mockImplementationOnce(async () => ({ + id: "test-completion", + choices: [{ message: { role: "assistant", content: "Cached response" } }], + usage: { + prompt_tokens: 4_621, + completion_tokens: 16, + total_tokens: 4_637, + prompt_tokens_details: { cached_tokens: 4_608 }, + }, + })) + + const chunks = await collectStream(nonStreamingHandler.createMessage("system prompt", [])) + + expect(chunks).toContainEqual({ + type: "usage", + inputTokens: 4_621, + outputTokens: 16, + cacheReadTokens: 4_608, + }) + }) + + it("reports cached prompt tokens for a streaming O3 response", async () => { + const o3Handler = new OpenAiHandler({ ...mockOptions, openAiModelId: "o3-mini" }) + mockCreate.mockImplementationOnce(async () => + asyncStreamFrom([ + { + choices: [{ delta: { content: "Cached response" }, index: 0 }], + usage: { + prompt_tokens: 5_053, + completion_tokens: 16, + total_tokens: 5_069, + prompt_tokens_details: { cached_tokens: 4_864 }, + }, + }, + ]), + ) + + const chunks = await collectStream(o3Handler.createMessage("system prompt", [])) + + expect(chunks).toContainEqual({ + type: "usage", + inputTokens: 5_053, + outputTokens: 16, + cacheReadTokens: 4_864, + }) + }) + + it.each([ + ["string", "10"], + ["object", { tokens: 10 }], + ["negative", -1], + ["fractional", 10.5], + ["non-finite", Number.POSITIVE_INFINITY], + ["greater than prompt tokens", 101], + ])("ignores invalid %s cached prompt tokens in streaming responses", async (_name, cachedTokens) => { + mockCreate.mockImplementationOnce(async () => + asyncStreamFrom([ + { + choices: [{ delta: { content: "Response" }, index: 0 }], + usage: { + prompt_tokens: 100, + completion_tokens: 5, + prompt_tokens_details: { cached_tokens: cachedTokens }, + }, + }, + ]), + ) + + const chunks = await collectStream(handler.createMessage("system prompt", [])) + + expect(chunks).toContainEqual({ type: "usage", inputTokens: 100, outputTokens: 5 }) + }) + + it.each([ + ["string", "10"], + ["object", { tokens: 10 }], + ["negative", -1], + ["fractional", 10.5], + ["non-finite", Number.POSITIVE_INFINITY], + ["greater than prompt tokens", 101], + ])("ignores invalid %s cached prompt tokens in non-streaming responses", async (_name, cachedTokens) => { + const nonStreamingHandler = new OpenAiHandler({ ...mockOptions, openAiStreamingEnabled: false }) + mockCreate.mockImplementationOnce(async () => ({ + id: "test-completion", + choices: [{ message: { role: "assistant", content: "Response" } }], + usage: { + prompt_tokens: 100, + completion_tokens: 5, + prompt_tokens_details: { cached_tokens: cachedTokens }, + }, + })) + + const chunks = await collectStream(nonStreamingHandler.createMessage("system prompt", [])) + + expect(chunks).toContainEqual({ type: "usage", inputTokens: 100, outputTokens: 5 }) + }) + + describe.each([true, false])("cache accounting with streaming=%s", (openAiStreamingEnabled) => { + it.each([ + ["authoritative count", 23, 23], + ["authoritative zero", 0, undefined], + ["fully cached prompt", 100, 100], + ["null falls back", null, 71], + ["absent falls back", undefined, 71], + ["negative", -1, undefined], + ["fractional", 10.5, undefined], + ["infinite", Infinity, undefined], + ["NaN", NaN, undefined], + ["string", "23", undefined], + ["object", { tokens: 23 }, undefined], + ["exceeds input", 101, undefined], + ])("respects %s without substituting a conflicting fallback", async (_name, reported, expected) => { + const cacheHandler = new OpenAiHandler({ ...mockOptions, openAiStreamingEnabled }) + const usage = { + prompt_tokens: 100, + completion_tokens: 5, + cache_creation_input_tokens: 7, + cache_read_input_tokens: reported, + prompt_tokens_details: { cached_tokens: 71 }, + } + mockCreate.mockResolvedValueOnce( + openAiStreamingEnabled + ? asyncStreamFrom([{ choices: [], usage }]) + : { choices: [{ message: { content: "Response" } }], usage }, + ) + + const chunks = await collectStream(cacheHandler.createMessage("system prompt", [])) + + expect(chunks.filter((chunk) => chunk.type === "usage")).toEqual([ + { + type: "usage", + inputTokens: 100, + outputTokens: 5, + cacheWriteTokens: 7, + cacheReadTokens: expected, + }, + ]) + }) + }) + + it.each([null, undefined, {}, { prompt_tokens_details: null }])( + "defaults missing non-streaming usage counters to zero: %j", + async (usage) => { + const nonStreamingHandler = new OpenAiHandler({ ...mockOptions, openAiStreamingEnabled: false }) + mockCreate.mockResolvedValueOnce({ choices: [{ message: { content: "Response" } }], usage }) + + const chunks = await collectStream(nonStreamingHandler.createMessage("system prompt", [])) + + expect(chunks.filter((chunk) => chunk.type === "usage")).toEqual([ + { type: "usage", inputTokens: 0, outputTokens: 0 }, + ]) + }, + ) }) diff --git a/src/api/providers/openai.ts b/src/api/providers/openai.ts index 619d05d28a..eb04e04312 100644 --- a/src/api/providers/openai.ts +++ b/src/api/providers/openai.ts @@ -28,6 +28,17 @@ import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata, Complete import { handleOpenAIError } from "./utils/error-handler" import { extractReasoningFromDelta } from "./utils/extract-reasoning" +type OpenAiUsage = + | { + prompt_tokens?: number + completion_tokens?: number + cache_creation_input_tokens?: number + cache_read_input_tokens?: unknown + prompt_tokens_details?: { cached_tokens?: unknown } | null + } + | null + | undefined + // TODO: Rename this to OpenAICompatibleHandler. Also, I think the // `OpenAINativeHandler` can subclass from this, since it's obviously // compatible with the OpenAI API. We can also rename it to `OpenAIHandler`. @@ -281,13 +292,24 @@ export class OpenAiHandler extends BaseProvider implements SingleCompletionHandl } } - protected processUsageMetrics(usage: any, _modelInfo?: ModelInfo): ApiStreamUsageChunk { + protected processUsageMetrics(usage: OpenAiUsage, _modelInfo?: ModelInfo): ApiStreamUsageChunk { + const inputTokens = usage?.prompt_tokens || 0 + const reportedCacheReadTokens = usage?.cache_read_input_tokens ?? usage?.prompt_tokens_details?.cached_tokens + const cacheReadTokens = + typeof reportedCacheReadTokens === "number" && + Number.isFinite(reportedCacheReadTokens) && + Number.isInteger(reportedCacheReadTokens) && + reportedCacheReadTokens >= 0 && + reportedCacheReadTokens <= inputTokens + ? reportedCacheReadTokens || undefined + : undefined + return { type: "usage", - inputTokens: usage?.prompt_tokens || 0, + inputTokens, outputTokens: usage?.completion_tokens || 0, cacheWriteTokens: usage?.cache_creation_input_tokens || undefined, - cacheReadTokens: usage?.cache_read_input_tokens || undefined, + cacheReadTokens, } } @@ -471,11 +493,7 @@ export class OpenAiHandler extends BaseProvider implements SingleCompletionHandl } if (chunk.usage) { - yield { - type: "usage", - inputTokens: chunk.usage.prompt_tokens || 0, - outputTokens: chunk.usage.completion_tokens || 0, - } + yield this.processUsageMetrics(chunk.usage) } } } diff --git a/src/eslint-suppressions.json b/src/eslint-suppressions.json index 93741e9174..bcb5947288 100644 --- a/src/eslint-suppressions.json +++ b/src/eslint-suppressions.json @@ -401,7 +401,7 @@ }, "api/providers/openai.ts": { "@typescript-eslint/no-explicit-any": { - "count": 3 + "count": 2 } }, "api/providers/openrouter.ts": {