diff --git a/packages/cli/src/ui/hooks/useAgentStreamingState.ts b/packages/cli/src/ui/hooks/useAgentStreamingState.ts index 881f715b2ca..8094a0fa24a 100644 --- a/packages/cli/src/ui/hooks/useAgentStreamingState.ts +++ b/packages/cli/src/ui/hooks/useAgentStreamingState.ts @@ -84,14 +84,13 @@ export function useAgentStreamingState( // Dedicated listener for usage metadata — updates React state directly // so the token count is available immediately (even if no other event - // triggers a re-render). Prefers totalTokenCount (prompt + output) - // because output becomes history for the next round, matching - // geminiChat.ts. + // triggers a re-render). Context usage tracks prompt size; output + // isn't in history yet. const usageHandler = (event: { usage?: { totalTokenCount?: number; promptTokenCount?: number }; }) => { const count = - event?.usage?.totalTokenCount ?? event?.usage?.promptTokenCount; + event?.usage?.promptTokenCount ?? event?.usage?.totalTokenCount; if (typeof count === 'number' && count > 0) { setLastPromptTokenCount(count); } diff --git a/packages/cli/src/ui/hooks/useSessionPicker.test.tsx b/packages/cli/src/ui/hooks/useSessionPicker.test.tsx index be70200393a..4dc00487aaa 100644 --- a/packages/cli/src/ui/hooks/useSessionPicker.test.tsx +++ b/packages/cli/src/ui/hooks/useSessionPicker.test.tsx @@ -275,5 +275,4 @@ describe('useSessionPicker multi-select state', () => { expect(onConfirmMulti).toHaveBeenCalledWith(['s2']); expect(onSelect).not.toHaveBeenCalled(); }); - }); diff --git a/packages/core/src/agents/runtime/agent-core.ts b/packages/core/src/agents/runtime/agent-core.ts index d0c7704ce11..075c507cf3a 100644 --- a/packages/core/src/agents/runtime/agent-core.ts +++ b/packages/core/src/agents/runtime/agent-core.ts @@ -1637,10 +1637,11 @@ Important Rules: const thoughtTok = Number(usage.thoughtsTokenCount || 0); const cachedTok = Number(usage.cachedContentTokenCount || 0); const totalTok = Number(usage.totalTokenCount || 0); - // Prefer totalTokenCount (prompt + output) for context usage — the - // output from this round becomes history for the next, matching - // the approach in geminiChat.ts. - const contextTok = isFinite(totalTok) && totalTok > 0 ? totalTok : inTok; + // Context usage tracks prompt size; output isn't in history yet. + // Guard against malformed provider values (`Infinity`/`NaN`) so the + // downstream compaction math doesn't get poisoned — `Infinity` is + // truthy and would otherwise overwrite a valid prior reading. + const contextTok = inTok || totalTok; if (isFinite(contextTok) && contextTok > 0) { this.lastPromptTokenCount = contextTok; } diff --git a/packages/core/src/core/anthropicContentGenerator/anthropicContentGenerator.test.ts b/packages/core/src/core/anthropicContentGenerator/anthropicContentGenerator.test.ts index 276f1a3a310..084eab4cfb6 100644 --- a/packages/core/src/core/anthropicContentGenerator/anthropicContentGenerator.test.ts +++ b/packages/core/src/core/anthropicContentGenerator/anthropicContentGenerator.test.ts @@ -2269,8 +2269,8 @@ describe('AnthropicContentGenerator', () => { delta: { stop_reason: 'end_turn' }, usage: { output_tokens: 5, - input_tokens: 7, - cache_read_input_tokens: 2, + input_tokens: 2, + cache_read_input_tokens: 7, }, }; yield { type: 'message_stop' }; @@ -2331,11 +2331,88 @@ describe('AnthropicContentGenerator', () => { const last = chunks[chunks.length - 1]!; expect(last.candidates?.[0]?.finishReason).toBe(FinishReason.STOP); expect(last.usageMetadata).toEqual({ - cachedContentTokenCount: 2, - promptTokenCount: 9, // cached(2) + input(7) + cachedContentTokenCount: 7, + promptTokenCount: 9, // input(2) + cached(7) — Anthropic-true (input < cache_read) candidatesTokenCount: 5, totalTokenCount: 14, }); }); + + it('accumulates cache_creation_input_tokens through the streaming pipeline', async () => { + // Real Anthropic mid-conversation: `message_start` reports the warm + // prefix bucket (cache_read), the new cache write bucket + // (cache_creation), and the fresh tail (input). The streaming + // accumulator must hold onto cache_creation alongside the other + // buckets so the final chunk's usageMetadata reflects the full + // prompt size — otherwise the cache_creation portion is silently + // dropped from the displayed total and the Footer under-reports by + // exactly that many tokens. + const { AnthropicContentGenerator } = await importGenerator(); + anthropicState.createImpl.mockResolvedValue( + (async function* () { + yield { + type: 'message_start', + message: { + id: 'msg-1', + model: 'claude-test', + usage: { + input_tokens: 2_500, + cache_read_input_tokens: 32_088, + cache_creation_input_tokens: 8_700, + }, + }, + }; + yield { + type: 'content_block_start', + index: 0, + content_block: { type: 'text' }, + }; + yield { + type: 'content_block_delta', + index: 0, + delta: { type: 'text_delta', text: 'ok' }, + }; + yield { type: 'content_block_stop', index: 0 }; + yield { + type: 'message_delta', + delta: { stop_reason: 'end_turn' }, + usage: { output_tokens: 400 }, + }; + yield { type: 'message_stop' }; + })(), + ); + + const generator = new AnthropicContentGenerator( + { + model: 'claude-test', + apiKey: 'test-key', + timeout: 10_000, + maxRetries: 2, + samplingParams: { max_tokens: 123 }, + schemaCompliance: 'auto', + }, + mockConfig, + ); + + const stream = await generator.generateContentStream({ + model: 'models/ignored', + contents: 'Hello', + } as unknown as GenerateContentParameters); + + const chunks: GenerateContentResponse[] = []; + for await (const chunk of stream) { + chunks.push(chunk); + } + + const last = chunks[chunks.length - 1]!; + expect(last.usageMetadata).toEqual({ + // Sum of all three prompt buckets: 2,500 + 32,088 + 8,700 = 43,288. + // cachedContentTokenCount reports cache_read only. + promptTokenCount: 43_288, + candidatesTokenCount: 400, + totalTokenCount: 43_688, + cachedContentTokenCount: 32_088, + }); + }); }); }); diff --git a/packages/core/src/core/anthropicContentGenerator/anthropicContentGenerator.ts b/packages/core/src/core/anthropicContentGenerator/anthropicContentGenerator.ts index 813f5dd4591..987766b15da 100644 --- a/packages/core/src/core/anthropicContentGenerator/anthropicContentGenerator.ts +++ b/packages/core/src/core/anthropicContentGenerator/anthropicContentGenerator.ts @@ -28,6 +28,7 @@ type RawMessageStreamEvent = Anthropic.RawMessageStreamEvent; import { RequestTokenEstimator } from '../../utils/request-tokenizer/index.js'; import { safeJsonParse } from '../../utils/safeJsonParse.js'; import { AnthropicContentConverter } from './converter.js'; +import { buildAnthropicUsageMetadata } from './usage.js'; import { buildRuntimeFetchOptions, redactProxyError, @@ -770,6 +771,7 @@ export class AnthropicContentGenerator implements ContentGenerator { let messageId: string | undefined; let model = this.contentGeneratorConfig.model; let cachedTokens = 0; + let cacheCreationTokens = 0; let promptTokens = 0; let completionTokens = 0; let finishReason: string | undefined; @@ -784,6 +786,9 @@ export class AnthropicContentGenerator implements ContentGenerator { model = event.message.model ?? model; cachedTokens = event.message.usage?.cache_read_input_tokens ?? cachedTokens; + cacheCreationTokens = + event.message.usage?.cache_creation_input_tokens ?? + cacheCreationTokens; promptTokens = event.message.usage?.input_tokens ?? promptTokens; break; } @@ -910,6 +915,12 @@ export class AnthropicContentGenerator implements ContentGenerator { cachedTokens = cacheRead; } } + if (usageRecord?.['cache_creation_input_tokens'] !== undefined) { + const cacheCreate = usageRecord['cache_creation_input_tokens']; + if (typeof cacheCreate === 'number') { + cacheCreationTokens = cacheCreate; + } + } if (finishReason || event.usage) { const chunk = this.buildGeminiChunk( @@ -917,12 +928,12 @@ export class AnthropicContentGenerator implements ContentGenerator { messageId, model, finishReason, - { - cachedContentTokenCount: cachedTokens, - promptTokenCount: cachedTokens + promptTokens, - candidatesTokenCount: completionTokens, - totalTokenCount: cachedTokens + promptTokens + completionTokens, - }, + buildAnthropicUsageMetadata({ + inputTokens: promptTokens, + cacheReadTokens: cachedTokens, + cacheCreationTokens, + outputTokens: completionTokens, + }), ); collectedResponses.push(chunk); yield chunk; @@ -936,12 +947,12 @@ export class AnthropicContentGenerator implements ContentGenerator { messageId, model, finishReason, - { - cachedContentTokenCount: cachedTokens, - promptTokenCount: cachedTokens + promptTokens, - candidatesTokenCount: completionTokens, - totalTokenCount: cachedTokens + promptTokens + completionTokens, - }, + buildAnthropicUsageMetadata({ + inputTokens: promptTokens, + cacheReadTokens: cachedTokens, + cacheCreationTokens, + outputTokens: completionTokens, + }), ); collectedResponses.push(chunk); yield chunk; diff --git a/packages/core/src/core/anthropicContentGenerator/converter.test.ts b/packages/core/src/core/anthropicContentGenerator/converter.test.ts index 0cf25c225b5..05d0612c754 100644 --- a/packages/core/src/core/anthropicContentGenerator/converter.test.ts +++ b/packages/core/src/core/anthropicContentGenerator/converter.test.ts @@ -1306,6 +1306,7 @@ describe('AnthropicContentConverter', () => { promptTokenCount: 3, candidatesTokenCount: 5, totalTokenCount: 8, + cachedContentTokenCount: 0, }); const parts = response.candidates?.[0]?.content?.parts || []; @@ -1332,6 +1333,35 @@ describe('AnthropicContentConverter', () => { { functionCall: { id: 't1', name: 'tool', args: { x: 1 } } }, ]); }); + + it('forwards cache_read_input_tokens and cache_creation_input_tokens through to usageMetadata', () => { + // A real Anthropic mid-conversation response carries all three prompt + // buckets simultaneously: `input_tokens` (the non-cached tail), + // `cache_read_input_tokens` (the warm prefix served from cache), and + // `cache_creation_input_tokens` (the new region being written). The + // converter must forward both cache fields so the normalizer can sum + // them — dropping either silently undercounts the Footer reading by + // the size of the dropped bucket. + const response = converter.convertAnthropicResponseToGemini({ + id: 'msg-1', + model: 'claude-test', + stop_reason: 'end_turn', + content: [{ type: 'text', text: 'ok' }], + usage: { + input_tokens: 2_500, + cache_read_input_tokens: 32_088, + cache_creation_input_tokens: 8_700, + output_tokens: 400, + }, + } as unknown as Anthropic.Message); + + expect(response.usageMetadata).toEqual({ + promptTokenCount: 43_288, + candidatesTokenCount: 400, + totalTokenCount: 43_688, + cachedContentTokenCount: 32_088, + }); + }); }); describe('mapAnthropicFinishReasonToGemini', () => { diff --git a/packages/core/src/core/anthropicContentGenerator/converter.ts b/packages/core/src/core/anthropicContentGenerator/converter.ts index 81c996b31c6..f51e9b18461 100644 --- a/packages/core/src/core/anthropicContentGenerator/converter.ts +++ b/packages/core/src/core/anthropicContentGenerator/converter.ts @@ -18,6 +18,7 @@ import type { ToolListUnion, } from '@google/genai'; import { FinishReason, GenerateContentResponse } from '@google/genai'; +import { buildAnthropicUsageMetadata } from './usage.js'; import type Anthropic from '@anthropic-ai/sdk'; import { safeJsonParse } from '../../utils/safeJsonParse.js'; import { @@ -325,13 +326,12 @@ export class AnthropicContentConverter { geminiResponse.promptFeedback = { safetyRatings: [] }; if (response.usage) { - const promptTokens = response.usage.input_tokens || 0; - const completionTokens = response.usage.output_tokens || 0; - geminiResponse.usageMetadata = { - promptTokenCount: promptTokens, - candidatesTokenCount: completionTokens, - totalTokenCount: promptTokens + completionTokens, - }; + geminiResponse.usageMetadata = buildAnthropicUsageMetadata({ + inputTokens: response.usage.input_tokens || 0, + cacheReadTokens: response.usage.cache_read_input_tokens || 0, + cacheCreationTokens: response.usage.cache_creation_input_tokens || 0, + outputTokens: response.usage.output_tokens || 0, + }); } return geminiResponse; diff --git a/packages/core/src/core/anthropicContentGenerator/usage.test.ts b/packages/core/src/core/anthropicContentGenerator/usage.test.ts new file mode 100644 index 00000000000..a2831462944 --- /dev/null +++ b/packages/core/src/core/anthropicContentGenerator/usage.test.ts @@ -0,0 +1,131 @@ +/** + * @license + * Copyright 2025 Qwen + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, expect, it } from 'vitest'; +import { buildAnthropicUsageMetadata } from './usage.js'; + +describe('buildAnthropicUsageMetadata', () => { + it('sums all three prompt fields under standard Anthropic semantics', () => { + expect( + buildAnthropicUsageMetadata({ + inputTokens: 5_000, + cacheReadTokens: 25_000, + cacheCreationTokens: 0, + outputTokens: 1_000, + }), + ).toEqual({ + promptTokenCount: 30_000, + candidatesTokenCount: 1_000, + totalTokenCount: 31_000, + cachedContentTokenCount: 25_000, + }); + }); + + it('sums when only cache_creation is set (first cache write)', () => { + expect( + buildAnthropicUsageMetadata({ + inputTokens: 10_000, + cacheReadTokens: 0, + cacheCreationTokens: 20_000, + outputTokens: 500, + }), + ).toEqual({ + promptTokenCount: 30_000, + candidatesTokenCount: 500, + totalTokenCount: 30_500, + cachedContentTokenCount: 0, + }); + }); + + it('uses inputTokens alone when it already covers cache fields (OpenAI semantics on Anthropic protocol)', () => { + expect( + buildAnthropicUsageMetadata({ + inputTokens: 30_000, + cacheReadTokens: 25_000, + cacheCreationTokens: 0, + outputTokens: 800, + }), + ).toEqual({ + promptTokenCount: 30_000, + candidatesTokenCount: 800, + totalTokenCount: 30_800, + cachedContentTokenCount: 25_000, + }); + }); + + it('reports inputTokens directly when no cache fields are present', () => { + expect( + buildAnthropicUsageMetadata({ + inputTokens: 12_345, + cacheReadTokens: 0, + cacheCreationTokens: 0, + outputTokens: 678, + }), + ).toEqual({ + promptTokenCount: 12_345, + candidatesTokenCount: 678, + totalTokenCount: 13_023, + cachedContentTokenCount: 0, + }); + }); + + it('keeps summing when inputTokens grows past cache_creation in a long Anthropic conversation', () => { + // Regression: an earlier guard mis-classified this as OpenAI-style + // (because input >= cache_creation) and dropped the cache_creation + // portion, producing a one-shot Footer "drop" at the crossover point. + expect( + buildAnthropicUsageMetadata({ + inputTokens: 50_000, + cacheReadTokens: 0, + cacheCreationTokens: 32_088, + outputTokens: 200, + }), + ).toEqual({ + promptTokenCount: 82_088, + candidatesTokenCount: 200, + totalTokenCount: 82_288, + cachedContentTokenCount: 0, + }); + }); + + it('sums all three buckets when a warm turn both reads and writes cache (real Anthropic mid-conversation)', () => { + // Mid-conversation turn on real Anthropic: the system+tools prefix is + // served from cache (cache_read) AND a new cache breakpoint extends + // the cached region further into the conversation (cache_creation > 0). + // input_tokens carries only the still-non-cached tail. All three buckets + // are mutually exclusive on real Anthropic so the prompt total is their + // sum, and `cachedContentTokenCount` reports the read portion only. + expect( + buildAnthropicUsageMetadata({ + inputTokens: 2_500, + cacheReadTokens: 32_088, + cacheCreationTokens: 8_700, + outputTokens: 400, + }), + ).toEqual({ + promptTokenCount: 43_288, + candidatesTokenCount: 400, + totalTokenCount: 43_688, + cachedContentTokenCount: 32_088, + }); + }); + + it('handles all-zero usage cleanly', () => { + expect( + buildAnthropicUsageMetadata({ + inputTokens: 0, + cacheReadTokens: 0, + cacheCreationTokens: 0, + outputTokens: 0, + }), + ).toEqual({ + promptTokenCount: 0, + candidatesTokenCount: 0, + totalTokenCount: 0, + cachedContentTokenCount: 0, + }); + }); +}); diff --git a/packages/core/src/core/anthropicContentGenerator/usage.ts b/packages/core/src/core/anthropicContentGenerator/usage.ts new file mode 100644 index 00000000000..c6b8dafe32a --- /dev/null +++ b/packages/core/src/core/anthropicContentGenerator/usage.ts @@ -0,0 +1,59 @@ +/** + * @license + * Copyright 2025 Qwen + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { GenerateContentResponseUsageMetadata } from '@google/genai'; + +export interface AnthropicTokenParts { + inputTokens: number; + cacheReadTokens: number; + cacheCreationTokens: number; + outputTokens: number; +} + +/** + * Normalize Anthropic-side token counts into Gemini's `usageMetadata` shape. + * + * Anthropic reports the prompt across three mutually-exclusive fields: + * `input_tokens`, `cache_read_input_tokens`, `cache_creation_input_tokens`. + * The full prompt is the sum. + * + * `cache_creation_input_tokens` is unique to Anthropic's protocol — OpenAI + * has no equivalent — so its presence is a strong signal the response + * follows real Anthropic semantics. Use that as the primary discriminator: + * + * - cache_creation > 0 → Anthropic semantics, sum all three + * - else if cache_read > 0 and input ≥ cache_read → OpenAI-style on the + * Anthropic protocol (input already covers the cached portion), trust + * input alone + * - else → sum (when no cache activity, sum equals input) + * + * An earlier version of this guard compared `inputTokens` to *both* cache + * fields and fell back to "input alone" whenever input was the larger + * value. That mis-fired on long real Anthropic conversations: once enough + * history accumulates, `inputTokens` naturally grows past + * `cache_creation_input_tokens`, which would silently drop the cache + * portion from the displayed prompt size and produce a one-shot Footer + * "drop" at the crossover point. + */ +export function buildAnthropicUsageMetadata( + parts: AnthropicTokenParts, +): GenerateContentResponseUsageMetadata { + const { inputTokens, cacheReadTokens, cacheCreationTokens, outputTokens } = + parts; + const looksLikeOpenAi = + cacheCreationTokens === 0 && + cacheReadTokens > 0 && + inputTokens >= cacheReadTokens; + const promptTotal = looksLikeOpenAi + ? inputTokens + : inputTokens + cacheReadTokens + cacheCreationTokens; + return { + promptTokenCount: promptTotal, + candidatesTokenCount: outputTokens, + totalTokenCount: promptTotal + outputTokens, + cachedContentTokenCount: cacheReadTokens, + }; +} diff --git a/packages/core/src/core/geminiChat.test.ts b/packages/core/src/core/geminiChat.test.ts index d9bb4eb6370..6a9ff791127 100644 --- a/packages/core/src/core/geminiChat.test.ts +++ b/packages/core/src/core/geminiChat.test.ts @@ -944,9 +944,12 @@ describe('GeminiChat', async () => { 'prompt-id-1', ); - // Verify that token counting is called when usageMetadata is present + // Verify that token counting is called when usageMetadata is present. + // The Footer-driving counter must reflect *prompt* size only — output + // tokens for the in-flight round are not yet in history. The mock + // returns promptTokenCount=42, so that's what should be reported. expect(uiTelemetryService.setLastPromptTokenCount).toHaveBeenCalledWith( - 57, + 42, ); expect(uiTelemetryService.setLastPromptTokenCount).toHaveBeenCalledTimes( 1, diff --git a/packages/core/src/core/geminiChat.ts b/packages/core/src/core/geminiChat.ts index 4f06666c405..03b35fb2604 100644 --- a/packages/core/src/core/geminiChat.ts +++ b/packages/core/src/core/geminiChat.ts @@ -1217,10 +1217,9 @@ export class GeminiChat { // Collect token usage for consolidated recording if (chunk.usageMetadata) { usageMetadata = chunk.usageMetadata; - // Use || instead of ?? so that totalTokenCount=0 falls back to promptTokenCount. - // Some providers omit total_tokens or return 0 in streaming usage chunks. + // Context usage tracks prompt size; output isn't in history yet. const lastPromptTokenCount = - usageMetadata.totalTokenCount || usageMetadata.promptTokenCount; + usageMetadata.promptTokenCount || usageMetadata.totalTokenCount; if (lastPromptTokenCount) { // Always update the per-chat counter so this chat (including // subagents) can make its own compaction decisions. diff --git a/packages/core/src/services/sessionService.test.ts b/packages/core/src/services/sessionService.test.ts index 204806472e1..0fdce0b84ff 100644 --- a/packages/core/src/services/sessionService.test.ts +++ b/packages/core/src/services/sessionService.test.ts @@ -895,6 +895,21 @@ describe('SessionService', () => { ).toBe(450); }); + it('should prefer promptTokenCount over totalTokenCount when both are present', () => { + const assistant: ChatRecord = { + ...baseRecord, + uuid: 'a1', + parentUuid: 'comp', + type: 'assistant', + usageMetadata: { promptTokenCount: 200, totalTokenCount: 450 }, + }; + expect( + getResumePromptTokenCount( + makeConversation([compressionRecord, assistant]), + ), + ).toBe(200); + }); + it('should fall back to compression when latest assistant has zero usage', () => { const assistant: ChatRecord = { ...baseRecord, diff --git a/packages/core/src/services/sessionService.ts b/packages/core/src/services/sessionService.ts index dd0372ff8ca..0940066909e 100644 --- a/packages/core/src/services/sessionService.ts +++ b/packages/core/src/services/sessionService.ts @@ -1280,7 +1280,7 @@ export function replayUiTelemetryFromConversation( /** * Returns the best available prompt token count for resuming telemetry. * Walks backward through messages and returns the first valid value: - * - The latest assistant's non-zero usage (totalTokenCount ?? promptTokenCount). + * - The latest assistant's non-zero usage (promptTokenCount ?? totalTokenCount). * - The most recent chat compression checkpoint's newTokenCount. */ export function getResumePromptTokenCount( @@ -1291,7 +1291,7 @@ export function getResumePromptTokenCount( if (record.type === 'assistant') { const usage = record.usageMetadata; - const candidate = usage?.totalTokenCount ?? usage?.promptTokenCount; + const candidate = usage?.promptTokenCount ?? usage?.totalTokenCount; if (candidate) { return candidate; }