fix(agents): preserve usage cost in usage snapshot normalization

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);
 

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<string, unknown>;
+  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;
     }