Skip to content

feat(caching): capture and price cached tokens across providers #1114

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
murdore merged 1 commit into from
Jul 8, 2026
Merged

feat(caching): capture and price cached tokens across providers #1114

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
37 changes: 33 additions & 4 deletions src/lib/providers/googleAiStudio.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -825,6 +825,7 @@ export class GoogleAIStudioProvider extends BaseProvider {
let lastStepText = "";
let totalInputTokens = 0;
let totalOutputTokens = 0;
let totalCacheReadTokens = 0;
let step = 0;
let completedWithFinalAnswer = false;
const failedTools = new Map<
Expand Down Expand Up @@ -866,6 +867,7 @@ export class GoogleAIStudioProvider extends BaseProvider {
);
totalInputTokens += chunkResult.inputTokens;
totalOutputTokens += chunkResult.outputTokens;
totalCacheReadTokens += chunkResult.cacheReadTokens ?? 0;

const stepText = extractTextFromParts(
chunkResult.rawResponseParts,
Expand Down Expand Up @@ -1002,13 +1004,27 @@ export class GoogleAIStudioProvider extends BaseProvider {
hitStepLimitWithoutFinalAnswer ? "max_steps" : "stop",
);

// Gemini promptTokenCount is OVERLAPPING: it already includes
// cachedContentTokenCount. Subtract once here so calculateCost
// bills the cached portion at the cheaper cacheRead rate without
// double-counting. Total billable tokens are conserved.
const adjustedInputTokens = Math.max(
0,
totalInputTokens - totalCacheReadTokens,
);
analyticsResolve({
provider: this.providerName,
model: modelName,
tokenUsage: {
input: totalInputTokens,
input: adjustedInputTokens,
output: totalOutputTokens,
total: totalInputTokens + totalOutputTokens,
total:
adjustedInputTokens +
totalCacheReadTokens +
totalOutputTokens,
...(totalCacheReadTokens > 0
? { cacheReadTokens: totalCacheReadTokens }
: {}),
},
requestDuration: responseTime,
timestamp: new Date().toISOString(),
Expand Down Expand Up @@ -1189,6 +1205,7 @@ export class GoogleAIStudioProvider extends BaseProvider {
let lastStepText = "";
let totalInputTokens = 0;
let totalOutputTokens = 0;
let totalCacheReadTokens = 0;
const allToolCalls: Array<{
toolName: string;
args: Record<string, unknown>;
Expand Down Expand Up @@ -1229,6 +1246,7 @@ export class GoogleAIStudioProvider extends BaseProvider {
const chunkResult = await collectStreamChunks(stream);
totalInputTokens += chunkResult.inputTokens;
totalOutputTokens += chunkResult.outputTokens;
totalCacheReadTokens += chunkResult.cacheReadTokens ?? 0;

const stepText = extractTextFromParts(
chunkResult.rawResponseParts,
Expand Down Expand Up @@ -1351,14 +1369,25 @@ export class GoogleAIStudioProvider extends BaseProvider {
// analytics / evaluation / tracing stay attached. The native AI
// Studio generate path bypasses BaseProvider.generate(), so
// skipping enhanceResult would silently drop those features.
// Gemini promptTokenCount is OVERLAPPING (already includes
// cachedContentTokenCount). Subtract once so the cached portion is
// billed at the cheaper cacheRead rate without double-counting.
const adjustedInputTokens = Math.max(
0,
totalInputTokens - totalCacheReadTokens,
);
const baseResult: EnhancedGenerateResult = {
content: finalText,
provider: this.providerName,
model: modelName,
usage: {
input: totalInputTokens,
input: adjustedInputTokens,
output: totalOutputTokens,
total: totalInputTokens + totalOutputTokens,
total:
adjustedInputTokens + totalCacheReadTokens + totalOutputTokens,
...(totalCacheReadTokens > 0
? { cacheReadTokens: totalCacheReadTokens }
: {}),
},
responseTime,
toolsUsed: allToolCalls.map((tc) => tc.toolName),
Expand Down
42 changes: 38 additions & 4 deletions src/lib/providers/googleNativeGemini3.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -695,6 +695,7 @@ export async function collectStreamChunks(
const stepFunctionCalls: NativeFunctionCall[] = [];
let inputTokens = 0;
let outputTokens = 0;
let cacheReadTokens = 0;

for await (const chunk of stream) {
// Extract raw parts from candidates FIRST
Expand All @@ -717,15 +718,31 @@ export async function collectStreamChunks(

// Accumulate usage metadata from chunks
const usage = chunkRecord.usageMetadata as
| { promptTokenCount?: number; candidatesTokenCount?: number }
| {
promptTokenCount?: number;
candidatesTokenCount?: number;
cachedContentTokenCount?: number;
}
| undefined;
if (usage) {
inputTokens = Math.max(inputTokens, usage.promptTokenCount || 0);
outputTokens = Math.max(outputTokens, usage.candidatesTokenCount || 0);
// cachedContentTokenCount is OVERLAPPING (a subset already inside
// promptTokenCount). Surface it so the call site subtracts once.
cacheReadTokens = Math.max(
cacheReadTokens,
usage.cachedContentTokenCount || 0,
);
}
}

return { rawResponseParts, stepFunctionCalls, inputTokens, outputTokens };
return {
rawResponseParts,
stepFunctionCalls,
inputTokens,
outputTokens,
cacheReadTokens,
};
}

/**
Expand Down Expand Up @@ -829,6 +846,7 @@ export async function collectStreamChunksIncremental(
const stepFunctionCalls: NativeFunctionCall[] = [];
let inputTokens = 0;
let outputTokens = 0;
let cacheReadTokens = 0;

for await (const chunk of stream) {
const chunkRecord = chunk as Record<string, unknown>;
Expand All @@ -853,15 +871,31 @@ export async function collectStreamChunksIncremental(
}

const usage = chunkRecord.usageMetadata as
| { promptTokenCount?: number; candidatesTokenCount?: number }
| {
promptTokenCount?: number;
candidatesTokenCount?: number;
cachedContentTokenCount?: number;
}
| undefined;
if (usage) {
inputTokens = Math.max(inputTokens, usage.promptTokenCount || 0);
outputTokens = Math.max(outputTokens, usage.candidatesTokenCount || 0);
// cachedContentTokenCount is OVERLAPPING (a subset already inside
// promptTokenCount). Surface it so the call site subtracts once.
cacheReadTokens = Math.max(
cacheReadTokens,
usage.cachedContentTokenCount || 0,
);
}
}

return { rawResponseParts, stepFunctionCalls, inputTokens, outputTokens };
return {
rawResponseParts,
stepFunctionCalls,
inputTokens,
outputTokens,
cacheReadTokens,
};
}

/**
Expand Down
46 changes: 42 additions & 4 deletions src/lib/providers/googleVertex.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1806,6 +1806,7 @@ export class GoogleVertexProvider extends BaseProvider {
// promptTokenCount is typically in the final chunk, candidatesTokenCount accumulates
let totalInputTokens = 0;
let totalOutputTokens = 0;
let totalCacheReadTokens = 0;

// Track text parts as they arrive from the SDK so the returned async
// iterable yields multiple chunks instead of a single buffered chunk.
Expand Down Expand Up @@ -1941,6 +1942,7 @@ export class GoogleVertexProvider extends BaseProvider {
| {
promptTokenCount?: number;
candidatesTokenCount?: number;
cachedContentTokenCount?: number;
totalTokenCount?: number;
}
| undefined;
Expand All @@ -1956,6 +1958,13 @@ export class GoogleVertexProvider extends BaseProvider {
usageMetadata.promptTokenCount,
usageMetadata.candidatesTokenCount ?? 0,
);
// cachedContentTokenCount is OVERLAPPING (a subset already inside
// promptTokenCount). Clamp to the prompt count so a later uncached
// step resets it to 0 instead of leaving a stale cached value.
totalCacheReadTokens = Math.min(
usageMetadata.cachedContentTokenCount ?? 0,
usageMetadata.promptTokenCount,
);
}
// Take the latest candidatesTokenCount (accumulates through chunks)
if (
Expand Down Expand Up @@ -2503,6 +2512,13 @@ export class GoogleVertexProvider extends BaseProvider {
(te) => te.name !== "final_result",
);

// Gemini promptTokenCount is OVERLAPPING (already includes
// cachedContentTokenCount). Subtract once so the cached portion is billed at
// the cheaper cacheRead rate without double-counting; total is conserved.
const adjustedInputTokens = Math.max(
0,
totalInputTokens - totalCacheReadTokens,
);
const result: StreamResult = {
stream: createTextStream(),
provider: this.providerName,
Expand All @@ -2511,9 +2527,12 @@ export class GoogleVertexProvider extends BaseProvider {
stopReason,
rawFinishReason: lastFinishReason,
usage: {
input: totalInputTokens,
input: adjustedInputTokens,
output: totalOutputTokens,
total: totalInputTokens + totalOutputTokens,
total: adjustedInputTokens + totalCacheReadTokens + totalOutputTokens,
...(totalCacheReadTokens > 0 && {
cacheReadTokens: totalCacheReadTokens,
}),
},
toolCalls: externalToolCalls.map((tc) => ({
toolName: tc.toolName,
Expand Down Expand Up @@ -2977,6 +2996,7 @@ export class GoogleVertexProvider extends BaseProvider {
// promptTokenCount is typically in the final chunk, candidatesTokenCount accumulates
let totalInputTokens = 0;
let totalOutputTokens = 0;
let totalCacheReadTokens = 0;

// Abort scaffolding (mirrors executeNativeAnthropicStream). The native
// Gemini SDK cancels via config.abortSignal, so drive an internal
Expand Down Expand Up @@ -3101,6 +3121,7 @@ export class GoogleVertexProvider extends BaseProvider {
| {
promptTokenCount?: number;
candidatesTokenCount?: number;
cachedContentTokenCount?: number;
totalTokenCount?: number;
}
| undefined;
Expand All @@ -3116,6 +3137,13 @@ export class GoogleVertexProvider extends BaseProvider {
usageMetadata.promptTokenCount,
usageMetadata.candidatesTokenCount ?? 0,
);
// cachedContentTokenCount is OVERLAPPING (a subset already inside
// promptTokenCount). Clamp to the prompt count so a later uncached
// step resets it to 0 instead of leaving a stale cached value.
totalCacheReadTokens = Math.min(
usageMetadata.cachedContentTokenCount ?? 0,
usageMetadata.promptTokenCount,
);
}
// Take the latest candidatesTokenCount (accumulates through chunks)
if (
Expand Down Expand Up @@ -3643,6 +3671,13 @@ export class GoogleVertexProvider extends BaseProvider {
);

// Build EnhancedGenerateResult
// Gemini promptTokenCount is OVERLAPPING (already includes
// cachedContentTokenCount). Subtract once so the cached portion is billed at
// the cheaper cacheRead rate without double-counting; total is conserved.
const adjustedInputTokens = Math.max(
0,
totalInputTokens - totalCacheReadTokens,
);
const result: EnhancedGenerateResult = {
content: finalText,
provider: this.providerName,
Expand All @@ -3652,9 +3687,12 @@ export class GoogleVertexProvider extends BaseProvider {
rawFinishReason: lastFinishReason,
stepsUsed: step,
usage: {
input: totalInputTokens,
input: adjustedInputTokens,
output: totalOutputTokens,
total: totalInputTokens + totalOutputTokens,
total: adjustedInputTokens + totalCacheReadTokens + totalOutputTokens,
...(totalCacheReadTokens > 0 && {
cacheReadTokens: totalCacheReadTokens,
}),
},
responseTime,
toolsUsed: externalToolCalls.map((tc) => tc.toolName),
Expand Down
4 changes: 4 additions & 0 deletions src/lib/types/common.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -357,6 +357,10 @@ export type RawUsageObject = {
cacheCreationTokens?: number;
cacheReadTokens?: number;

// OpenAI/DeepSeek/NIM/OpenAI-compatible nested cache field (overlapping:
// cached_tokens is a SUBSET already included in prompt_tokens)
prompt_tokens_details?: { cached_tokens?: number };

// OpenAI o1/Anthropic reasoning tokens
reasoningTokens?: number;
reasoning?: number;
Expand Down
8 changes: 8 additions & 0 deletions src/lib/types/providers.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1861,6 +1861,14 @@ export type CollectedChunkResult = {
stepFunctionCalls: NativeFunctionCall[];
inputTokens: number;
outputTokens: number;
/**
* Gemini cached-content tokens (overlapping: included in promptTokenCount).
* Surfaced so the call site can subtract from input and bill at cacheRead
* rate. Subtraction happens at the call site, not in the collector.
*/
cacheReadTokens?: number;
/** Cache creation tokens (symmetry; Gemini does not emit this). */
cacheCreationTokens?: number;
};

/** Push-based text channel for incremental streaming. */
Expand Down
Loading
Loading