From f5bb89c8e55502fb6c4726a592b7a7f60d1e0566 Mon Sep 17 00:00:00 2001
From: Daniel Smolsky <dannysmo@gmail.com>
Date: Thu, 22 Jan 2026 00:04:55 -0500
Subject: [PATCH] feat: show extracted token count in prune notifications

Add token count display for extractions in both minimal and detailed
notification formats. The extracted token count always appears when
distillation is present, while the full distillation text remains
controlled by the showDistillation config setting.
---
 lib/ui/notification.ts | 36 ++++++++++++++++++++++++------------
 lib/ui/utils.ts        |  6 ++++++
 2 files changed, 30 insertions(+), 12 deletions(-)

diff --git a/lib/ui/notification.ts b/lib/ui/notification.ts
index 81bba142..acb948cd 100644
--- a/lib/ui/notification.ts
+++ b/lib/ui/notification.ts
@@ -1,6 +1,7 @@
 import type { Logger } from "../logger"
 import type { SessionState } from "../state"
 import {
+    countDistillationTokens,
     formatExtracted,
     formatPrunedItemsList,
     formatStatsHeader,
@@ -19,14 +20,19 @@ export const PRUNE_REASON_LABELS: Record<PruneReason, string> = {
 function buildMinimalMessage(
     state: SessionState,
     reason: PruneReason | undefined,
-    distillation?: string[],
+    distillation: string[] | undefined,
+    showDistillation: boolean,
 ): string {
-    const reasonSuffix = reason ? ` — ${PRUNE_REASON_LABELS[reason]}` : ""
+    const extractedTokens = countDistillationTokens(distillation)
+    const extractedSuffix =
+        extractedTokens > 0 ? ` (extracted ${formatTokenCount(extractedTokens)})` : ""
+    const reasonSuffix = reason && extractedTokens === 0 ? ` — ${PRUNE_REASON_LABELS[reason]}` : ""
     let message =
         formatStatsHeader(state.stats.totalPruneTokens, state.stats.pruneTokenCounter) +
-        reasonSuffix
+        reasonSuffix +
+        extractedSuffix
 
-    return message + formatExtracted(distillation)
+    return message + formatExtracted(showDistillation ? distillation : undefined)
 }
 
 function buildDetailedMessage(
@@ -34,21 +40,26 @@ function buildDetailedMessage(
     reason: PruneReason | undefined,
     pruneToolIds: string[],
     toolMetadata: Map<string, ToolParameterEntry>,
-    workingDirectory?: string,
-    distillation?: string[],
+    workingDirectory: string,
+    distillation: string[] | undefined,
+    showDistillation: boolean,
 ): string {
     let message = formatStatsHeader(state.stats.totalPruneTokens, state.stats.pruneTokenCounter)
 
     if (pruneToolIds.length > 0) {
         const pruneTokenCounterStr = `~${formatTokenCount(state.stats.pruneTokenCounter)}`
-        const reasonLabel = reason ? ` — ${PRUNE_REASON_LABELS[reason]}` : ""
-        message += `\n\n▣ Pruning (${pruneTokenCounterStr})${reasonLabel}`
+        const extractedTokens = countDistillationTokens(distillation)
+        const extractedSuffix =
+            extractedTokens > 0 ? `, extracted ${formatTokenCount(extractedTokens)}` : ""
+        const reasonLabel =
+            reason && extractedTokens === 0 ? ` — ${PRUNE_REASON_LABELS[reason]}` : ""
+        message += `\n\n▣ Pruning (${pruneTokenCounterStr}${extractedSuffix})${reasonLabel}`
 
         const itemLines = formatPrunedItemsList(pruneToolIds, toolMetadata, workingDirectory)
         message += "\n" + itemLines.join("\n")
     }
 
-    return (message + formatExtracted(distillation)).trim()
+    return (message + formatExtracted(showDistillation ? distillation : undefined)).trim()
 }
 
 export async function sendUnifiedNotification(
@@ -73,18 +84,19 @@ export async function sendUnifiedNotification(
         return false
     }
 
-    const showExtraction = config.tools.extract.showDistillation ? distillation : undefined
+    const showDistillation = config.tools.extract.showDistillation
 
     const message =
         config.pruneNotification === "minimal"
-            ? buildMinimalMessage(state, reason, showExtraction)
+            ? buildMinimalMessage(state, reason, distillation, showDistillation)
             : buildDetailedMessage(
                   state,
                   reason,
                   pruneToolIds,
                   toolMetadata,
                   workingDirectory,
-                  showExtraction,
+                  distillation,
+                  showDistillation,
               )
 
     await sendIgnoredMessage(client, sessionId, message, params, logger)
diff --git a/lib/ui/utils.ts b/lib/ui/utils.ts
index 9bee8d5b..b1e00ed9 100644
--- a/lib/ui/utils.ts
+++ b/lib/ui/utils.ts
@@ -1,5 +1,11 @@
 import { ToolParameterEntry } from "../state"
 import { extractParameterKey } from "../messages/utils"
+import { countTokens } from "../strategies/utils"
+
+export function countDistillationTokens(distillation?: string[]): number {
+    if (!distillation || distillation.length === 0) return 0
+    return countTokens(distillation.join("\n"))
+}
 
 export function formatExtracted(distillation?: string[]): string {
     if (!distillation || distillation.length === 0) {