pi-condense 2.6.0 → 2.7.0

package/index.ts

@@ -24,13 +24,23 @@ import { isProtected } from "./src/protected.js";
 import { registerQueryTool } from "./src/query-tool.js";
 import { registerCommands, setPruneStatusWidget } from "./src/commands.js";
 import { formatSummaryToolCallRefs, makeSummaryDetails, substituteInlineRefs } from "./src/summary-refs.js";
-import type { ContextPruneConfig, CapturedBatch, PruneFrontier, FlushOptions } from "./src/types.js";
+import type {
+  ContextPruneConfig,
+  CapturedBatch,
+  PruneFrontier,
+  FlushOptions,
+  ContextMetricsSnapshot,
+  FlushMetricsEntry,
+  FlushTrigger,
+} from "./src/types.js";
 import {
   DEFAULT_CONFIG,
   CUSTOM_TYPE_SUMMARY,
   CUSTOM_TYPE_STATS,
   CUSTOM_TYPE_FRONTIER,
+  CUSTOM_TYPE_FLUSH_METRICS,
 } from "./src/types.js";
+import { computeContextMetrics } from "./src/context-metrics.js";
 import { StatsAccumulator, emitExternalCost } from "./src/stats.js";
 import { PruneFrontierTracker } from "./src/frontier.js";
 import { BlockRefIssuer } from "./src/block-refs.js";
@@ -42,6 +52,8 @@ import { spillOversizedBatch } from "./src/spill.js";
 import { occKey } from "./src/occurrence-key.js";
 import { DiagnosticSink } from "./src/diagnostics.js";
 
