🤖 fix: clear stale usage state when history is rewritten so the next send does not auto-compact

src/node/services/agentSession.ts

@@ -2733,6 +2733,7 @@ export class AgentSession {
       // when the edit target is outside the active context window.
       const truncateTargetId = await this.getEditTruncateTargetId(editMessageId);
 
+      this.clearUsageState();
       const truncateResult = await this.historyService.truncateAfterMessage(
         this.workspaceId,
         truncateTargetId
@@ -3418,6 +3419,11 @@ export class AgentSession {
     };
   }
 
+  /** Prevent cached usage from auto-compacting a rewritten context. */
+  clearUsageState(): void {
+    this.lastUsageState = undefined;
+  }
+
   /**
    * Persist a manual user message + emit a stream-error chat event when a
    * pre-stream gate (e.g. the unpriced-model budget gate) rejects a send.
@@ -4583,6 +4589,7 @@ export class AgentSession {
       });
     }
 
+    this.clearUsageState();
     const clearResult = this.clearHistoryForHardRestart
       ? await this.clearHistoryForHardRestart({
           monitorHistoryLockHeld: context.monitorHistoryLockHeld === true,
@@ -5196,7 +5203,7 @@ export class AgentSession {
 
           // Compaction collapses history to a boundary summary, so prior context-usage snapshots
           // are stale. Clear them to prevent immediate re-trigger loops on the follow-up turn.
-          this.lastUsageState = undefined;
+          this.clearUsageState();
 
           if (completedCompactionRequest?.source === "auto-compaction") {
             this.emitChatEvent({
@@ -6724,6 +6731,7 @@ export class AgentSession {
       pendingFollowUp: params.pendingFollowUp,
     });
     if (result.success) {
+      this.clearUsageState();
       this.onPostCompactionStateChange?.();
     }
     return result;

src/node/services/historyService.test.ts

@@ -1,12 +1,12 @@
 import { describe, it, expect, beforeEach, afterEach } from "bun:test";
 import { CONTEXT_BOUNDARY_KINDS } from "@/common/constants/contextBoundary";
 import { HistoryService } from "./historyService";
-import { Config } from "@/node/config";
+import type { Config } from "@/node/config";
+import { createTestHistoryService } from "./testHistoryService";
 import { createMuxMessage, type MuxMessage } from "@/common/types/message";
 import assert from "node:assert";
 import * as fs from "fs/promises";
 import * as path from "path";
-import * as os from "os";
 
 /** Collect all messages via iterateFullHistory (replaces removed getFullHistory). */
 async function collectFullHistory(service: HistoryService, workspaceId: string) {
@@ -48,25 +48,17 @@ async function appendNumberedMessages(
 describe("HistoryService", () => {
   let service: HistoryService;
   let config: Config;
-  let tempDir: string;
+  let cleanup: () => Promise<void>;
 
   beforeEach(async () => {
-    // Create a temporary directory for test files
-    tempDir = path.join(os.tmpdir(), `mux-test-${Date.now()}-${Math.random()}`);
-    await fs.mkdir(tempDir, { recursive: true });
-
-    // Create a Config with the temp directory
-    config = new Config(tempDir);
-    service = new HistoryService(config);
+    const testService = await createTestHistoryService();
+    service = testService.historyService;
+    config = testService.config;
+    cleanup = testService.cleanup;
   });
 
   afterEach(async () => {
-    // Clean up temp directory
-    try {
-      await fs.rm(tempDir, { recursive: true, force: true });
-    } catch {
-      // Ignore cleanup errors
-    }
+    await cleanup();
   });
 
   describe("getHistory", () => {
@@ -2037,6 +2029,67 @@ describe("HistoryService", () => {
       expect(await fs.readFile(archivePath(wsId), "utf-8")).toBe(archiveBefore);
     });
 
+    it("does not reseed usage from before a partial prefix truncation", async () => {
+      await appendNumberedMessages(service, wsId, 8);
+      await service.appendToHistory(
+        wsId,
+        createMuxMessage("assistant-usage", "assistant", "reply", {
+          contextUsage: { inputTokens: 95_000, outputTokens: 100, totalTokens: 95_100 },
+          contextProviderMetadata: { openai: {} },
+          model: "openai:gpt-4o",
+        })
+      );
+      await service.appendToHistory(
+        wsId,
+        createMuxMessage("assistant-provider-metadata", "assistant", "reply", {
+          contextProviderMetadata: { openai: {} },
+          model: "openai:gpt-4o",
+        })
+      );
+
+      const truncateResult = await service.truncateHistory(wsId, 0.5);
+      expect(truncateResult.success).toBe(true);
+
+      const restarted = new HistoryService(config);
+      const remaining = await restarted.getHistoryFromLatestBoundary(wsId);
+      expect(remaining.success).toBe(true);
+      if (remaining.success) {
+        const retainedAssistant = remaining.data.find(
+          (message) => message.id === "assistant-usage"
+        );
+        expect(retainedAssistant).toBeDefined();
+        expect(retainedAssistant?.metadata?.contextUsage).toBeUndefined();
+        expect(retainedAssistant?.metadata?.contextProviderMetadata).toBeUndefined();
+        const providerMetadataOnly = remaining.data.find(
+          (message) => message.id === "assistant-provider-metadata"
+        );
+        expect(providerMetadataOnly).toBeDefined();
+        expect(providerMetadataOnly?.metadata?.contextProviderMetadata).toBeUndefined();
+      }
+    });
+
+    it("preserves active usage when truncation removes only sealed rows", async () => {
+      await appendNumberedMessages(service, wsId, 8);
+      await service.appendToHistory(wsId, boundaryMessage("boundary-1", 1));
+      await service.appendToHistory(
+        wsId,
+        createMuxMessage("active-usage", "assistant", "reply", {
+          contextUsage: { inputTokens: 95_000, outputTokens: 100, totalTokens: 95_100 },
+          model: "openai:gpt-4o",
+        })
+      );
+
+      expect((await service.truncateHistory(wsId, 0.2)).success).toBe(true);
+
+      const active = await service.getHistoryFromLatestBoundary(wsId);
+      expect(active.success).toBe(true);
+      if (active.success) {
+        expect(
+          active.data.find((message) => message.id === "active-usage")?.metadata?.contextUsage
+        ).toBeDefined();
+      }
+    });
+
     it("hasHistory sees archive-only workspaces", async () => {
       await appendNumberedMessages(service, wsId, 1);
       await service.appendToHistory(wsId, boundaryMessage("boundary-1", 1));

src/node/services/historyService.ts

@@ -20,6 +20,8 @@ import { safeStringifyForCounting } from "@/common/utils/tokens/safeStringifyFor
 import { normalizeLegacyMuxMetadata } from "@/node/utils/messages/legacy";
 import { CONTEXT_BOUNDARY_KINDS } from "@/common/constants/contextBoundary";
 import {
+  findLatestContextBoundaryIndex,
+  getContextBoundaryKind,
   isDurableCompactedMarker,
   isDurableContextBoundaryMarker,
 } from "@/common/utils/messages/compactionBoundary";
@@ -41,6 +43,32 @@ function hasDurableCompactionBoundary(metadata: MuxMetadata | undefined): boolea
   return isPositiveInteger(metadata.compactionEpoch);
 }
 
+function prefixCutChangesActiveContext(messages: MuxMessage[], removeCount: number): boolean {
+  const boundaryIndex = findLatestContextBoundaryIndex(messages);
+  if (boundaryIndex < 0) {
+    return removeCount > 0;
+  }
+  const activeStart =
+    getContextBoundaryKind(messages[boundaryIndex]) === CONTEXT_BOUNDARY_KINDS.RESET
+      ? boundaryIndex + 1
+      : boundaryIndex;
+  return removeCount > activeStart;
+}
+
+function stripContextUsage(message: MuxMessage): MuxMessage {
+  if (!message.metadata) {
+    return message;
+  }
+  return {
+    ...message,
+    metadata: {
+      ...message.metadata,
+      contextUsage: undefined,
+      contextProviderMetadata: undefined,
+    },
+  };
+}
+
 function getCompactionMetadataToPreserve(
   workspaceId: string,
   existingMessage: MuxMessage,
@@ -1800,6 +1828,8 @@ export class HistoryService {
         const historyPath = this.getChatHistoryPath(workspaceId);
         const historyEntries = this.serializeHistoryEntries(truncatedMessages, workspaceId);
 
+        const archiveMaxSeq = await this.getArchiveTailMaxSequence(workspaceId);
+
         // Atomic write prevents corruption if app crashes mid-write
         await writeFileAtomic(historyPath, historyEntries);
 
@@ -1829,7 +1859,6 @@ export class HistoryService {
         // truncation. When the truncation empties the active file, floor the
         // counter with the archive max so new appends can never reuse archived
         // sequence numbers.
-        const archiveMaxSeq = await this.getArchiveTailMaxSequence(workspaceId);
         const nextSeq = Math.max(maxTruncatedSeq, archiveMaxSeq) + 1;
         assert(
           isNonNegativeInteger(nextSeq),
@@ -1869,11 +1898,25 @@ export class HistoryService {
         keepTargetMessage ? messageIndex + 1 : messageIndex
       );
 
+      const historyPath = this.getChatHistoryPath(workspaceId);
+      const archivePath = this.getChatArchivePath(workspaceId);
+      await fs.rm(historyPath, { force: true });
       await writeFileAtomic(
-        this.getChatHistoryPath(workspaceId),
-        this.serializeHistoryEntries(truncatedMessages, workspaceId)
+        archivePath,
+        this.serializeHistoryEntries(truncatedMessages.map(stripContextUsage), workspaceId)
       );
-      await fs.rm(this.getChatArchivePath(workspaceId), { force: true });
+      await fs.rename(archivePath, historyPath);
+      try {
+        await writeFileAtomic(
+          historyPath,
+          this.serializeHistoryEntries(truncatedMessages, workspaceId)
+        );
+      } catch (error) {
+        log.warn("Failed to restore usage after archived history truncation", {
+          workspaceId,
+          error,
+        });
+      }
       // chat.jsonl may contain sealed epochs again — allow the lazy check to re-run.
       this.sealedRotationChecked.delete(workspaceId);
 
@@ -1950,10 +1993,9 @@ export class HistoryService {
         // Structural rewrite requires full history content (oldest rows live in
         // the sealed archive). Percentage truncation is a rare recovery path
         // (compaction-failure retry), so the O(total-history) read is acceptable.
-        const messages = [
-          ...(await this.readArchivedHistory(workspaceId)),
-          ...(await this.readChatHistory(workspaceId)),
-        ];
+        const archivedMessages = await this.readArchivedHistory(workspaceId);
+        const chatMessages = await this.readChatHistory(workspaceId);
+        const messages = [...archivedMessages, ...chatMessages];
         if (messages.length === 0) {
           return Ok([]); // Nothing to truncate
         }
@@ -2003,21 +2045,37 @@ export class HistoryService {
           return Ok(deletedSequences);
         }
 
-        // Keep messages after removeCount
-        const remainingMessages = messages.slice(removeCount);
+        const activeContextChanged = prefixCutChangesActiveContext(messages, removeCount);
+        const sanitize = activeContextChanged
+          ? stripContextUsage
+          : (message: MuxMessage) => message;
+        const remainingMessages = messages.slice(removeCount).map(sanitize);
         const deletedMessages = messages.slice(0, removeCount);
         const deletedSequences = deletedMessages
           .map((msg) => msg.metadata?.historySequence)
           .filter((s): s is number => isNonNegativeInteger(s));
+        const remainingArchiveCount = Math.max(0, archivedMessages.length - removeCount);
+        const remainingArchive = remainingMessages.slice(0, remainingArchiveCount);
+        const remainingChat = remainingMessages.slice(remainingArchiveCount);
 
-        // Collapse the remainder into chat.jsonl and drop the archive (the cut
-        // may fall anywhere inside it). It may contain old boundaries; a later
-        // boundary write re-seals it.
-        const historyEntries = this.serializeHistoryEntries(remainingMessages, workspaceId);
-
-        // Atomic write prevents corruption if app crashes mid-write
-        await writeFileAtomic(historyPath, historyEntries);
-        await fs.rm(archivePath, { force: true });
+        if (activeContextChanged && archivedMessages.length > 0) {
+          await writeFileAtomic(
+            historyPath,
+            this.serializeHistoryEntries(chatMessages.map(stripContextUsage), workspaceId)
+          );
+        }
+        if (remainingArchive.length > 0) {
+          await writeFileAtomic(
+            archivePath,
+            this.serializeHistoryEntries(remainingArchive, workspaceId)
+          );
+        } else {
+          await fs.rm(archivePath, { force: true });
+        }
+        await writeFileAtomic(
+          historyPath,
+          this.serializeHistoryEntries(remainingChat, workspaceId)
+        );
         this.sealedRotationChecked.delete(workspaceId);
 
         // Update sequence counter to continue from where we are.