From 7e0b3f16e3bd37efdb45c90afa2647617bf48507 Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Thu, 26 Feb 2026 21:35:13 +0000 Subject: [PATCH] fix: preserve assistant usage snapshots during compaction cleanup --- ...ed-runner.sanitize-session-history.test.ts | 9 +++-- src/agents/pi-embedded-runner/google.ts | 9 ++++- ...-embedded-subscribe.handlers.compaction.ts | 8 ++-- ...g-single-line-fenced-blocks-reopen.test.ts | 37 +++++++++++++++++++ src/agents/usage.ts | 32 ++++++++++++++++ 5 files changed, 85 insertions(+), 10 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 20ea0905d91..fc1a2cec801 100644 --- a/src/agents/pi-embedded-runner.sanitize-session-history.test.ts +++ b/src/agents/pi-embedded-runner.sanitize-session-history.test.ts @@ -14,6 +14,7 @@ import { sanitizeWithOpenAIResponses, TEST_SESSION_ID, } from "./pi-embedded-runner.sanitize-session-history.test-harness.js"; +import { makeZeroUsageSnapshot } from "./usage.js"; vi.mock("./pi-embedded-helpers.js", async () => ({ ...(await vi.importActual("./pi-embedded-helpers.js")), @@ -210,7 +211,7 @@ describe("sanitizeSessionHistory", () => { | (AgentMessage & { usage?: unknown }) | undefined; expect(staleAssistant).toBeDefined(); - expect(staleAssistant?.usage).toBeUndefined(); + expect(staleAssistant?.usage).toEqual(makeZeroUsageSnapshot()); }); it("preserves fresh assistant usage snapshots created after latest compaction summary", async () => { @@ -264,7 +265,7 @@ describe("sanitizeSessionHistory", () => { AgentMessage & { usage?: unknown } >; expect(assistants).toHaveLength(2); - expect(assistants[0]?.usage).toBeUndefined(); + expect(assistants[0]?.usage).toEqual(makeZeroUsageSnapshot()); expect(assistants[1]?.usage).toBeDefined(); }); @@ -306,7 +307,7 @@ describe("sanitizeSessionHistory", () => { const assistant = result.find((message) => message.role === "assistant") as | (AgentMessage & { usage?: unknown }) | undefined; - expect(assistant?.usage).toBeUndefined(); + expect(assistant?.usage).toEqual(makeZeroUsageSnapshot()); }); it("keeps fresh usage after compaction timestamp in summary-first ordering", async () => { @@ -368,7 +369,7 @@ describe("sanitizeSessionHistory", () => { const freshAssistant = assistants.find((message) => JSON.stringify(message.content).includes("fresh answer"), ); - expect(keptAssistant?.usage).toBeUndefined(); + expect(keptAssistant?.usage).toEqual(makeZeroUsageSnapshot()); expect(freshAssistant?.usage).toBeDefined(); }); diff --git a/src/agents/pi-embedded-runner/google.ts b/src/agents/pi-embedded-runner/google.ts index 5e8f546cd09..429c1ddd9d9 100644 --- a/src/agents/pi-embedded-runner/google.ts +++ b/src/agents/pi-embedded-runner/google.ts @@ -24,6 +24,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 { log } from "./logger.js"; import { dropThinkingBlocks } from "./thinking.js"; import { describeUnknownError } from "./utils.js"; @@ -186,9 +187,13 @@ function stripStaleAssistantUsageBeforeLatestCompaction(messages: AgentMessage[] continue; } + // pi-coding-agent expects assistant usage to always be present during context + // accounting. Keep stale snapshots structurally valid, but zeroed out. const candidateRecord = candidate as unknown as Record; - const { usage: _droppedUsage, ...rest } = candidateRecord; - out[i] = rest as unknown as AgentMessage; + out[i] = { + ...candidateRecord, + usage: makeZeroUsageSnapshot(), + } as unknown as AgentMessage; touched = true; } return touched ? out : messages; diff --git a/src/agents/pi-embedded-subscribe.handlers.compaction.ts b/src/agents/pi-embedded-subscribe.handlers.compaction.ts index 8ae5d1ef465..f25d05f0065 100644 --- a/src/agents/pi-embedded-subscribe.handlers.compaction.ts +++ b/src/agents/pi-embedded-subscribe.handlers.compaction.ts @@ -2,6 +2,7 @@ import type { AgentEvent } from "@mariozechner/pi-agent-core"; import { emitAgentEvent } from "../infra/agent-events.js"; import { getGlobalHookRunner } from "../plugins/hook-runner-global.js"; import type { EmbeddedPiSubscribeContext } from "./pi-embedded-subscribe.handlers.types.js"; +import { makeZeroUsageSnapshot } from "./usage.js"; export function handleAutoCompactionStart(ctx: EmbeddedPiSubscribeContext) { ctx.state.compactionInFlight = true; @@ -96,9 +97,8 @@ function clearStaleAssistantUsageOnSessionMessages(ctx: EmbeddedPiSubscribeConte if (candidate.role !== "assistant") { continue; } - if (!("usage" in candidate)) { - continue; - } - delete (candidate as { usage?: unknown }).usage; + // pi-coding-agent expects assistant usage to exist when computing context usage. + // Reset stale snapshots to zeros instead of deleting the field. + candidate.usage = makeZeroUsageSnapshot(); } } diff --git a/src/agents/pi-embedded-subscribe.subscribe-embedded-pi-session.splits-long-single-line-fenced-blocks-reopen.test.ts b/src/agents/pi-embedded-subscribe.subscribe-embedded-pi-session.splits-long-single-line-fenced-blocks-reopen.test.ts index bbc2a019286..bff7046cc80 100644 --- a/src/agents/pi-embedded-subscribe.subscribe-embedded-pi-session.splits-long-single-line-fenced-blocks-reopen.test.ts +++ b/src/agents/pi-embedded-subscribe.subscribe-embedded-pi-session.splits-long-single-line-fenced-blocks-reopen.test.ts @@ -6,6 +6,7 @@ import { expectFencedChunks, } from "./pi-embedded-subscribe.e2e-harness.js"; import { subscribeEmbeddedPiSession } from "./pi-embedded-subscribe.js"; +import { makeZeroUsageSnapshot } from "./usage.js"; type SessionEventHandler = (evt: unknown) => void; @@ -115,4 +116,40 @@ describe("subscribeEmbeddedPiSession", () => { expect(resolved).toBe(true); expect(subscription.isCompacting()).toBe(false); }); + + it("resets assistant usage to a zero snapshot after compaction without retry", () => { + const listeners: SessionEventHandler[] = []; + const session = { + messages: [ + { + role: "assistant", + content: [{ type: "text", text: "old" }], + usage: { + input: 120, + output: 30, + cacheRead: 5, + cacheWrite: 0, + totalTokens: 155, + cost: { input: 0.001, output: 0.002, cacheRead: 0, cacheWrite: 0, total: 0.003 }, + }, + }, + ], + subscribe: (listener: SessionEventHandler) => { + listeners.push(listener); + return () => {}; + }, + } as unknown as Parameters[0]["session"]; + + subscribeEmbeddedPiSession({ + session, + runId: "run-3", + }); + + for (const listener of listeners) { + listener({ type: "auto_compaction_end", willRetry: false }); + } + + const usage = (session.messages?.[0] as { usage?: unknown } | undefined)?.usage; + expect(usage).toEqual(makeZeroUsageSnapshot()); + }); }); diff --git a/src/agents/usage.ts b/src/agents/usage.ts index be23df97116..703df4ad7e7 100644 --- a/src/agents/usage.ts +++ b/src/agents/usage.ts @@ -34,6 +34,38 @@ export type NormalizedUsage = { total?: number; }; +export type AssistantUsageSnapshot = { + input: number; + output: number; + cacheRead: number; + cacheWrite: number; + totalTokens: number; + cost: { + input: number; + output: number; + cacheRead: number; + cacheWrite: number; + total: number; + }; +}; + +export function makeZeroUsageSnapshot(): AssistantUsageSnapshot { + return { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + total: 0, + }, + }; +} + const asFiniteNumber = (value: unknown): number | undefined => { if (typeof value !== "number") { return undefined;