+const EMPTY_METRICS_SNAPSHOT: ContextMetricsSnapshot = { openCycleThinkingTokens: 0, largestChainSharePct: 0, frontierGapTokens: 0 };
+
 export default function (pi: ExtensionAPI) {
   // Shared mutable config reference — updated by /pruner commands
   const currentConfig: { value: ContextPruneConfig } = {
@@ -74,6 +86,45 @@ export default function (pi: ExtensionAPI) {
   const pendingBatches: CapturedBatch[] = [];
   let isFlushing = false;
   let previousFraction: number | null = null;
+  // Set on session_start/session_tree when the branch rescan finds recoverable
+  // work but pendingBatches was just zeroed (reload/tree-switch). Lets the
+  // turn_end budget gate fire without a freshly pushed batch. Boolean only —
+  // no queue reconstruction; flushPending's own rescan is the data path.
+  // Cleared on every non-concurrent flushPending invocation.
+  let rearmedPending = false;
+
+  // Latest ContextMetricsSnapshot, recomputed at reload probes, batch capture,
+  // and flush entry. Cached (rather than recomputed on every widget refresh)
+  // because computeContextMetrics walks the full branch.
+  let metricsCache: ContextMetricsSnapshot | undefined;
+  const computeMetricsSnapshot = (ctx: any): ContextMetricsSnapshot | undefined => {
+    try {
+      // Includes persisted custom_message entries (e.g. this extension's own
+      // summary messages) alongside plain "message" entries: both are retained
+      // LLM context, so both belong in the largest-chain-share denominator.
+      // Projected inline (rather than importing pi-coding-agent's
+      // createCustomMessage) because that helper isn't re-exported from the
+      // package's "." export map -- shape mirrors createCustomMessage's output
+      // (role "custom"), which never matches the user/assistant/toolResult
+      // roles computeContextMetrics keys off, so it only inflates totalChars.
+      const branch = ctx.sessionManager.getBranch()
+        .filter((e: any) => (e.type === "message" && e.message) || e.type === "custom_message")
+        .map((e: any) =>
+          e.type === "custom_message"
+            ? { role: "custom", customType: e.customType, content: e.content, display: e.display, details: e.details, timestamp: new Date(e.timestamp).getTime() }
+            : e.message,
+        );
+      metricsCache = computeContextMetrics(
+        branch,
+        frontier.get(),
+        (k: string) => indexer.isSummarized(k),
+        protectionPredicate,
+      );
+    } catch (err) {
+      console.error("pi-condense: context metrics computation failed", err);
+    }
+    return metricsCache;
+  };
 
   type FlushResult =
     | { ok: true; reason: "flushed" | "skipped-oversized" | "skipped-trivial" | "skipped-deduped"; batchCount: number; toolCallCount: number; rawCharCount: number; summaryCharCount: number; dedupedCount?: number }
@@ -135,12 +186,18 @@ export default function (pi: ExtensionAPI) {
   // ── Helper: capture + trim + group pending batches (no LLM work) ──────────
   // Exposed to commands.ts via registerCommands so /pruner now can preview the
   // queue before opening the multi-row progress overlay.
-  const capturePendingBatches = (ctx: any): CapturedBatch[] => {
+  // `rethrow` is for the reload rearm probe only (session_start/session_tree):
+  // it needs to observe a rescan failure so it can console.error and leave
+  // rearmedPending false, per spec. Every other caller (turn_end capture path,
+  // flushPending, /pruner commands) keeps the existing swallow-and-fall-back
+  // behavior so a transient getBranch failure there never blocks the turn.
+  const capturePendingBatches = (ctx: any, opts?: { rethrow?: boolean }): CapturedBatch[] => {
     let batches: CapturedBatch[] = [];
     try {
       const branch = ctx.sessionManager.getBranch();
       batches = captureUnindexedBatchesFromSession(branch, indexer, protectionPredicate);
-    } catch {
+    } catch (err) {
+      if (opts?.rethrow) throw err;
       batches = pendingBatches.slice();
     }
     batches = batches
@@ -175,47 +232,97 @@ export default function (pi: ExtensionAPI) {
   const flushPending = async (ctx: any, options: FlushOptions = {}): Promise<FlushResult> => {
     if (isFlushing) return { ok: false, reason: "already-flushing" };
 
-    // Use pre-captured batches if provided (avoids double-capture when the
-    // caller previewed the queue before opening the progress overlay).
-    const batches: CapturedBatch[] = options.previewedBatches ?? capturePendingBatches(ctx);
+    // Clear on every non-concurrent invocation, regardless of outcome — the
+    // rearm is a one-shot nudge for the very next eligible gate check.
+    rearmedPending = false;
 
-    if (batches.length === 0) return { ok: false, reason: "empty" };
+    // Pre-flush pressure snapshot — recorded once at flush entry so the
+    // observability entry reflects what triggered this attempt, not what's
+    // left after it ran.
+    const entryMetrics: ContextMetricsSnapshot = computeMetricsSnapshot(ctx) ?? EMPTY_METRICS_SNAPSHOT;
+    const trigger: FlushTrigger = options.trigger ?? "manual";
+    const delivery = options.delivery ?? "runtime";
 
-    // Bail out before we drain pendingBatches so they don't need restoring.
-    if (options.signal?.aborted) return { ok: false, reason: "aborted" };
+    // One-entry-per-attempt tracking, emitted once from the outer `finally`
+    // below. `appendEntry` is assigned only once `sessionManager` is captured
+    // (session delivery); until then (empty/aborted/pre-capture-failure exits)
+    // the emitter falls back to pi.appendEntry.
+    let capturedBatches = 0;
+    let processedCount = 0;
+    let outcome: FlushMetricsEntry["outcome"] = "empty";
+    let appendEntry: ((customType: string, data?: unknown) => void) | undefined;
+
+    // Non-fatal by construction: observability must never affect the flush outcome.
+    const emitFlushMetricsOnce = () => {
+      const entry: FlushMetricsEntry = {
+        ts: Date.now(),
+        trigger,
+        capturedBatches,
+        processedBatches: processedCount,
+        outcome,
+        metrics: entryMetrics,
+      };
+      const appender: (type: string, data: unknown) => void = appendEntry
+        ? delivery === "runtime" ? (type, data) => pi.appendEntry(type, data) : appendEntry
+        : (type, data) => pi.appendEntry(type, data);
+      try {
+        appender(CUSTOM_TYPE_FLUSH_METRICS, entry);
+      } catch {
+        // non-fatal: observability must never fail the flush
+      }
+    };
 
-    // Draining the queue since we've captured the state via session or slice.
-    // We drain BEFORE the await so concurrent calls (though guarded by isFlushing)
-    // or rapid turn-ends don't result in double-summarization.
-    pendingBatches.length = 0;
+    let batches: CapturedBatch[] = [];
+    let sessionManager: SessionAppender | undefined;
+    try {
+      // Bind the session appender as soon as delivery is known, BEFORE the
+      // empty-capture/aborted exits below — so emitFlushMetricsOnce's finally
+      // emit routes through sessionManager for those exits too, instead of
+      // falling back to the (possibly stale, print-mode) pi.appendEntry.
+      if (delivery === "session") {
+        try {
+          sessionManager = ctx.sessionManager as unknown as SessionAppender;
+          appendEntry = (customType: string, data?: unknown) => sessionManager!.appendCustomEntry(customType, data);
+        } catch (err) {
+          outcome = "error";
+          return { ok: false, reason: isStaleContextError(err) ? "stale-context" : "failed", error: errorMessage(err) };
+        }
+      }
 
-    isFlushing = true;
+      // Use pre-captured batches if provided (avoids double-capture when the
+      // caller previewed the queue before opening the progress overlay).
+      batches = options.previewedBatches ?? capturePendingBatches(ctx);
+      capturedBatches = batches.length;
 
-    const delivery = options.delivery ?? "runtime";
-    let sessionManager: SessionAppender | undefined;
-    if (delivery === "session") {
-      try {
-        sessionManager = ctx.sessionManager as unknown as SessionAppender;
-      } catch (err) {
-        restoreBatches(batches);
-        isFlushing = false;
-        return { ok: false, reason: isStaleContextError(err) ? "stale-context" : "failed", error: errorMessage(err) };
+      if (batches.length === 0) {
+        outcome = "empty";
+        return { ok: false, reason: "empty" };
       }
-    }
 
-    const appendEntry = (customType: string, data?: unknown) => sessionManager!.appendCustomEntry(customType, data);
-    const appendSummaryMessage = (content: string, details: unknown) =>
-      sessionManager!.appendCustomMessageEntry(CUSTOM_TYPE_SUMMARY, content, false, details);
+      // Bail out before we drain pendingBatches so they don't need restoring.
+      if (options.signal?.aborted) {
+        outcome = "error";
+        return { ok: false, reason: "aborted" };
+      }
 
-    // Routes alias persistence through whichever delivery is active so the
-    // dedup pre-flush pass writes CUSTOM_TYPE_DEDUP_ALIAS entries via the
-    // same path the rest of the flush uses.
-    const persistAlias: (customType: string, data?: unknown) => void =
-      delivery === "runtime"
-        ? (type, data) => pi.appendEntry(type, data)
-        : appendEntry;
+      // Draining the queue since we've captured the state via session or slice.
+      // We drain BEFORE the await so concurrent calls (though guarded by isFlushing)
+      // or rapid turn-ends don't result in double-summarization.
+      pendingBatches.length = 0;
+
+      isFlushing = true;
+
+      const appendSummaryMessage = (content: string, details: unknown) =>
+        sessionManager!.appendCustomMessageEntry(CUSTOM_TYPE_SUMMARY, content, false, details);
+
+      // Routes alias persistence through whichever delivery is active so the
+      // dedup pre-flush pass writes CUSTOM_TYPE_DEDUP_ALIAS entries via the
+      // same path the rest of the flush uses.
+      const persistAlias: (customType: string, data?: unknown) => void =
+        delivery === "runtime"
+          ? (type, data) => pi.appendEntry(type, data)
+          : appendEntry!;
 
-    try {
       // ── Pre-flush content-hash dedup pass ────────────────────────────
       // For each tool call, check the indexer's contentHashToOriginal map.
       // A hit means an identical (toolName, normalized resultText) pair has
@@ -431,7 +538,7 @@ export default function (pi: ExtensionAPI) {
             } else {
               appendSummaryMessage(summaryText, batchDetails);
               indexer.registerSummaryRefs(summaryRefs);
-              indexer.addBatch(batch, appendEntry);
+              indexer.addBatch(batch, appendEntry!);
             }
             // Keep the in-memory summary-body registry current so chain compression
             // can build synthetic chain messages without rescanning session entries.
@@ -459,7 +566,8 @@ export default function (pi: ExtensionAPI) {
 
       if (processedBatches.length === 0) {
         // Nothing was persisted (all calls failed or first call failed)
-        setPruneStatusWidget(ctx, currentConfig.value, statsAccum.getLiveReclaim(), diagnostics.counts());
+        setPruneStatusWidget(ctx, currentConfig.value, statsAccum.getLiveReclaim(), diagnostics.counts(), metricsCache);
+        outcome = "error";
         return { ok: false, reason: "summarizer-failed" };
       }
 
@@ -521,18 +629,22 @@ export default function (pi: ExtensionAPI) {
           statsAccum.persist(pi);
         } else {
           frontier.advance(frontierSnapshot);
-          appendEntry(CUSTOM_TYPE_FRONTIER, frontierSnapshot);
+          appendEntry!(CUSTOM_TYPE_FRONTIER, frontierSnapshot);
           try {
-            appendEntry(CUSTOM_TYPE_STATS, statsAccum.getStats());
+            appendEntry!(CUSTOM_TYPE_STATS, statsAccum.getStats());
           } catch {
             // Ignore stats persistence failures; the prune result and frontier are the contract.
           }
         }
       } catch (err) {
+        // Batches were summarized/persisted before the frontier/stats write failed;
+        // reflect that in processedBatches rather than reporting 0.
+        processedCount = processedBatches.length;
+        outcome = "error";
         return { ok: false, reason: isStaleContextError(err) ? "stale-context" : "failed", error: errorMessage(err) };
       }
 
-      setPruneStatusWidget(ctx, currentConfig.value, statsAccum.getLiveReclaim(), diagnostics.counts());
+      setPruneStatusWidget(ctx, currentConfig.value, statsAccum.getLiveReclaim(), diagnostics.counts(), metricsCache);
       emitExternalCost(pi, statsAccum);
 
       // Chain compression — compress closed chains beyond the rolling window.
@@ -620,6 +732,12 @@ export default function (pi: ExtensionAPI) {
         }
       }
 
+      // Very end of the try block, deliberately after (and outside) the
+      // chain-compression block's own try/catch above: a compression failure
+      // must not eat this entry — the summarization phase already succeeded.
+      processedCount = processedBatches.length;
+      outcome = flushOutcome;
+
       const returnReason: "flushed" | "skipped-oversized" | "skipped-trivial" | "skipped-deduped" =
         actuallyFlushedCount > 0
           ? "flushed"
@@ -640,10 +758,11 @@ export default function (pi: ExtensionAPI) {
       };
     } catch (err) {
       restoreBatches(batches);
+      outcome = "error";
       // When the abort signal fired, summarizeBatch rethrows rather than
       // swallowing the error.  Don't show a UI error — the user intended this.
       if (options.signal?.aborted) {
-        setPruneStatusWidget(ctx, currentConfig.value, statsAccum.getLiveReclaim(), diagnostics.counts());
+        setPruneStatusWidget(ctx, currentConfig.value, statsAccum.getLiveReclaim(), diagnostics.counts(), metricsCache);
         return { ok: false, reason: "aborted" };
       }
       if (isStaleContextError(err)) {
@@ -653,6 +772,7 @@ export default function (pi: ExtensionAPI) {
       return { ok: false, reason: "failed", error: errorMessage(err) };
     } finally {
       isFlushing = false;
+      emitFlushMetricsOnce();
     }
   };
 
@@ -678,9 +798,19 @@ export default function (pi: ExtensionAPI) {
     // Clear any batches queued before the session reload
     pendingBatches.length = 0;
     previousFraction = null;
+    rearmedPending = false;
+    if (currentConfig.value.enabled) {
+      try {
+        rearmedPending = capturePendingBatches(ctx, { rethrow: true }).length > 0;
+      } catch (err) {
+        console.error("pi-condense: reload rearm probe failed", err);
+      }
+    }
+
+    computeMetricsSnapshot(ctx);
 
     // Update footer status
-    setPruneStatusWidget(ctx, currentConfig.value, statsAccum.getLiveReclaim(), diagnostics.counts());
+    setPruneStatusWidget(ctx, currentConfig.value, statsAccum.getLiveReclaim(), diagnostics.counts(), metricsCache);
 
     ctx.ui.setWidget(
       "pruner-boot",
@@ -708,6 +838,17 @@ export default function (pi: ExtensionAPI) {
     // Pending batches belong to the old branch — discard them
     pendingBatches.length = 0;
     previousFraction = null;
+    rearmedPending = false;
+    if (currentConfig.value.enabled) {
+      try {
+        rearmedPending = capturePendingBatches(ctx, { rethrow: true }).length > 0;
+      } catch (err) {
+        console.error("pi-condense: reload rearm probe failed", err);
+      }
+    }
+
+    computeMetricsSnapshot(ctx);
+    setPruneStatusWidget(ctx, currentConfig.value, statsAccum.getLiveReclaim(), diagnostics.counts(), metricsCache);
   });
 
   // ── turn_end: capture batch, flush immediately or queue ──────────────────
@@ -716,73 +857,91 @@ export default function (pi: ExtensionAPI) {
 
     const hasToolResults = event.toolResults && event.toolResults.length > 0;
 
-    if (!hasToolResults) {
-      // Text-only final turns are handled by message_end in agent-message mode.
-      // In print mode, turn_end can fire after session shutdown, so do not start
-      // deferred LLM work from this late lifecycle event.
-      return;
-    }
+    // Text-only final turns are handled by message_end in agent-message mode.
+    // In print mode, turn_end can fire after session shutdown, so do not start
+    // deferred LLM work from this late lifecycle event — UNLESS a reload probe
+    // (session_start/session_tree) found recoverable pending work: that flag
+    // must still reach the budget/delta gate below without a freshly captured
+    // batch on this turn.
+    if (!hasToolResults && !rearmedPending) return;
+
+    let pushedBatch = false;
+    if (hasToolResults) {
+      const capturedBatch = captureBatch(
+        event.message,
+        event.toolResults,
+        event.turnIndex,
+        Date.now()
+      );
+      // Drop user-protected tool/path results so they stay verbatim in context.
+      // Filtering at capture time keeps the
+      // underlying assistant `toolCall` block AND its `ToolResultMessage`
+      // untouched in Pi's session/event stream — only the in-memory
+      // CapturedBatch is pruned, which is exactly what we want.
+      const filtered = {
+        ...capturedBatch,
+        toolCalls: capturedBatch.toolCalls.filter((tc) => !isProtected(tc.toolName, tc.args, currentConfig.value)),
+      };
 
-    const capturedBatch = captureBatch(
-      event.message,
-      event.toolResults,
-      event.turnIndex,
-      Date.now()
-    );
-    // Drop user-protected tool/path results so they stay verbatim in context.
-    // Filtering at capture time keeps the
-    // underlying assistant `toolCall` block AND its `ToolResultMessage`
-    // untouched in Pi's session/event stream — only the in-memory
-    // CapturedBatch is pruned, which is exactly what we want.
-    const filtered = {
-      ...capturedBatch,
-      toolCalls: capturedBatch.toolCalls.filter((tc) => !isProtected(tc.toolName, tc.args, currentConfig.value)),
-    };
+      // Eager spill: offload oversized single results to sidecar files before they
+      // ever reach a request. addBatch inside marks them isSummarized, so
+      // trimBatchToPendingRange drops them from the pending set below. Best-effort:
+      // a spill failure leaves the result inline for the normal flush pipeline.
+      try {
+        await spillOversizedBatch({
+          batch: filtered,
+          indexer,
+          config: {
+            spillThreshold: currentConfig.value.spillThreshold,
+            spillPreviewBytes: currentConfig.value.spillPreviewBytes,
+            dedupByContentHash: currentConfig.value.dedupByContentHash,
+          },
+          sessionDir: ctx.sessionManager.getSessionDir(),
+          sessionId: ctx.sessionManager.getSessionId(),
+          appendEntry: (type, data) => (ctx.sessionManager as unknown as SessionAppender).appendCustomEntry(type, data),
+        });
+      } catch {
+        // best-effort; never block the turn
+      }
 
-    // Eager spill: offload oversized single results to sidecar files before they
-    // ever reach a request. addBatch inside marks them isSummarized, so
-    // trimBatchToPendingRange drops them from the pending set below. Best-effort:
-    // a spill failure leaves the result inline for the normal flush pipeline.
-    try {
-      await spillOversizedBatch({
-        batch: filtered,
-        indexer,
-        config: {
-          spillThreshold: currentConfig.value.spillThreshold,
-          spillPreviewBytes: currentConfig.value.spillPreviewBytes,
-          dedupByContentHash: currentConfig.value.dedupByContentHash,
-        },
-        sessionDir: ctx.sessionManager.getSessionDir(),
-        sessionId: ctx.sessionManager.getSessionId(),
-        appendEntry: (type, data) => (ctx.sessionManager as unknown as SessionAppender).appendCustomEntry(type, data),
-      });
-    } catch {
-      // best-effort; never block the turn
+      const batch = trimBatchToPendingRange(filtered);
+      if (batch) {
+        pushedBatch = true;
+        pendingBatches.push(batch);
+
+        // Let the user know a batch is queued
+        const n = pendingBatches.length;
+        const trigger = currentConfig.value.pruneOn === "agent-message"
+          ? "agent's next text response"
+          : "/pruner now";
+        if (currentConfig.value.showPruneStatusLine) {
+          setPruneStatusWidget(ctx, currentConfig.value, `prune: ${n} pending`);
+          safeNotify(
+            ctx,
+            `pruner: ${n} turn${n === 1 ? "" : "s"} queued — will summarize on ${trigger}`,
+            "info"
+          );
+        }
+      }
     }
 
-    const batch = trimBatchToPendingRange(filtered);
-    if (!batch) return;
+    // Recompute regardless of whether trim produced a batch: a turn whose
+    // toolResults are all protected/spilled/summarized/trimmed-empty still
+    // changes the branch (thinking, open-segment size), so the cache must not
+    // go stale on it. Placed before the pushedBatch/rearmedPending early
+    // return below — a cache write is not gate evaluation.
+    if (hasToolResults) computeMetricsSnapshot(ctx);
 
-    pendingBatches.push(batch);
-
-    // Let the user know a batch is queued
-    const n = pendingBatches.length;
-    const trigger = currentConfig.value.pruneOn === "agent-message"
-      ? "agent's next text response"
-      : "/pruner now";
-    if (currentConfig.value.showPruneStatusLine) {
-      setPruneStatusWidget(ctx, currentConfig.value, `prune: ${n} pending`);
-      safeNotify(
-        ctx,
-        `pruner: ${n} turn${n === 1 ? "" : "s"} queued — will summarize on ${trigger}`,
-        "info"
-      );
-    }
+    // Mirrors main's `if (!batch) return;`: no freshly pushed batch this turn
+    // means no gate evaluation, regardless of leftover pendingBatches from an
+    // earlier turn — UNLESS a reload probe armed rearmedPending, in which case
+    // the gate below must still run.
+    if (!pushedBatch && !rearmedPending) return;
 
     // Token-budget auto-flush: an additional, mode-independent trigger. When context
     // usage crosses autoBudgetThreshold, compact the queued batches now instead of
-    // waiting for this mode's flush boundary. The pendingBatches.length guard makes
-    // an already-drained queue a no-op.
+    // waiting for this mode's flush boundary. The pendingBatches.length-or-rearmed
+    // guard makes an already-drained, non-rearmed queue a no-op.
     const usage = ctx.getContextUsage?.();
     const budgetHit = shouldBudgetFlush(usage, currentConfig.value.autoBudgetThreshold);
     const deltaHit = shouldDeltaFlush(usage, previousFraction, currentConfig.value.budgetTurnDelta);
@@ -791,16 +950,19 @@ export default function (pi: ExtensionAPI) {
     const f = usageFraction(usage);
     if (f != null) previousFraction = f;
 
-    if (pendingBatches.length > 0 && !isFlushing && (budgetHit || deltaHit)) {
+    const n = pendingBatches.length;
+    if ((n > 0 || rearmedPending) && !isFlushing && (budgetHit || deltaHit)) {
       // Always surface this flush (even when the routine status line is off): it's a
       // significant, infrequent event — context crossed a threshold or jumped sharply
       // this turn — and it self-throttles because pendingBatches is drained right after.
       safeNotify(
         ctx,
-        `pruner: ${budgetHit ? "context budget reached" : "context jumped this turn"} — compacting ${n} pending turn${n === 1 ? "" : "s"}`,
+        n > 0
+          ? `pruner: ${budgetHit ? "context budget reached" : "context jumped this turn"} — compacting ${n} pending turn${n === 1 ? "" : "s"}`
+          : `pruner: ${budgetHit ? "context budget reached" : "context jumped this turn"} — compacting work recovered after reload`,
         "info",
       );
-      await flushPending(ctx, { delivery: "session" });
+      await flushPending(ctx, { delivery: "session", trigger: n === 0 ? "rearmed" : budgetHit ? "budget" : "delta" });
     }
   });
 
@@ -813,7 +975,7 @@ export default function (pi: ExtensionAPI) {
     if (!currentConfig.value.enabled) return;
     if (currentConfig.value.pruneOn !== "agent-message") return;
     if (!isFinalAssistantMessage(event.message)) return;
-    await flushPending(ctx, { delivery: "session", closingMessage: event.message });
+    await flushPending(ctx, { delivery: "session", closingMessage: event.message, trigger: "message-end" });
   });
 
   // ── agent_end: last-chance cleanup only ─────────────────────────────────────
@@ -821,8 +983,12 @@ export default function (pi: ExtensionAPI) {
   // already be disposing the session, so avoid starting a best-effort LLM call here.
   pi.on("agent_end", async (_event, ctx) => {
     if (!currentConfig.value.enabled) return;
-    if (pendingBatches.length === 0) return;
-    setPruneStatusWidget(ctx, currentConfig.value, `prune: ${pendingBatches.length} pending`);
+    if (pendingBatches.length === 0 && !rearmedPending) return;
+    setPruneStatusWidget(
+      ctx,
+      currentConfig.value,
+      pendingBatches.length > 0 ? `prune: ${pendingBatches.length} pending` : "prune: recovered pending (reload)",
+    );
   });
 
   // ── context: prune summarized tool results from next LLM call ─────────────
@@ -851,7 +1017,7 @@ export default function (pi: ExtensionAPI) {
       changed = true;
       statsAccum.setLiveReclaim(result.beforeChars, result.afterChars);
     }
-    setPruneStatusWidget(ctx, currentConfig.value, statsAccum.getLiveReclaim(), diagnostics.counts());
+    setPruneStatusWidget(ctx, currentConfig.value, statsAccum.getLiveReclaim(), diagnostics.counts(), metricsCache);
 
     if (!changed) return undefined;
     return { messages };
@@ -888,5 +1054,18 @@ export default function (pi: ExtensionAPI) {
     return { compressedEntries: result.compressedEntries, skipped: result.skipped.filter((s) => s.reason === "no-summary").length };
   };
 
-  registerCommands(pi, currentConfig, flushPending, capturePendingBatches, () => statsAccum.getStats(), () => statsAccum.getLiveReclaim(), indexer, compactChains, () => diagnostics.counts());
+  registerCommands(
+    pi,
+    currentConfig,
+    flushPending,
+    capturePendingBatches,
+    () => statsAccum.getStats(),
+    () => statsAccum.getLiveReclaim(),
+    indexer,
+    compactChains,
+    () => diagnostics.counts(),
+    (ctx: any) => computeMetricsSnapshot(ctx) ?? EMPTY_METRICS_SNAPSHOT,
+    () => metricsCache,
+    () => rearmedPending,
+  );
 }

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "pi-condense",
-  "version": "2.6.0",
+  "version": "2.7.0",
   "description": "Pi coding-agent extension that summarizes completed tool-call batches, replaces raw outputs with short stubs, compresses closed tool-call chains, and recovers any original on demand via context_tree_query.",
   "author": "Jacek Juraszek",
   "license": "MIT",