From fe11328b0d31d8666c8e9b662c00097291132483 Mon Sep 17 00:00:00 2001 From: Mario Zechner Date: Tue, 6 Oct 2026 18:12:13 +0000 Subject: [PATCH] fix(ai): compare faux prompt-cache usage per message The faux provider's prompt-cache usage estimate compared the previous and current prompt as two joined strings, character by character, on every call. It now compares the prompts message by message and compares characters only from the first differing message. Usage numbers are unchanged. --- packages/ai/CHANGELOG.md | 4 +++ packages/ai/src/providers/faux.ts | 37 +++++++++++++++++++------- packages/ai/test/faux-provider.test.ts | 21 +++++++++++++++ 3 files changed, 52 insertions(+), 10 deletions(-) diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index efc4386d8..7c1ed3af6 100644 --- a/packages/ai/CHANGELOG.md +++ b/packages/ai/CHANGELOG.md @@ -11,6 +11,10 @@ - Added `durationMs` to `AssistantMessage`: `AssistantMessageEventStream` measures each response with a monotonic clock from the start of the request to its final message, for every API implementation, including direct calls. Deferred results fetched later stay untimed ([#10549](https://github.com/earendil-works/pi/issues/10549)) - Added an optional `durationMs` to `ToolResultMessage` for the execution time of the tool ([#10549](https://github.com/earendil-works/pi/issues/10549)) +### Changed + +- The faux provider's prompt-cache usage estimate compares the previous and current prompt message by message and compares characters only from the first differing message; the usage numbers are unchanged. + ### Fixed - Fixed `server_busy` and `servers are currently busy` provider errors ending the turn instead of being retried ([#10543](https://github.com/earendil-works/pi/issues/10543)) diff --git a/packages/ai/src/providers/faux.ts b/packages/ai/src/providers/faux.ts index a38c18774..c3771cc2c 100644 --- a/packages/ai/src/providers/faux.ts +++ b/packages/ai/src/providers/faux.ts @@ -214,8 +214,23 @@ function messageToText(message: Message): string { return toolResultToText(message); } -function serializeContext(context: TranscriptContext): string { - return context.messages.map((message) => `${message.role}:${messageToText(message)}`).join("\n\n"); +/** Length of the prompt text that joins `messages` with blank lines. */ +function joinedLength(messages: readonly string[], count = messages.length): number { + let length = count > 0 ? (count - 1) * 2 : 0; + for (let index = 0; index < count; index++) length += messages[index]!.length; + return length; +} + +/** + * Length of the common prefix of the two joined prompts. Equal messages are compared whole; characters are compared + * only from the first message that differs. + */ +function commonPromptPrefixLength(previous: readonly string[], current: readonly string[]): number { + let index = 0; + while (index < previous.length && index < current.length && previous[index] === current[index]) index++; + const rest = (messages: readonly string[]) => + index === messages.length ? "" : (index > 0 ? "\n\n" : "") + messages.slice(index).join("\n\n"); + return joinedLength(previous, index) + commonPrefixLength(rest(previous), rest(current)); } function commonPrefixLength(a: string, b: string): number { @@ -231,10 +246,12 @@ function withUsageEstimate( message: AssistantMessage, context: TranscriptContext, options: StreamOptions | undefined, - promptCache: Map, + promptCache: Map, ): AssistantMessage { - const promptText = serializeContext(context); - const promptTokens = estimateTokens(promptText); + // One text per message; the whole prompt joins them with blank lines. + const prompt = context.messages.map((message) => `${message.role}:${messageToText(message)}`); + const promptLength = joinedLength(prompt); + const promptTokens = Math.ceil(promptLength / 4); const outputTokens = estimateTokens(assistantContentToText(message.content)); let input = promptTokens; let cacheRead = 0; @@ -244,14 +261,14 @@ function withUsageEstimate( if (sessionId && options?.cacheRetention !== "none") { const previousPrompt = promptCache.get(sessionId); if (previousPrompt) { - const cachedChars = commonPrefixLength(previousPrompt, promptText); - cacheRead = estimateTokens(previousPrompt.slice(0, cachedChars)); - cacheWrite = estimateTokens(promptText.slice(cachedChars)); + const cachedChars = commonPromptPrefixLength(previousPrompt, prompt); + cacheRead = Math.ceil(cachedChars / 4); + cacheWrite = Math.ceil((promptLength - cachedChars) / 4); input = Math.max(0, promptTokens - cacheRead); } else { cacheWrite = promptTokens; } - promptCache.set(sessionId, promptText); + promptCache.set(sessionId, prompt); } return { @@ -445,7 +462,7 @@ export function createFauxCore(options: RegisterFauxProviderOptions) { let pendingResponses: FauxResponseStep[] = []; const tokensPerSecond = options.tokensPerSecond; const state: FauxProviderState = { callCount: 0, deferredFetchCount: 0, cancelledDeferred: [] }; - const promptCache = new Map(); + const promptCache = new Map(); const deferredResponses = new Map< string, { diff --git a/packages/ai/test/faux-provider.test.ts b/packages/ai/test/faux-provider.test.ts index 557885f23..6f4178111 100644 --- a/packages/ai/test/faux-provider.test.ts +++ b/packages/ai/test/faux-provider.test.ts @@ -324,6 +324,27 @@ describe("faux provider", () => { expect(second.usage.input + second.usage.cacheRead).toBeGreaterThan(second.usage.input); }); + it("counts cached characters up to the first difference in the joined prompt", async () => { + const registration = registerFauxProvider(); + registrations.push(registration); + registration.setResponses([fauxAssistantMessage("a"), fauxAssistantMessage("b"), fauxAssistantMessage("c")]); + const options = { sessionId: "session-1", cacheRetention: "short" } as const; + const user = (content: string) => ({ role: "user" as const, content, timestamp: 1 }); + + // Prompt texts: "user:hello world" (16 characters), then 16 + 2 + "user:next" (9) = 27. + await complete(registration.getModel(), { messages: [user("hello world")] }, options); + const extended = await complete( + registration.getModel(), + { messages: [user("hello world"), user("next")] }, + options, + ); + expect(extended.usage).toMatchObject({ input: 3, cacheRead: 4, cacheWrite: 3 }); + + // The first message now differs after "user:hello w" (12 characters): "user:hello wide" + 2 + 9 = 26. + const edited = await complete(registration.getModel(), { messages: [user("hello wide"), user("next")] }, options); + expect(edited.usage).toMatchObject({ input: 4, cacheRead: 3, cacheWrite: 4 }); + }); + it("does not simulate caching when cacheRetention is none", async () => { const registration = registerFauxProvider(); registrations.push(registration);