fix(compaction): preserve fresh usage after compaction

wyf027 · wyf027 · commit 06c47e7f9683 · 2026-05-17T11:34:29.000+08:00
diff --git a/src/agents/pi-embedded-subscribe.handlers.compaction.test.ts b/src/agents/pi-embedded-subscribe.handlers.compaction.test.ts
@@ -14,19 +14,32 @@ import {
   reconcileSessionStoreCompactionCountAfterSuccess,
 } from "./pi-embedded-subscribe.handlers.compaction.js";
 import type { EmbeddedPiSubscribeContext } from "./pi-embedded-subscribe.handlers.types.js";
+import { makeZeroUsageSnapshot } from "./usage.js";
+
+function makeUsage(input: number, output: number, totalTokens: number) {
+  return {
+    input,
+    output,
+    cacheRead: 0,
+    cacheWrite: 0,
+    totalTokens,
+    cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
+  };
+}
 
 function createCompactionContext(params: {
   storePath: string;
   sessionKey: string;
   agentId?: string;
   initialCount: number;
   info?: (message: string, meta?: Record<string, unknown>) => void;
+  messages?: unknown[];
 }): EmbeddedPiSubscribeContext {
   let compactionCount = params.initialCount;
   return {
     params: {
       runId: "run-test",
-      session: { messages: [] } as never,
+      session: { messages: params.messages ?? [] } as never,
       config: { session: { store: params.storePath } } as never,
       sessionKey: params.sessionKey,
       sessionId: "session-1",
@@ -306,4 +319,59 @@ describe("handleCompactionEnd", () => {
     expect(await readCompactionCount(storePath, sessionKey)).toBe(2);
     expect(ctx.noteCompactionTokensAfter).toHaveBeenCalledWith(undefined);
   });
+
+  it("keeps fresh assistant usage after the latest compaction summary", async () => {
+    const tmp = await fs.mkdtemp(path.join(os.tmpdir(), "openclaw-compaction-usage-"));
+    const storePath = path.join(tmp, "sessions.json");
+    const sessionKey = "main";
+    await seedSessionStore({
+      storePath,
+      sessionKey,
+      compactionCount: 0,
+    });
+    const staleUsage = makeUsage(120_000, 3_000, 123_000);
+    const freshUsage = makeUsage(1_000, 250, 1_250);
+    const messages: Array<Record<string, unknown> & { usage?: unknown }> = [
+      {
+        role: "assistant",
+        content: "pre-compaction answer",
+        timestamp: "2026-03-20T00:00:00.000Z",
+        usage: staleUsage,
+      },
+      {
+        role: "compactionSummary",
+        summary: "compressed",
+        timestamp: "2026-03-20T00:00:10.000Z",
+      },
+      {
+        role: "user",
+        content: "new question",
+        timestamp: "2026-03-20T00:00:20.000Z",
+      },
+      {
+        role: "assistant",
+        content: "fresh answer",
+        timestamp: "2026-03-20T00:00:30.000Z",
+        usage: freshUsage,
+      },
+    ];
+
+    const ctx = createCompactionContext({
+      storePath,
+      sessionKey,
+      initialCount: 0,
+      messages,
+    });
+
+    handleCompactionEnd(ctx, {
+      type: "compaction_end",
+      reason: "threshold",
+      result: { kept: 12 },
+      willRetry: false,
+      aborted: false,
+    });
+
+    expect(messages[0]?.usage).toEqual(makeZeroUsageSnapshot());
+    expect(messages[3]?.usage).toEqual(freshUsage);
+  });
 });
diff --git a/src/agents/pi-embedded-subscribe.handlers.compaction.ts b/src/agents/pi-embedded-subscribe.handlers.compaction.ts
@@ -183,16 +183,61 @@ function clearStaleAssistantUsageOnSessionMessages(ctx: EmbeddedPiSubscribeConte
   if (!Array.isArray(messages)) {
     return;
   }
-  for (const message of messages) {
+
+  let latestCompactionSummaryIndex = -1;
+  let latestCompactionTimestamp: number | null = null;
+  for (let i = 0; i < messages.length; i += 1) {
+    const entry = messages[i];
+    if (!entry || typeof entry !== "object") {
+      continue;
+    }
+    const candidate = entry as { role?: unknown; timestamp?: unknown };
+    if (candidate.role !== "compactionSummary") {
+      continue;
+    }
+    latestCompactionSummaryIndex = i;
+    latestCompactionTimestamp = parseMessageTimestamp(candidate.timestamp);
+  }
+
+  for (let i = 0; i < messages.length; i += 1) {
+    const message = messages[i];
     if (!message || typeof message !== "object") {
       continue;
     }
-    const candidate = message as { role?: unknown; usage?: unknown };
+    const candidate = message as {
+      role?: unknown;
+      timestamp?: unknown;
+      usage?: unknown;
+    };
     if (candidate.role !== "assistant") {
       continue;
     }
+    const staleByLegacyNoSummary = latestCompactionSummaryIndex === -1;
+    const messageTimestamp = parseMessageTimestamp(candidate.timestamp);
+    const staleByTimestamp =
+      latestCompactionTimestamp !== null &&
+      messageTimestamp !== null &&
+      messageTimestamp <= latestCompactionTimestamp;
+    const staleByLegacyOrdering =
+      latestCompactionSummaryIndex !== -1 && i < latestCompactionSummaryIndex;
+    if (!staleByLegacyNoSummary && !staleByTimestamp && !staleByLegacyOrdering) {
+      continue;
+    }
     // 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();
   }
 }
+
+function parseMessageTimestamp(value: unknown): number | null {
+  if (typeof value === "number" && Number.isFinite(value)) {
+    return value;
+  }
+  if (typeof value === "string") {
+    const parsed = Date.parse(value);
+    if (Number.isFinite(parsed)) {
+      return parsed;
+    }
+  }
+  return null;
+}
