fix: keep session token totals in sync (#1440) (thanks @robbyczgw-cla)

CHANGELOG.md

@@ -7,6 +7,7 @@ Docs: https://docs.clawd.bot
 ### Fixes
 - BlueBubbles: stop typing indicator on idle/no-reply. (#1439) Thanks @Nicell.
 - Auto-reply: only report a model switch when session state is available. (#1465) Thanks @robbyczgw-cla.
+- Auto-reply: keep cached context token count in sync after compaction. (#1440) Thanks @robbyczgw-cla.
 - Control UI: resolve local avatar URLs with basePath across injection + identity RPC. (#1457) Thanks @dlauer.
 - Agents: surface concrete API error details instead of generic AI service errors.
 - Docs: fix gog auth services example to include docs scope. (#1454) Thanks @zerone0x.

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,26 @@ export async function compactEmbeddedPiSession(params: {
               session.agent.replaceMessages(limited);
             }
             const result = await session.compact(params.customInstructions);
+            // Estimate tokens after compaction with the same context-usage heuristics.
+            let tokensAfter: number | undefined;
+            try {
+              const usage =
+                typeof session.getContextUsage === "function"
+                  ? session.getContextUsage()
+                  : undefined;
+              let estimate = usage?.tokens;
+              if (!Number.isFinite(estimate) || !estimate || estimate <= 0) {
+                estimate = 0;
+                for (const message of session.messages) {
+                  estimate += estimateTokens(message);
+                }
+              }
+              if (Number.isFinite(estimate) && estimate > 0 && estimate <= result.tokensBefore) {
+                tokensAfter = estimate;
+              }
+            } catch {
+              tokensAfter = undefined;
+            }
             return {
               ok: true,
               compacted: true,
@@ -377,6 +402,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/agents/subagent-registry.ts

@@ -31,7 +31,8 @@ const subagentRuns = new Map<string, SubagentRunRecord>();
 let sweeper: NodeJS.Timeout | null = null;
 let listenerStarted = false;
 let listenerStop: (() => void) | null = null;
-let restoreAttempted = false;
+// Use var to avoid TDZ on circular init paths that can call restoreSubagentRunsOnce early.
+var restoreAttempted = false;
 
 function persistSubagentRuns() {
   try {

src/auto-reply/reply.block-streaming.test.ts

@@ -42,7 +42,7 @@ describe("block streaming", () => {
   });
 
   async function waitForCalls(fn: () => number, calls: number) {
-    const deadline = Date.now() + 1500;
+    const deadline = Date.now() + 15000;
     while (fn() < calls) {
       if (Date.now() > deadline) {
         throw new Error(`Expected ${calls} call(s), got ${fn()}`);

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.test.ts

@@ -2,7 +2,7 @@ import { describe, expect, it, vi } from "vitest";
 
 import type { ClawdbotConfig } from "../../config/config.js";
 import { enqueueSystemEvent, resetSystemEventsForTest } from "../../infra/system-events.js";
-import { prependSystemEvents } from "./session-updates.js";
+import { incrementCompactionCount, prependSystemEvents } from "./session-updates.js";
 
 describe("prependSystemEvents", () => {
   it("adds a local timestamp to queued system events by default", async () => {
@@ -29,3 +29,37 @@ describe("prependSystemEvents", () => {
     vi.useRealTimers();
   });
 });
+
+describe("incrementCompactionCount", () => {
+  it("updates cached total tokens after compaction without clearing input/output", async () => {
+    const sessionKey = "agent:main:main";
+    const sessionStore = {
+      [sessionKey]: {
+        sessionId: "s1",
+        updatedAt: 10,
+        compactionCount: 1,
+        totalTokens: 9_000,
+        inputTokens: 111,
+        outputTokens: 222,
+      },
+    };
+    const now = 1234;
+
+    const nextCount = await incrementCompactionCount({
+      sessionEntry: sessionStore[sessionKey],
+      sessionStore,
+      sessionKey,
+      now,
+      tokensAfter: 2_000,
+    });
+
+    expect(nextCount).toBe(2);
+    expect(sessionStore[sessionKey]).toMatchObject({
+      compactionCount: 2,
+      totalTokens: 2_000,
+      inputTokens: 111,
+      outputTokens: 222,
+      updatedAt: now,
+    });
+  });
+});

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

@@ -237,23 +237,39 @@ export async function incrementCompactionCount(params: {
   sessionKey?: string;
   storePath?: string;
   now?: number;
+  /** Token count after compaction - if provided, updates cached context usage */
+  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 context usage.
+  const updates: Partial<SessionEntry> = {
     compactionCount: nextCount,
     updatedAt: now,
   };
+  // If tokensAfter is provided, update the cached total to reflect post-compaction context size.
+  if (tokensAfter != null && tokensAfter > 0) {
+    updates.totalTokens = tokensAfter;
+  }
+  sessionStore[sessionKey] = {
+    ...entry,
+    ...updates,
+  };
   if (storePath) {
     await updateSessionStore(storePath, (store) => {
       store[sessionKey] = {
         ...store[sessionKey],
-        compactionCount: nextCount,
-        updatedAt: now,
+        ...updates,
       };
     });
   }