Sessions: persist prompt-token totals without usage

This commit is contained in:
Vignesh Natarajan
2026-02-21 23:37:27 -08:00
parent 76828e8dc8
commit aab20e58d7
4 changed files with 86 additions and 16 deletions

View File

@@ -57,25 +57,25 @@ export async function persistSessionUsageUpdate(params: {
}
const label = params.logLabel ? `${params.logLabel} ` : "";
if (hasNonzeroUsage(params.usage)) {
const hasUsage = hasNonzeroUsage(params.usage);
const hasPromptTokens =
typeof params.promptTokens === "number" &&
Number.isFinite(params.promptTokens) &&
params.promptTokens > 0;
const hasFreshContextSnapshot = Boolean(params.lastCallUsage) || hasPromptTokens;
if (hasUsage || hasFreshContextSnapshot) {
try {
await updateSessionStoreEntry({
storePath,
sessionKey,
update: async (entry) => {
const input = params.usage?.input ?? 0;
const output = params.usage?.output ?? 0;
const resolvedContextTokens = params.contextTokensUsed ?? entry.contextTokens;
const hasPromptTokens =
typeof params.promptTokens === "number" &&
Number.isFinite(params.promptTokens) &&
params.promptTokens > 0;
const hasFreshContextSnapshot = Boolean(params.lastCallUsage) || hasPromptTokens;
// Use last-call usage for totalTokens when available. The accumulated
// `usage.input` sums input tokens from every API call in the run
// (tool-use loops, compaction retries), overstating actual context.
// `lastCallUsage` reflects only the final API call — the true context.
const usageForContext = params.lastCallUsage ?? params.usage;
const usageForContext = params.lastCallUsage ?? (hasUsage ? params.usage : undefined);
const totalTokens = hasFreshContextSnapshot
? deriveSessionTotalTokens({
usage: usageForContext,
@@ -84,19 +84,22 @@ export async function persistSessionUsageUpdate(params: {
})
: undefined;
const patch: Partial<SessionEntry> = {
inputTokens: input,
outputTokens: output,
cacheRead: params.usage?.cacheRead ?? 0,
cacheWrite: params.usage?.cacheWrite ?? 0,
// Missing a last-call snapshot means context utilization is stale/unknown.
totalTokens,
totalTokensFresh: typeof totalTokens === "number",
modelProvider: params.providerUsed ?? entry.modelProvider,
model: params.modelUsed ?? entry.model,
contextTokens: resolvedContextTokens,
systemPromptReport: params.systemPromptReport ?? entry.systemPromptReport,
updatedAt: Date.now(),
};
if (hasUsage) {
patch.inputTokens = params.usage?.input ?? 0;
patch.outputTokens = params.usage?.output ?? 0;
patch.cacheRead = params.usage?.cacheRead ?? 0;
patch.cacheWrite = params.usage?.cacheWrite ?? 0;
}
// Missing a last-call snapshot (and promptTokens fallback) means
// context utilization is stale/unknown.
patch.totalTokens = totalTokens;
patch.totalTokensFresh = typeof totalTokens === "number";
return applyCliSessionIdToSessionPatch(params, entry, patch);
},
});