fix: invalidate retained context usage after compaction

packages/agent-modules/context-manager/src/manager.ts

@@ -1,8 +1,4 @@
-import {
-  generateSummary,
-  type AgentMessage,
-  type CompactionSummaryMessage,
-} from '@earendil-works/pi-agent-core';
+import { generateSummary, type AgentMessage } from '@earendil-works/pi-agent-core';
 import type {
   PiBeforeLlmCallHook,
   PiBeforeLlmCallHookDecision,
@@ -13,6 +9,7 @@ import { computeCompactionTriggerAt, DEFAULT_CONTEXT_MANAGER_SETTINGS } from './
 import { createDefaultContextTokenEstimator, type TokenEstimator } from './token-estimator.js';
 import type {
   ContextCompactionPlan,
+  ContextCompactionSummaryMessage,
   ContextManagerCheckpointOptions,
   ContextManagerLock,
   ContextManagerObserver,
@@ -116,7 +113,7 @@ export class ContextManager {
 
     const replacementId = this.idGenerator();
     const replacementMessages = [
-      createCompactionSummary(summary, plan.tokensBefore, this.nowMs()),
+      createCompactionSummary(summary, plan.tokensBefore, this.nowMs(), plan.keptMessages.length),
       ...plan.keptMessages,
     ];
     const tokensAfter = this.tokenEstimator.estimateContextTokens(replacementMessages).tokens;
@@ -510,12 +507,14 @@ function createCompactionSummary(
   summary: string,
   tokensBefore: number,
   timestamp: number,
-): CompactionSummaryMessage {
+  keptMessageCount: number,
+): ContextCompactionSummaryMessage {
   return {
     role: 'compactionSummary',
     summary,
     tokensBefore,
     timestamp,
+    keptMessageCount,
   };
 }
 

packages/agent-modules/context-manager/src/token-estimator.ts

@@ -20,7 +20,7 @@
  * Ground-truth-aware estimate
  * ---------------------------
  * `estimateContextTokens` preserves pi-agent-core's good pattern:
- *  1. Find the last successful assistant message with a usage block;
+ *  1. Find the last successful assistant message with usage for the current context;
  *  2. Trust its provider-reported total as the prefix sum;
  *  3. Estimate only the trailing messages after that point.
  * This bounds estimator error to the trailing window — typically a handful of
@@ -48,6 +48,7 @@
 import type { AgentMessage } from '@earendil-works/pi-agent-core';
 import { Buffer } from 'node:buffer';
 import { countTokens as countO200kBase } from 'gpt-tokenizer/model/gpt-4o';
+import type { ContextCompactionSummaryMessage } from './types.js';
 
 export interface ContextTokenEstimate {
   /** Estimated total tokens consumed by `messages`. */
@@ -163,8 +164,36 @@ function getAssistantUsage(message: AgentMessage): AssistantUsageRef | undefined
 function getLastAssistantUsageInfo(
   messages: AgentMessage[],
 ): { usage: AssistantUsageRef; index: number } | undefined {
+  let firstFreshIndex = 0;
+  let compactedAt: number | undefined;
   for (let i = messages.length - 1; i >= 0; i -= 1) {
-    const usage = getAssistantUsage(messages[i]!);
+    if (messages[i]!.role !== 'compactionSummary') continue;
+    const summary = messages[i] as ContextCompactionSummaryMessage;
+    const keptCount = summary.keptMessageCount;
+    if (keptCount !== undefined) {
+      // Kept assistants appear AFTER the summary too. Only appended responses
+      // describe the replacement context; their wall-clock timestamps may tie
+      // or precede the summary if the clock changed.
+      if (!Number.isSafeInteger(keptCount) || keptCount < 0) return undefined;
+      firstFreshIndex = i + 1 + keptCount;
+    } else {
+      // Legacy persisted transcripts have no explicit retained-tail boundary.
+      // Match the local runtime's conservative timestamp freshness rule.
+      if (!Number.isFinite(summary.timestamp)) return undefined;
+      compactedAt = summary.timestamp;
+      firstFreshIndex = i + 1;
+    }
+    break;
+  }
+  for (let i = messages.length - 1; i >= firstFreshIndex; i -= 1) {
+    const message = messages[i]!;
+    if (
+      compactedAt !== undefined &&
+      (!Number.isFinite(message.timestamp) || message.timestamp <= compactedAt)
+    ) {
+      continue;
+    }
+    const usage = getAssistantUsage(message);
     if (usage) return { usage, index: i };
   }
   return undefined;

packages/agent-modules/context-manager/src/types.ts

@@ -1,6 +1,20 @@
-import type { AgentMessage, StreamFn, ThinkingLevel } from '@earendil-works/pi-agent-core';
+import type {
+  AgentMessage,
+  CompactionSummaryMessage,
+  StreamFn,
+  ThinkingLevel,
+} from '@earendil-works/pi-agent-core';
 import type { Api, Model, Tool } from '@earendil-works/pi-ai';
 
+export interface ContextCompactionSummaryMessage extends CompactionSummaryMessage {
+  /**
+   * Number of following messages retained from the old context, whose usage is stale.
+   * Persist this with the summary so a restored transcript has the same usage boundary.
+   * Absent on legacy summaries, which must use timestamp-based freshness instead.
+   */
+  keptMessageCount?: number;
+}
+
 export interface ContextManagerSettings {
   enabled: boolean;
   reserveTokens: number;