diff --git a/packages/ai/src/providers/google-vertex.ts b/packages/ai/src/providers/google-vertex.ts index 7dad38f27..6271fc331 100644 --- a/packages/ai/src/providers/google-vertex.ts +++ b/packages/ai/src/providers/google-vertex.ts @@ -234,11 +234,17 @@ export const streamGoogleVertex: StreamFunction<"google-vertex"> = ( } if (chunk.usageMetadata) { + // promptTokenCount includes cachedContentTokenCount when cached content is used. + // Subtract to get non-cached input, matching the OpenAI convention where + // input = uncached prompt tokens and cacheRead = cached tokens so that + // input + cacheRead = total prompt tokens (no double-counting). + // Ref: https://ai.google.dev/api/generate-content#v1beta.GenerateContentResponse.UsageMetadata + const cachedTokens = chunk.usageMetadata.cachedContentTokenCount || 0; output.usage = { - input: chunk.usageMetadata.promptTokenCount || 0, + input: (chunk.usageMetadata.promptTokenCount || 0) - cachedTokens, output: (chunk.usageMetadata.candidatesTokenCount || 0) + (chunk.usageMetadata.thoughtsTokenCount || 0), - cacheRead: chunk.usageMetadata.cachedContentTokenCount || 0, + cacheRead: cachedTokens, cacheWrite: 0, totalTokens: chunk.usageMetadata.totalTokenCount || 0, cost: { diff --git a/packages/ai/src/providers/google.ts b/packages/ai/src/providers/google.ts index 90395ea47..cb25c71d4 100644 --- a/packages/ai/src/providers/google.ts +++ b/packages/ai/src/providers/google.ts @@ -217,11 +217,17 @@ export const streamGoogle: StreamFunction<"google-generative-ai"> = ( } if (chunk.usageMetadata) { + // promptTokenCount includes cachedContentTokenCount when cached content is used. + // Subtract to get non-cached input, matching the OpenAI convention where + // input = uncached prompt tokens and cacheRead = cached tokens so that + // input + cacheRead = total prompt tokens (no double-counting). + // Ref: https://ai.google.dev/api/generate-content#v1beta.GenerateContentResponse.UsageMetadata + const cachedTokens = chunk.usageMetadata.cachedContentTokenCount || 0; output.usage = { - input: chunk.usageMetadata.promptTokenCount || 0, + input: (chunk.usageMetadata.promptTokenCount || 0) - cachedTokens, output: (chunk.usageMetadata.candidatesTokenCount || 0) + (chunk.usageMetadata.thoughtsTokenCount || 0), - cacheRead: chunk.usageMetadata.cachedContentTokenCount || 0, + cacheRead: cachedTokens, cacheWrite: 0, totalTokens: chunk.usageMetadata.totalTokenCount || 0, cost: { diff --git a/packages/ai/src/utils/overflow.ts b/packages/ai/src/utils/overflow.ts index d986a7691..38af73cbe 100644 --- a/packages/ai/src/utils/overflow.ts +++ b/packages/ai/src/utils/overflow.ts @@ -115,7 +115,7 @@ export function isContextOverflow(message: AssistantMessage, contextWindow?: num // Case 2: Usage-based overflow (silent or provider-specific) if (contextWindow) { - const inputTokens = message.usage.input + message.usage.cacheRead; + const inputTokens = message.usage.input + message.usage.cacheRead + message.usage.cacheWrite; if (inputTokens > contextWindow) { return true; } diff --git a/packages/coding-agent/src/modes/components/status-line.ts b/packages/coding-agent/src/modes/components/status-line.ts index 747a3ee53..68dfcc866 100644 --- a/packages/coding-agent/src/modes/components/status-line.ts +++ b/packages/coding-agent/src/modes/components/status-line.ts @@ -7,6 +7,7 @@ import { settings } from "../../config/settings"; import type { StatusLinePreset, StatusLineSegmentId, StatusLineSeparatorStyle } from "../../config/settings-schema"; import { theme } from "../../modes/theme/theme"; import type { AgentSession } from "../../session/agent-session"; +import { calculatePromptTokens } from "../../session/compaction/compaction"; import { findGitHeadPathSync, sanitizeStatusText } from "../shared"; import { canReuseCachedPr, @@ -365,12 +366,7 @@ export class StatusLineComponent implements Component { .reverse() .find(m => m.role === "assistant" && m.stopReason !== "aborted") as AssistantMessage | undefined; - const contextTokens = lastAssistantMessage - ? lastAssistantMessage.usage.input + - lastAssistantMessage.usage.output + - lastAssistantMessage.usage.cacheRead + - lastAssistantMessage.usage.cacheWrite - : 0; + const contextTokens = lastAssistantMessage ? calculatePromptTokens(lastAssistantMessage.usage) : 0; const contextWindow = state.model?.contextWindow || 0; const contextPercent = contextWindow > 0 ? (contextTokens / contextWindow) * 100 : 0; diff --git a/packages/coding-agent/src/session/agent-session.ts b/packages/coding-agent/src/session/agent-session.ts index 62da992b5..678d87b89 100644 --- a/packages/coding-agent/src/session/agent-session.ts +++ b/packages/coding-agent/src/session/agent-session.ts @@ -108,6 +108,7 @@ import { extractFileMentions, generateFileMentionMessages } from "../utils/file- import { type CompactionResult, calculateContextTokens, + calculatePromptTokens, collectEntriesForBranchSummary, compact, estimateTokens, @@ -5076,7 +5077,7 @@ Be thorough - include exact file paths, function names, error messages, and tech }; } - const usageTokens = calculateContextTokens(lastUsage); + const usageTokens = calculatePromptTokens(lastUsage); let trailingTokens = 0; for (let i = lastUsageIndex + 1; i < messages.length; i++) { trailingTokens += estimateTokens(messages[i]);