Merge pull request #1440 from robbyczgw-cla/fix/token-count-after-compaction

fix: update token count display after compaction

src/agents/pi-embedded-runner/compact.ts

@@ -1,7 +1,12 @@
 import fs from "node:fs/promises";
 import os from "node:os";
 
-import { createAgentSession, SessionManager, SettingsManager } from "@mariozechner/pi-coding-agent";
+import {
+  createAgentSession,
+  estimateTokens,
+  SessionManager,
+  SettingsManager,
+} from "@mariozechner/pi-coding-agent";
 
 import { resolveHeartbeatPrompt } from "../../auto-reply/heartbeat.js";
 import type { ReasoningLevel, ThinkLevel } from "../../auto-reply/thinking.js";
@@ -370,6 +375,21 @@ export async function compactEmbeddedPiSession(params: {
               session.agent.replaceMessages(limited);
             }
             const result = await session.compact(params.customInstructions);
+            // Estimate tokens after compaction by summing token estimates for remaining messages
+            let tokensAfter: number | undefined;
+            try {
+              tokensAfter = 0;
+              for (const message of session.messages) {
+                tokensAfter += estimateTokens(message);
+              }
+              // Sanity check: tokensAfter should be less than tokensBefore
+              if (tokensAfter > result.tokensBefore) {
+                tokensAfter = undefined; // Don't trust the estimate
+              }
+            } catch {
+              // If estimation fails, leave tokensAfter undefined
+              tokensAfter = undefined;
+            }
             return {
               ok: true,
               compacted: true,
@@ -377,6 +397,7 @@ export async function compactEmbeddedPiSession(params: {
                 summary: result.summary,
                 firstKeptEntryId: result.firstKeptEntryId,
                 tokensBefore: result.tokensBefore,
+                tokensAfter,
                 details: result.details,
               },
             };

src/agents/pi-embedded-runner/types.ts

@@ -59,6 +59,7 @@ export type EmbeddedPiCompactResult = {
     summary: string;
     firstKeptEntryId: string;
     tokensBefore: number;
+    tokensAfter?: number;
     details?: unknown;
   };
 };

src/auto-reply/reply/commands-compact.ts

@@ -83,18 +83,13 @@ export const handleCompactCommand: CommandHandler = async (params) => {
     ownerNumbers: params.command.ownerList.length > 0 ? params.command.ownerList : undefined,
   });
 
-  const totalTokens =
-    params.sessionEntry.totalTokens ??
-    (params.sessionEntry.inputTokens ?? 0) + (params.sessionEntry.outputTokens ?? 0);
-  const contextSummary = formatContextUsageShort(
-    totalTokens > 0 ? totalTokens : null,
-    params.contextTokens ?? params.sessionEntry.contextTokens ?? null,
-  );
   const compactLabel = result.ok
     ? result.compacted
-      ? result.result?.tokensBefore
-        ? `Compacted (${formatTokenCount(result.result.tokensBefore)} before)`
-        : "Compacted"
+      ? result.result?.tokensBefore != null && result.result?.tokensAfter != null
+        ? `Compacted (${formatTokenCount(result.result.tokensBefore)} → ${formatTokenCount(result.result.tokensAfter)})`
+        : result.result?.tokensBefore
+          ? `Compacted (${formatTokenCount(result.result.tokensBefore)} before)`
+          : "Compacted"
       : "Compaction skipped"
     : "Compaction failed";
   if (result.ok && result.compacted) {
@@ -103,8 +98,20 @@ export const handleCompactCommand: CommandHandler = async (params) => {
       sessionStore: params.sessionStore,
       sessionKey: params.sessionKey,
       storePath: params.storePath,
+      // Update token counts after compaction
+      tokensAfter: result.result?.tokensAfter,
     });
   }
+  // Use the post-compaction token count for context summary if available
+  const tokensAfterCompaction = result.result?.tokensAfter;
+  const totalTokens =
+    tokensAfterCompaction ??
+    params.sessionEntry.totalTokens ??
+    (params.sessionEntry.inputTokens ?? 0) + (params.sessionEntry.outputTokens ?? 0);
+  const contextSummary = formatContextUsageShort(
+    totalTokens > 0 ? totalTokens : null,
+    params.contextTokens ?? params.sessionEntry.contextTokens ?? null,
+  );
   const reason = result.reason?.trim();
   const line = reason
     ? `${compactLabel}: ${reason} • ${contextSummary}`

src/auto-reply/reply/session-updates.ts

@@ -237,23 +237,42 @@ export async function incrementCompactionCount(params: {
   sessionKey?: string;
   storePath?: string;
   now?: number;
+  /** Token count after compaction - if provided, updates session token counts */
+  tokensAfter?: number;
 }): Promise<number | undefined> {
-  const { sessionEntry, sessionStore, sessionKey, storePath, now = Date.now() } = params;
+  const {
+    sessionEntry,
+    sessionStore,
+    sessionKey,
+    storePath,
+    now = Date.now(),
+    tokensAfter,
+  } = params;
   if (!sessionStore || !sessionKey) return undefined;
   const entry = sessionStore[sessionKey] ?? sessionEntry;
   if (!entry) return undefined;
   const nextCount = (entry.compactionCount ?? 0) + 1;
-  sessionStore[sessionKey] = {
-    ...entry,
+  // Build update payload with compaction count and optionally updated token counts
+  const updates: Partial<SessionEntry> = {
     compactionCount: nextCount,
     updatedAt: now,
   };
+  // If tokensAfter is provided, update the cached token counts to reflect post-compaction state
+  if (tokensAfter != null && tokensAfter > 0) {
+    updates.totalTokens = tokensAfter;
+    // Clear input/output breakdown since we only have the total estimate after compaction
+    updates.inputTokens = undefined;
+    updates.outputTokens = undefined;
+  }
+  sessionStore[sessionKey] = {
+    ...entry,
+    ...updates,
+  };
   if (storePath) {
     await updateSessionStore(storePath, (store) => {
       store[sessionKey] = {
         ...store[sessionKey],
-        compactionCount: nextCount,
-        updatedAt: now,
+        ...updates,
       };
     });
   }