From 70c7d1ba3f67088ad9420da3e5599bbd799cc291 Mon Sep 17 00:00:00 2001 From: Maurice Date: Wed, 21 Oct 2026 21:26:21 +0000 Subject: [PATCH] feat(assistant): read the conversation from Anthropic's prompt cache The Anthropic request marked only the instructions and the trip card for the prompt cache, so every step of a turn paid the whole history again. It now also marks the end of the history, which the next step starts from, and the end of what came before the user's question, which every step of the turn sends unchanged, within the four marks a request takes. Every provider now reports how much of the input it read from its cache (Anthropic cache_read_input_tokens, OpenAI cached_tokens), the turn adds it up, and the usage log line says it, so the saving can be measured. --- .../ai-chat-client.interface.ts | 3 +- .../nest/ai-connections/ai-usage.service.ts | 15 +++++- .../providers/anthropic-messages.client.ts | 43 +++++++++++++-- .../providers/openai-chat.client.ts | 9 +++- .../providers/openai-responses.client.ts | 8 ++- server/src/nest/assistant/assistant-agent.ts | 7 ++- .../src/nest/assistant/assistant.service.ts | 1 + .../ai-connections/provider-adapters.test.ts | 52 +++++++++++++++++-- .../nest/assistant/assistant-agent.test.ts | 10 ++-- 9 files changed, 126 insertions(+), 22 deletions(-) diff --git a/server/src/nest/ai-connections/ai-chat-client.interface.ts b/server/src/nest/ai-connections/ai-chat-client.interface.ts index 94caba6a1..a9c80b91d 100644 --- a/server/src/nest/ai-connections/ai-chat-client.interface.ts +++ b/server/src/nest/ai-connections/ai-chat-client.interface.ts @@ -51,7 +51,8 @@ export type AiFinishReason = 'stop' | 'tool-calls' | 'length'; export type AiChatEvent = | { type: 'text-delta'; delta: string } | { type: 'tool-call'; call: AiToolCall } - | { type: 'usage'; inputTokens: number; outputTokens: number } + /** `cachedInputTokens` is the part of `inputTokens` the provider read from its prompt cache. */ + | { type: 'usage'; inputTokens: number; cachedInputTokens: number; outputTokens: number } | { type: 'provider-state'; state: AiProviderState } | { type: 'finish'; reason: AiFinishReason }; diff --git a/server/src/nest/ai-connections/ai-usage.service.ts b/server/src/nest/ai-connections/ai-usage.service.ts index ca75ad79d..88277ac26 100644 --- a/server/src/nest/ai-connections/ai-usage.service.ts +++ b/server/src/nest/ai-connections/ai-usage.service.ts @@ -1,5 +1,6 @@ import { AssistantUsage } from '../../db/entities/AssistantUsage.entity'; import type { AssistantUsageRepository } from '../../db/repositories/AssistantUsage.repository'; +import { logDebug } from '../audit/audit-log.logger'; import { InjectRepository } from '@mikro-orm/nestjs'; import { Injectable } from '@nestjs/common'; import { todayUtc, type AiConnectionScope } from '@trek/shared'; @@ -13,12 +14,22 @@ import { todayUtc, type AiConnectionScope } from '@trek/shared'; export class AiUsageService { constructor(@InjectRepository(AssistantUsage) private readonly usage: AssistantUsageRepository) {} - /** Adds a turn's usage to today's row; the turn itself is counted once, when it starts. */ + /** + * Adds a turn's usage to today's row; the turn itself is counted once, when + * it starts. How much of the input came from the provider's prompt cache is + * only logged, to see what the cache saves. + */ async record( userId: number, paidBy: AiConnectionScope, - delta: { turns?: number; inputTokens?: number; outputTokens?: number }, + delta: { turns?: number; inputTokens?: number; cachedInputTokens?: number; outputTokens?: number }, ): Promise { + if (delta.inputTokens) { + logDebug( + `[assistant] turn of user ${userId}: ${delta.inputTokens} input tokens, ` + + `${delta.cachedInputTokens ?? 0} of them from the prompt cache, ${delta.outputTokens ?? 0} output tokens`, + ); + } await this.usage.add(userId, todayUtc(), paidBy, { turns: delta.turns ?? 0, inputTokens: delta.inputTokens ?? 0, diff --git a/server/src/nest/ai-connections/providers/anthropic-messages.client.ts b/server/src/nest/ai-connections/providers/anthropic-messages.client.ts index b8dfb6dd2..d81f7ca11 100644 --- a/server/src/nest/ai-connections/providers/anthropic-messages.client.ts +++ b/server/src/nest/ai-connections/providers/anthropic-messages.client.ts @@ -67,12 +67,38 @@ export function toAnthropicMessages(messages: AiChatMessage[]): AnthropicMessage return out; } +/** The most cache breakpoints Anthropic takes in one request; the system blocks use the first ones. */ +const MAX_CACHE_BREAKPOINTS = 4; +const CACHED = { type: 'ephemeral' } as const; + +/** + * Marks the history for the prompt cache, with what the system blocks left of + * the breakpoints: the end of the whole history, which the next step of the + * turn starts from, and the end of what came before the user's question of + * this turn, which every step of the turn sends unchanged. The second one + * still hits when a step added more blocks (many tool calls and their + * results) than the cache looks back over from the first. The turns are the + * request's own copies, so they are marked in place. + */ +export function markHistoryForCache(messages: AnthropicMessage[], breakpoints: number): void { + const blocks = messages.flatMap((message) => message.content.map((block) => ({ role: message.role, block }))); + const last = blocks.at(-1); + if (!last || breakpoints < 1) return; + last.block.cache_control = CACHED; + let question = blocks.length - 1; + while (question >= 0 && !(blocks[question]?.role === 'user' && blocks[question]?.block.type === 'text')) question--; + const before = question > 0 ? blocks[question - 1] : undefined; + if (breakpoints >= 2 && before && before !== last) before.block.cache_control = CACHED; +} + const STOP_REASONS: Record = { tool_use: 'tool-calls', max_tokens: 'length' }; /** * Anthropic's Messages API, streamed. The instructions and the trip card are * two system blocks, each marked for the prompt cache, so the tools and the - * instructions are read from cache on every later step of a conversation. + * instructions are read from cache on every later step of a conversation; the + * history is marked too (markHistoryForCache), so a step reads the turns + * before it from the cache instead of paying for them again. */ export class AnthropicMessagesClient implements AiChatClient { readonly kind = 'anthropic' as const; @@ -86,13 +112,15 @@ export class AnthropicMessagesClient implements AiChatClient { async *stream(request: AiChatRequest, signal: AbortSignal): AsyncGenerator { const system = [request.system, request.context] .filter((text): text is string => !!text) - .map((text) => ({ type: 'text', text, cache_control: { type: 'ephemeral' } })); + .map((text) => ({ type: 'text', text, cache_control: CACHED })); + const messages = toAnthropicMessages(request.messages); + markHistoryForCache(messages, MAX_CACHE_BREAKPOINTS - system.length); const body = { model: this.model, max_tokens: MAX_TOKENS, stream: true, system, - messages: toAnthropicMessages(request.messages), + messages, ...(request.tools.length > 0 ? { tools: request.tools.map((tool) => ({ @@ -129,7 +157,12 @@ export class AnthropicMessagesClient implements AiChatClient { if (out) yield out; } if (!turn.finished) throw new AiProviderError('provider_unavailable', 0, true, undefined, 'stream_cut'); - yield { type: 'usage', inputTokens: turn.inputTokens, outputTokens: turn.outputTokens }; + yield { + type: 'usage', + inputTokens: turn.inputTokens, + cachedInputTokens: turn.cachedInputTokens, + outputTokens: turn.outputTokens, + }; yield { type: 'finish', reason: turn.reason }; } } @@ -137,6 +170,7 @@ export class AnthropicMessagesClient implements AiChatClient { /** What one streamed Messages API turn has said so far, fed one event at a time. */ class AnthropicTurn { inputTokens = 0; + cachedInputTokens = 0; outputTokens = 0; reason: AiFinishReason = 'stop'; finished = false; @@ -152,6 +186,7 @@ class AnthropicTurn { asCount(usage.input_tokens) + asCount(usage.cache_creation_input_tokens) + asCount(usage.cache_read_input_tokens); + this.cachedInputTokens = asCount(usage.cache_read_input_tokens); this.outputTokens = asCount(usage.output_tokens); return null; } diff --git a/server/src/nest/ai-connections/providers/openai-chat.client.ts b/server/src/nest/ai-connections/providers/openai-chat.client.ts index 227e9bd9f..a0b7526e6 100644 --- a/server/src/nest/ai-connections/providers/openai-chat.client.ts +++ b/server/src/nest/ai-connections/providers/openai-chat.client.ts @@ -115,7 +115,12 @@ export class OpenAiChatClient implements AiChatClient { } if (!turn.done && !turn.reason) throw new AiProviderError('provider_unavailable', 0, true, undefined, 'stream_cut'); for (const call of turn.calls()) yield { type: 'tool-call', call }; - yield { type: 'usage', inputTokens: turn.inputTokens, outputTokens: turn.outputTokens }; + yield { + type: 'usage', + inputTokens: turn.inputTokens, + cachedInputTokens: turn.cachedInputTokens, + outputTokens: turn.outputTokens, + }; yield { type: 'finish', reason: turn.finishReason() }; } } @@ -127,6 +132,7 @@ export class OpenAiChatClient implements AiChatClient { */ class ChatTurn { inputTokens = 0; + cachedInputTokens = 0; outputTokens = 0; reason: string | null = null; done = false; @@ -137,6 +143,7 @@ class ChatTurn { const usage = asRecord(chunk.usage); if (chunk.usage) { this.inputTokens = asCount(usage.prompt_tokens); + this.cachedInputTokens = asCount(asRecord(usage.prompt_tokens_details).cached_tokens); this.outputTokens = asCount(usage.completion_tokens); } const choice = asRecord(Array.isArray(chunk.choices) ? chunk.choices[0] : undefined); diff --git a/server/src/nest/ai-connections/providers/openai-responses.client.ts b/server/src/nest/ai-connections/providers/openai-responses.client.ts index bbbca9672..3790b9db6 100644 --- a/server/src/nest/ai-connections/providers/openai-responses.client.ts +++ b/server/src/nest/ai-connections/providers/openai-responses.client.ts @@ -116,7 +116,7 @@ export class OpenAiResponsesClient implements AiChatClient { const events = await this.open(body, request.cacheKey, signal); const items: Record[] = []; - let usage = { inputTokens: 0, outputTokens: 0 }; + let usage = { inputTokens: 0, cachedInputTokens: 0, outputTokens: 0 }; let completed = false; for await (const sse of readSseEvents(events)) { const event = parseEventData(sse.data); @@ -134,7 +134,11 @@ export class OpenAiResponsesClient implements AiChatClient { } } else if (event.type === 'response.completed') { const reported = asRecord(asRecord(event.response).usage); - usage = { inputTokens: asCount(reported.input_tokens), outputTokens: asCount(reported.output_tokens) }; + usage = { + inputTokens: asCount(reported.input_tokens), + cachedInputTokens: asCount(asRecord(reported.input_tokens_details).cached_tokens), + outputTokens: asCount(reported.output_tokens), + }; completed = true; } else if (event.type === 'error' || event.type === 'response.failed' || event.type === 'response.incomplete') { throw providerFailure(this.kind, 0, errorCodeOf(event) ?? errorCodeOf(event.response) ?? asText(event.type)); diff --git a/server/src/nest/assistant/assistant-agent.ts b/server/src/nest/assistant/assistant-agent.ts index 86d63316c..7eb805fa5 100644 --- a/server/src/nest/assistant/assistant-agent.ts +++ b/server/src/nest/assistant/assistant-agent.ts @@ -82,6 +82,8 @@ interface AgentOutcome { /** What failed (a provider refusal, the step limit), null for a turn that ended well or was stopped. */ code: AssistantErrorCode | null; inputTokens: number; + /** The part of inputTokens the provider read from its prompt cache. */ + cachedInputTokens: number; outputTokens: number; } @@ -179,7 +181,7 @@ interface Turn { w: UiStreamWriter; signal: AbortSignal; history: AiChatMessage[]; - usage: { inputTokens: number; outputTokens: number }; + usage: { inputTokens: number; cachedInputTokens: number; outputTokens: number }; knownPlaceIds: Set; /** Whether the turn used its one push (an empty, cut-off or announcing answer). */ nudged: boolean; @@ -224,6 +226,7 @@ async function streamAnswer(turn: Turn, rowId: number): Promise { expect([0, 1, 2, 3].map((i) => dig(body, 'messages', i, 'role'))).toEqual(['user', 'assistant', 'user', undefined]); }); + it('ANTH-006: marks the end of the history and the end of what came before the question for the cache', async () => { + const http = fetchAnswering(streamResponse(STREAM)); + await run(new AnthropicMessagesClient('m', 'k', http.fetcher).stream(REQUEST, signal())); + const marked = JSON.stringify(http.body()).match(/cache_control/g) ?? []; + expect(marked).toHaveLength(4); + // The question "Thanks" shares its turn with the tool result before it: both carry a mark. + expect(dig(http.body(), 'messages', 2, 'content')).toEqual([ + { type: 'tool_result', tool_use_id: 'call_1', content: '{"id":7}', cache_control: { type: 'ephemeral' } }, + { type: 'text', text: 'Thanks', cache_control: { type: 'ephemeral' } }, + ]); + expect(dig(http.body(), 'messages', 0, 'content', 0, 'cache_control')).toBeUndefined(); + }); + + it('ANTH-007: later steps of a turn keep the mark before the question; never more marks than allowed', () => { + const step = toAnthropicMessages([ + ...REQUEST.messages, + { role: 'assistant', text: '', toolCalls: [{ id: 'call_2', name: 'list_places', args: {} }] }, + { role: 'tool', toolCallId: 'call_2', name: 'list_places', output: '[]' }, + ]); + markHistoryForCache(step, 2); + const marks = step.flatMap((turn, i) => + turn.content.flatMap((block, j) => (block.cache_control ? [`${i}.${j}`] : [])), + ); + expect(marks).toEqual(['2.0', '4.0']); + + const first = toAnthropicMessages([{ role: 'user', text: 'Hi' }]); + markHistoryForCache(first, 2); + expect(first[0]?.content).toEqual([{ type: 'text', text: 'Hi', cache_control: { type: 'ephemeral' } }]); + + const none = toAnthropicMessages(REQUEST.messages); + markHistoryForCache(none, 0); + expect(JSON.stringify(none)).not.toContain('cache_control'); + const one = toAnthropicMessages(REQUEST.messages); + markHistoryForCache(one, 1); + expect(JSON.stringify(one).match(/cache_control/g)).toHaveLength(1); + markHistoryForCache([], 2); + }); + it('ANTH-002: a tool result and the next user text share one user turn', () => { expect(toAnthropicMessages(REQUEST.messages)[2]).toEqual({ role: 'user', @@ -216,7 +255,7 @@ describe('AnthropicMessagesClient', () => { expect(events).toEqual([ { type: 'text-delta', delta: 'On it' }, { type: 'tool-call', call: { id: 'toolu_01', name: 'create_place', args: { name: 'Ginkaku-ji' } } }, - { type: 'usage', inputTokens: 100, outputTokens: 42 }, + { type: 'usage', inputTokens: 100, cachedInputTokens: 90, outputTokens: 42 }, { type: 'finish', reason: 'tool-calls' }, ]); }); @@ -273,7 +312,7 @@ describe('OpenAiChatClient', () => { { choices: [{ delta: { tool_calls: [{ index: 0, function: { arguments: '{"name":' } }] } }] }, { choices: [{ delta: { tool_calls: [{ index: 0, function: { arguments: '"Nijo"}' } }] } }] }, { choices: [{ delta: {}, finish_reason: 'tool_calls' }] }, - { choices: [], usage: { prompt_tokens: 120, completion_tokens: 30 } }, + { choices: [], usage: { prompt_tokens: 120, prompt_tokens_details: { cached_tokens: 64 }, completion_tokens: 30 } }, '[DONE]', ]; @@ -314,7 +353,7 @@ describe('OpenAiChatClient', () => { { type: 'text-delta', delta: 'Sure' }, { type: 'text-delta', delta: ', adding.' }, { type: 'tool-call', call: { id: 'call_a', name: 'create_place', args: { name: 'Nijo' } } }, - { type: 'usage', inputTokens: 120, outputTokens: 30 }, + { type: 'usage', inputTokens: 120, cachedInputTokens: 64, outputTokens: 30 }, { type: 'finish', reason: 'tool-calls' }, ]); }); @@ -392,7 +431,10 @@ describe('OpenAiResponsesClient', () => { { type: 'response.output_text.delta', delta: 'Adding ' }, { type: 'response.output_text.delta', delta: 'it.' }, { type: 'response.output_item.done', item: CALL }, - { type: 'response.completed', response: { usage: { input_tokens: 500, output_tokens: 60 } } }, + { + type: 'response.completed', + response: { usage: { input_tokens: 500, input_tokens_details: { cached_tokens: 400 }, output_tokens: 60 } }, + }, ]; const tokens = (...values: string[]) => vi.fn(async (_rejected?: string) => values.shift() ?? 'none'); @@ -423,7 +465,7 @@ describe('OpenAiResponsesClient', () => { { type: 'text-delta', delta: 'Adding ' }, { type: 'text-delta', delta: 'it.' }, { type: 'tool-call', call: { id: 'call_x', name: 'create_place', args: { name: 'Ryoan-ji' } } }, - { type: 'usage', inputTokens: 500, outputTokens: 60 }, + { type: 'usage', inputTokens: 500, cachedInputTokens: 400, outputTokens: 60 }, { type: 'provider-state', state: { adapter: 'openai-responses', items: [REASONING, CALL] } }, { type: 'finish', reason: 'tool-calls' }, ]); diff --git a/server/tests/unit/nest/assistant/assistant-agent.test.ts b/server/tests/unit/nest/assistant/assistant-agent.test.ts index c44c170d5..0df6980a6 100644 --- a/server/tests/unit/nest/assistant/assistant-agent.test.ts +++ b/server/tests/unit/nest/assistant/assistant-agent.test.ts @@ -158,12 +158,12 @@ describe('runAgentLoop', () => { [ { type: 'text-delta', delta: 'Hello ' }, { type: 'text-delta', delta: 'there' }, - { type: 'usage', inputTokens: 10, outputTokens: 3 }, + { type: 'usage', inputTokens: 10, cachedInputTokens: 0, outputTokens: 3 }, { type: 'finish', reason: 'stop' }, ], ]); const outcome = await runAgentLoop(t.deps, t.input, t.writer, new AbortController().signal); - expect(outcome).toEqual({ end: 'stop', code: null, inputTokens: 10, outputTokens: 3 }); + expect(outcome).toEqual({ end: 'stop', code: null, inputTokens: 10, cachedInputTokens: 0, outputTokens: 3 }); expect(t.memory.saved).toEqual([ { rowId: 100, step: { parts: [{ type: 'text', text: 'Hello there' }], providerState: null }, status: 'complete' }, ]); @@ -184,16 +184,16 @@ describe('runAgentLoop', () => { [ { type: 'tool-call', call: { id: 'toolu_1', name: 'list_places', args: { tripId: 1 } } }, { type: 'provider-state', state: { adapter: 'openai-responses', items: [{ type: 'function_call' }] } }, - { type: 'usage', inputTokens: 5, outputTokens: 1 }, + { type: 'usage', inputTokens: 5, cachedInputTokens: 0, outputTokens: 1 }, { type: 'finish', reason: 'tool-calls' }, ], [ { type: 'text-delta', delta: 'Found it' }, - { type: 'usage', inputTokens: 7, outputTokens: 2 }, + { type: 'usage', inputTokens: 7, cachedInputTokens: 5, outputTokens: 2 }, ], ]); const outcome = await runAgentLoop(t.deps, t.input, t.writer, new AbortController().signal); - expect(outcome).toEqual({ end: 'stop', code: null, inputTokens: 12, outputTokens: 3 }); + expect(outcome).toEqual({ end: 'stop', code: null, inputTokens: 12, cachedInputTokens: 5, outputTokens: 3 }); expect(t.tools.call).toHaveBeenCalledWith('list_places', {}); const second = t.requests[1].messages; expect(second[1]).toEqual({