fix(compaction): preserve fresh usage tail pressure

jbetala7 · jbetala7 · commit 31b80410d4d7 · 2026-05-15T04:26:01.000+05:30
diff --git a/src/auto-reply/reply/agent-runner-memory.test.ts b/src/auto-reply/reply/agent-runner-memory.test.ts
@@ -746,6 +746,11 @@ describe("runMemoryFlushIfNeeded", () => {
             usage: { input: 240_000, output: 1_000 },
           },
         }),
+        JSON.stringify({
+          type: "compaction",
+          timestamp: new Date(1_700_000_000_000).toISOString(),
+          result: { summary: "compacted" },
+        }),
         JSON.stringify({
           message: {
             role: "assistant",
@@ -955,6 +960,79 @@ describe("runMemoryFlushIfNeeded", () => {
     expect(compactCall.currentTokenCount).toBeGreaterThanOrEqual(96_000);
   });
 
+  it("preserves latest usage plus tail pressure when no post-compaction marker exists", async () => {
+    const sessionFile = path.join(rootDir, "fresh-usage-tail-pressure-session.jsonl");
+    await fs.writeFile(
+      sessionFile,
+      [
+        JSON.stringify({
+          type: "session",
+          id: "session",
+        }),
+        JSON.stringify({
+          id: "u1",
+          parentId: null,
+          message: {
+            role: "user",
+            content: "continue the current long context turn",
+          },
+        }),
+        JSON.stringify({
+          id: "a1",
+          parentId: "u1",
+          message: {
+            role: "assistant",
+            content: "small latest answer",
+            usage: { input: 94_000, output: 500 },
+          },
+        }),
+        JSON.stringify({
+          id: "t1",
+          parentId: "a1",
+          message: {
+            role: "tool",
+            content: `post-usage tool tail ${"x".repeat(36_000)}`,
+          },
+        }),
+      ].join("\n"),
+      "utf8",
+    );
+    registerMemoryFlushPlanResolverForTest(() => ({
+      softThresholdTokens: 4_000,
+      forceFlushTranscriptBytes: 1_000_000_000,
+      reserveTokensFloor: 0,
+      prompt: "Pre-compaction memory flush.\nNO_REPLY",
+      systemPrompt: "Write memory to memory/YYYY-MM-DD.md.",
+      relativePath: "memory/2023-11-14.md",
+    }));
+    const sessionEntry: SessionEntry = {
+      sessionId: "session",
+      sessionFile,
+      updatedAt: Date.now(),
+      totalTokensFresh: false,
+    };
+
+    await runPreflightCompactionIfNeeded({
+      cfg: { agents: { defaults: { compaction: { memoryFlush: {} } } } },
+      followupRun: createTestFollowupRun({
+        sessionId: "session",
+        sessionFile,
+        sessionKey: "main",
+      }),
+      defaultModel: "anthropic/claude-opus-4-6",
+      agentCfgContextTokens: 100_000,
+      sessionEntry,
+      sessionStore: { main: sessionEntry },
+      sessionKey: "main",
+      storePath: path.join(rootDir, "sessions.json"),
+      isHeartbeat: false,
+      replyOperation: createReplyOperation(),
+    });
+
+    const compactCall = requireCompactEmbeddedPiSessionCall();
+    expect(compactCall.currentTokenCount).toBeGreaterThanOrEqual(100_000);
+  });
+
   it("does not count bytes from a large latest usage record as post-usage tail pressure", async () => {
     const sessionFile = path.join(rootDir, "large-usage-record-session.jsonl");
     await fs.writeFile(
diff --git a/src/auto-reply/reply/agent-runner-memory.ts b/src/auto-reply/reply/agent-runner-memory.ts
@@ -191,6 +191,7 @@ export type SessionTranscriptUsageSnapshot = {
   promptTokens?: number;
   outputTokens?: number;
   trailingBytesTokens?: number;
+  hasPostUsageCompactionMarker?: boolean;
 };
 
 // Keep a generous near-threshold window so large assistant outputs still trigger
@@ -220,6 +221,58 @@ function parseUsageFromTranscriptLine(line: string): ReturnType<typeof normalize
   return undefined;
 }
 
+function collectTranscriptText(value: unknown): string {
+  if (typeof value === "string") {
+    return value;
+  }
+  if (!Array.isArray(value)) {
+    return "";
+  }
+  return value
+    .map((item) => {
+      if (!item || typeof item !== "object" || Array.isArray(item)) {
+        return "";
+      }
+      const text = (item as { text?: unknown }).text;
+      return typeof text === "string" ? text : "";
+    })
+    .filter(Boolean)
+    .join("\n");
+}
+
+function transcriptLineHasPostUsageCompactionMarker(line: string): boolean {
+  const trimmed = line.trim();
+  if (!trimmed) {
+    return false;
+  }
+  try {
+    const parsed = JSON.parse(trimmed) as {
+      type?: unknown;
+      message?: { content?: unknown };
+      payload?: { type?: unknown; text?: unknown };
+    };
+    if (parsed.type === "compaction" || parsed.type === "session.compacted") {
+      return true;
+    }
+    const payloadType = parsed.payload?.type;
+    if (payloadType === "compaction" || payloadType === "session.compacted") {
+      return true;
+    }
+    const text = [
+      collectTranscriptText(parsed.message?.content),
+      collectTranscriptText(parsed.payload?.text),
+    ]
+      .filter(Boolean)
+      .join("\n");
+    return (
+      text.includes("[Post-compaction context refresh]") ||
+      text.includes("Session was just compacted.")
+    );
+  } catch {
+    return false;
+  }
+}
+
 function resolveSessionLogPath(
   sessionId?: string,
   sessionEntry?: SessionEntry,
@@ -257,6 +310,7 @@ function deriveTranscriptUsageSnapshot(
     | {
         usage: ReturnType<typeof normalizeUsage> | undefined;
         trailingBytes?: number;
+        hasPostUsageCompactionMarker?: boolean;
       }
     | undefined,
 ): SessionTranscriptUsageSnapshot | undefined {
@@ -276,6 +330,7 @@ function deriveTranscriptUsageSnapshot(
   return {
     promptTokens,
     outputTokens,
+    hasPostUsageCompactionMarker: snapshot.hasPostUsageCompactionMarker === true ? true : undefined,
     trailingBytesTokens:
       typeof snapshot.trailingBytes === "number" &&
       Number.isFinite(snapshot.trailingBytes) &&
@@ -360,6 +415,7 @@ async function readLastNonzeroUsageFromSessionLog(logPath: string) {
     const stat = await handle.stat();
     let position = stat.size;
     let leadingPartial = "";
+    let postUsageCompactionMarkerSeen = false;
     while (position > 0) {
       const chunkSize = Math.min(TRANSCRIPT_TAIL_CHUNK_BYTES, position);
       const start = position - chunkSize;
@@ -384,16 +440,23 @@ async function readLastNonzeroUsageFromSessionLog(logPath: string) {
           return {
             usage,
             trailingBytes: suffixBytesOutsideCombined + trailingBytesInChunk,
+            hasPostUsageCompactionMarker:
+              postUsageCompactionMarkerSeen ||
+              trailingLines.some(transcriptLineHasPostUsageCompactionMarker),
           };
         }
       }
+      if (!postUsageCompactionMarkerSeen) {
+        postUsageCompactionMarkerSeen = lines.some(transcriptLineHasPostUsageCompactionMarker);
+      }
       position = start;
     }
     const usage = parseUsageFromTranscriptLine(leadingPartial);
     return usage
       ? {
           usage,
           trailingBytes: Math.max(0, stat.size - Buffer.byteLength(leadingPartial, "utf8")),
+          hasPostUsageCompactionMarker: postUsageCompactionMarkerSeen,
         }
       : undefined;
   } finally {
@@ -461,6 +524,7 @@ async function estimatePromptTokensFromSessionTranscript(params: {
       const outputTokens = snapshot.usage?.outputTokens;
       const rawUsagePromptTokens = Math.ceil(promptTokens);
       const boundedUsagePromptTokens =
+        snapshot.usage?.hasPostUsageCompactionMarker === true &&
         typeof estimatedMessageTokens === "number" &&
         rawUsagePromptTokens > estimatedMessageTokens * 2 + 10_000
           ? estimatedMessageTokens
