From 73da3ecef42abcb23d1abb4d66ff8c826d273488 Mon Sep 17 00:00:00 2001
From: Stephen Parker <8068616+sp-hk2ldn@users.noreply.github.com>
Date: Thu, 5 Mar 2026 18:38:50 +0000
Subject: [PATCH] fix(agents): preserve unknown usage cost semantics

---
 ...ed-runner.sanitize-session-history.test.ts |  5 +-
 src/agents/pi-embedded-runner/google.ts       | 50 ++++++++++++-------
 2 files changed, 35 insertions(+), 20 deletions(-)

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 c0c184a43a7..e216a45f4f3 100644
--- a/src/agents/pi-embedded-runner.sanitize-session-history.test.ts
+++ b/src/agents/pi-embedded-runner.sanitize-session-history.test.ts
@@ -370,7 +370,6 @@ describe("sanitizeSessionHistory", () => {
       | undefined;
 
     expect(assistant?.usage).toEqual({
-      ...makeZeroUsageSnapshot(),
       input: 0,
       output: 3,
       cacheRead: 9,
@@ -423,7 +422,7 @@ describe("sanitizeSessionHistory", () => {
     });
   });
 
-  it("adds missing usage cost even when token fields already match", async () => {
+  it("preserves unknown cost when token fields already match", async () => {
     vi.mocked(helpers.isGoogleModelApi).mockReturnValue(false);
 
     const messages = castAgentMessages([
@@ -447,13 +446,13 @@ describe("sanitizeSessionHistory", () => {
       | undefined;
 
     expect(assistant?.usage).toEqual({
-      ...makeZeroUsageSnapshot(),
       input: 1,
       output: 2,
       cacheRead: 3,
       cacheWrite: 4,
       totalTokens: 10,
     });
+    expect((assistant?.usage as { cost?: unknown } | undefined)?.cost).toBeUndefined();
   });
 
   it("drops stale usage when compaction summary appears before kept assistant messages", async () => {
diff --git a/src/agents/pi-embedded-runner/google.ts b/src/agents/pi-embedded-runner/google.ts
index d3731795269..4daf30552b3 100644
--- a/src/agents/pi-embedded-runner/google.ts
+++ b/src/agents/pi-embedded-runner/google.ts
@@ -217,31 +217,44 @@ function normalizeAssistantUsageSnapshot(usage: unknown) {
   const totalTokens = normalized.total ?? input + output + cacheRead + cacheWrite;
   const cost = normalizeAssistantUsageCost(usage);
   return {
-    ...makeZeroUsageSnapshot(),
-    cost,
     input,
     output,
     cacheRead,
     cacheWrite,
     totalTokens,
+    ...(cost ? { cost } : {}),
   };
 }
 
-function normalizeAssistantUsageCost(usage: unknown): AssistantUsageSnapshot["cost"] {
+function normalizeAssistantUsageCost(usage: unknown): AssistantUsageSnapshot["cost"] | undefined {
   const base = makeZeroUsageSnapshot().cost;
   if (!usage || typeof usage !== "object") {
-    return base;
+    return undefined;
   }
   const rawCost = (usage as { cost?: unknown }).cost;
   if (!rawCost || typeof rawCost !== "object") {
-    return base;
+    return undefined;
   }
   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;
+  const inputRaw = toFiniteCostNumber(cost.input);
+  const outputRaw = toFiniteCostNumber(cost.output);
+  const cacheReadRaw = toFiniteCostNumber(cost.cacheRead);
+  const cacheWriteRaw = toFiniteCostNumber(cost.cacheWrite);
+  const totalRaw = toFiniteCostNumber(cost.total);
+  if (
+    inputRaw === undefined &&
+    outputRaw === undefined &&
+    cacheReadRaw === undefined &&
+    cacheWriteRaw === undefined &&
+    totalRaw === undefined
+  ) {
+    return undefined;
+  }
+  const input = inputRaw ?? base.input;
+  const output = outputRaw ?? base.output;
+  const cacheRead = cacheReadRaw ?? base.cacheRead;
+  const cacheWrite = cacheWriteRaw ?? base.cacheWrite;
+  const total = totalRaw ?? input + output + cacheRead + cacheWrite;
   return { input, output, cacheRead, cacheWrite, total };
 }
 
@@ -266,21 +279,24 @@ function ensureAssistantUsageSnapshots(messages: AgentMessage[]): AgentMessage[]
       message.usage && typeof message.usage === "object"
         ? (message.usage as { cost?: unknown }).cost
         : undefined;
+    const normalizedCost = normalizedUsage.cost;
     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 &&
-      (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
+      ((normalizedCost &&
+        usageCost &&
+        typeof usageCost === "object" &&
+        (usageCost as { input?: unknown }).input === normalizedCost.input &&
+        (usageCost as { output?: unknown }).output === normalizedCost.output &&
+        (usageCost as { cacheRead?: unknown }).cacheRead === normalizedCost.cacheRead &&
+        (usageCost as { cacheWrite?: unknown }).cacheWrite === normalizedCost.cacheWrite &&
+        (usageCost as { total?: unknown }).total === normalizedCost.total) ||
+        (!normalizedCost && usageCost === undefined))
     ) {
       continue;
     }