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.
This commit is contained in:
Mario Zechner 2026-10-06 18:12:13 +00:00
parent 4c28a6865c
commit fe11328b0d
3 changed files with 52 additions and 10 deletions

View file

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

View file

@ -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<string, string>,
promptCache: Map<string, readonly string[]>,
): 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<string, string>();
const promptCache = new Map<string, readonly string[]>();
const deferredResponses = new Map<
string,
{

View file

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