From 6bb2d9d79bce382d87ced7fea8bb07470ed52fbf Mon Sep 17 00:00:00 2001
From: Stephen Parker <8068616+sp-hk2ldn@users.noreply.github.com>
Date: Wed, 4 Mar 2026 20:50:52 +0000
Subject: [PATCH] fix(agents): harden compaction usage accounting when usage is
 missing

---
 ...ed-runner.sanitize-session-history.test.ts | 49 +++++++++++++++
 src/agents/pi-embedded-runner/google.ts       | 61 ++++++++++++++++++-
 2 files changed, 107 insertions(+), 3 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 13884cd904f..9db185e3eaa 100644
--- a/src/agents/pi-embedded-runner.sanitize-session-history.test.ts
+++ b/src/agents/pi-embedded-runner.sanitize-session-history.test.ts
@@ -330,6 +330,55 @@ describe("sanitizeSessionHistory", () => {
     expect(assistants[1]?.usage).toBeDefined();
   });
 
+  it("adds a zeroed assistant usage snapshot when usage is missing", async () => {
+    vi.mocked(helpers.isGoogleModelApi).mockReturnValue(false);
+
+    const messages = castAgentMessages([
+      { role: "user", content: "question" },
+      {
+        role: "assistant",
+        content: [{ type: "text", text: "answer without usage" }],
+      },
+    ]);
+
+    const result = await sanitizeOpenAIHistory(messages);
+    const assistant = result.find((message) => message.role === "assistant") as
+      | (AgentMessage & { usage?: unknown })
+      | undefined;
+
+    expect(assistant?.usage).toEqual(makeZeroUsageSnapshot());
+  });
+
+  it("normalizes mixed partial assistant usage fields to numeric totals", async () => {
+    vi.mocked(helpers.isGoogleModelApi).mockReturnValue(false);
+
+    const messages = castAgentMessages([
+      { role: "user", content: "question" },
+      {
+        role: "assistant",
+        content: [{ type: "text", text: "answer with partial usage" }],
+        usage: {
+          output: 3,
+          cache_read_input_tokens: 9,
+        },
+      },
+    ]);
+
+    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,
+    });
+  });
+
   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 094aa9142c3..07ffded6d02 100644
--- a/src/agents/pi-embedded-runner/google.ts
+++ b/src/agents/pi-embedded-runner/google.ts
@@ -25,7 +25,7 @@ import {
 } from "../session-transcript-repair.js";
 import type { TranscriptPolicy } from "../transcript-policy.js";
 import { resolveTranscriptPolicy } from "../transcript-policy.js";
-import { makeZeroUsageSnapshot } from "../usage.js";
+import { makeZeroUsageSnapshot, normalizeUsage, type UsageLike } from "../usage.js";
 import { log } from "./logger.js";
 import { dropThinkingBlocks } from "./thinking.js";
 import { describeUnknownError } from "./utils.js";
@@ -200,6 +200,60 @@ function stripStaleAssistantUsageBeforeLatestCompaction(messages: AgentMessage[]
   return touched ? out : messages;
 }
 
+function normalizeAssistantUsageSnapshot(usage: unknown) {
+  const normalized = normalizeUsage((usage ?? undefined) as UsageLike | undefined);
+  if (!normalized) {
+    return makeZeroUsageSnapshot();
+  }
+  const input = normalized.input ?? 0;
+  const output = normalized.output ?? 0;
+  const cacheRead = normalized.cacheRead ?? 0;
+  const cacheWrite = normalized.cacheWrite ?? 0;
+  const totalTokens = normalized.total ?? input + output + cacheRead + cacheWrite;
+  return {
+    ...makeZeroUsageSnapshot(),
+    input,
+    output,
+    cacheRead,
+    cacheWrite,
+    totalTokens,
+  };
+}
+
+function ensureAssistantUsageSnapshots(messages: AgentMessage[]): AgentMessage[] {
+  if (messages.length === 0) {
+    return messages;
+  }
+
+  let touched = false;
+  const out = [...messages];
+  for (let i = 0; i < out.length; i += 1) {
+    const message = out[i] as (AgentMessage & { role?: unknown; usage?: unknown }) | undefined;
+    if (!message || message.role !== "assistant") {
+      continue;
+    }
+    const normalizedUsage = normalizeAssistantUsageSnapshot(message.usage);
+    if (
+      message.usage &&
+      typeof message.usage === "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
+    ) {
+      continue;
+    }
+    out[i] = {
+      ...(message as unknown as Record<string, unknown>),
+      usage: normalizedUsage,
+    } as AgentMessage;
+    touched = true;
+  }
+
+  return touched ? out : messages;
+}
+
 export function findUnsupportedSchemaKeywords(schema: unknown, path: string): string[] {
   if (!schema || typeof schema !== "object") {
     return [];
@@ -449,8 +503,9 @@ export async function sanitizeSessionHistory(params: {
     ? sanitizeToolUseResultPairing(sanitizedToolCalls)
     : sanitizedToolCalls;
   const sanitizedToolResults = stripToolResultDetails(repairedTools);
-  const sanitizedCompactionUsage =
-    stripStaleAssistantUsageBeforeLatestCompaction(sanitizedToolResults);
+  const sanitizedCompactionUsage = ensureAssistantUsageSnapshots(
+    stripStaleAssistantUsageBeforeLatestCompaction(sanitizedToolResults),
+  );
 
   const isOpenAIResponsesApi =
     params.modelApi === "openai-responses" || params.modelApi === "openai-codex-responses";