diff --git a/src/agents/pi-embedded-runner.sanitize-session-history.test.ts b/src/agents/pi-embedded-runner.sanitize-session-history.test.ts index 9db185e3eaa..c0c184a43a7 100644 --- a/src/agents/pi-embedded-runner.sanitize-session-history.test.ts +++ b/src/agents/pi-embedded-runner.sanitize-session-history.test.ts @@ -379,6 +379,83 @@ describe("sanitizeSessionHistory", () => { }); }); + it("preserves existing usage cost while normalizing token fields", async () => { + vi.mocked(helpers.isGoogleModelApi).mockReturnValue(false); + + const messages = castAgentMessages([ + { role: "user", content: "question" }, + { + role: "assistant", + content: [{ type: "text", text: "answer with partial usage and cost" }], + usage: { + output: 3, + cache_read_input_tokens: 9, + cost: { + input: 1.25, + output: 2.5, + cacheRead: 0.25, + cacheWrite: 0, + total: 4, + }, + }, + }, + ]); + + const result = await sanitizeOpenAIHistory(messages); + const assistant = result.find((message) => message.role === "assistant") as + | (AgentMessage & { usage?: unknown }) + | undefined; + + expect(assistant?.usage).toEqual({ + ...makeZeroUsageSnapshot(), + input: 0, + output: 3, + cacheRead: 9, + cacheWrite: 0, + totalTokens: 12, + cost: { + input: 1.25, + output: 2.5, + cacheRead: 0.25, + cacheWrite: 0, + total: 4, + }, + }); + }); + + it("adds missing usage cost even when token fields already match", async () => { + vi.mocked(helpers.isGoogleModelApi).mockReturnValue(false); + + const messages = castAgentMessages([ + { role: "user", content: "question" }, + { + role: "assistant", + content: [{ type: "text", text: "answer with complete numeric usage but no cost" }], + usage: { + input: 1, + output: 2, + cacheRead: 3, + cacheWrite: 4, + totalTokens: 10, + }, + }, + ]); + + const result = await sanitizeOpenAIHistory(messages); + const assistant = result.find((message) => message.role === "assistant") as + | (AgentMessage & { usage?: unknown }) + | undefined; + + expect(assistant?.usage).toEqual({ + ...makeZeroUsageSnapshot(), + input: 1, + output: 2, + cacheRead: 3, + cacheWrite: 4, + totalTokens: 10, + }); + }); + it("drops stale usage when compaction summary appears before kept assistant messages", async () => { vi.mocked(helpers.isGoogleModelApi).mockReturnValue(false); diff --git a/src/agents/pi-embedded-runner/google.ts b/src/agents/pi-embedded-runner/google.ts index 07ffded6d02..d3731795269 100644 --- a/src/agents/pi-embedded-runner/google.ts +++ b/src/agents/pi-embedded-runner/google.ts @@ -25,7 +25,12 @@ import { } from "../session-transcript-repair.js"; import type { TranscriptPolicy } from "../transcript-policy.js"; import { resolveTranscriptPolicy } from "../transcript-policy.js"; -import { makeZeroUsageSnapshot, normalizeUsage, type UsageLike } from "../usage.js"; +import { + makeZeroUsageSnapshot, + normalizeUsage, + type AssistantUsageSnapshot, + type UsageLike, +} from "../usage.js"; import { log } from "./logger.js"; import { dropThinkingBlocks } from "./thinking.js"; import { describeUnknownError } from "./utils.js"; @@ -210,8 +215,10 @@ function normalizeAssistantUsageSnapshot(usage: unknown) { const cacheRead = normalized.cacheRead ?? 0; const cacheWrite = normalized.cacheWrite ?? 0; const totalTokens = normalized.total ?? input + output + cacheRead + cacheWrite; + const cost = normalizeAssistantUsageCost(usage); return { ...makeZeroUsageSnapshot(), + cost, input, output, cacheRead, @@ -220,6 +227,28 @@ function normalizeAssistantUsageSnapshot(usage: unknown) { }; } +function normalizeAssistantUsageCost(usage: unknown): AssistantUsageSnapshot["cost"] { + const base = makeZeroUsageSnapshot().cost; + if (!usage || typeof usage !== "object") { + return base; + } + const rawCost = (usage as { cost?: unknown }).cost; + if (!rawCost || typeof rawCost !== "object") { + return base; + } + const cost = rawCost as Record; + const input = toFiniteCostNumber(cost.input) ?? base.input; + const output = toFiniteCostNumber(cost.output) ?? base.output; + const cacheRead = toFiniteCostNumber(cost.cacheRead) ?? base.cacheRead; + const cacheWrite = toFiniteCostNumber(cost.cacheWrite) ?? base.cacheWrite; + const total = toFiniteCostNumber(cost.total) ?? input + output + cacheRead + cacheWrite; + return { input, output, cacheRead, cacheWrite, total }; +} + +function toFiniteCostNumber(value: unknown): number | undefined { + return typeof value === "number" && Number.isFinite(value) ? value : undefined; +} + function ensureAssistantUsageSnapshots(messages: AgentMessage[]): AgentMessage[] { if (messages.length === 0) { return messages; @@ -233,14 +262,25 @@ function ensureAssistantUsageSnapshots(messages: AgentMessage[]): AgentMessage[] continue; } const normalizedUsage = normalizeAssistantUsageSnapshot(message.usage); + const usageCost = + message.usage && typeof message.usage === "object" + ? (message.usage as { cost?: unknown }).cost + : undefined; if ( message.usage && typeof message.usage === "object" && + usageCost && + typeof usageCost === "object" && (message.usage as { input?: unknown }).input === normalizedUsage.input && (message.usage as { output?: unknown }).output === normalizedUsage.output && (message.usage as { cacheRead?: unknown }).cacheRead === normalizedUsage.cacheRead && (message.usage as { cacheWrite?: unknown }).cacheWrite === normalizedUsage.cacheWrite && - (message.usage as { totalTokens?: unknown }).totalTokens === normalizedUsage.totalTokens + (message.usage as { totalTokens?: unknown }).totalTokens === normalizedUsage.totalTokens && + (usageCost as { input?: unknown }).input === normalizedUsage.cost.input && + (usageCost as { output?: unknown }).output === normalizedUsage.cost.output && + (usageCost as { cacheRead?: unknown }).cacheRead === normalizedUsage.cost.cacheRead && + (usageCost as { cacheWrite?: unknown }).cacheWrite === normalizedUsage.cost.cacheWrite && + (usageCost as { total?: unknown }).total === normalizedUsage.cost.total ) { continue; }