refactor: use pre-computed token counts in prune state instead of re-scanning messages

Tarquinen · Tarquinen · commit d7a4ad2242aa · 2026-02-09T20:37:41.000-05:00
diff --git a/lib/commands/context.ts b/lib/commands/context.ts
@@ -74,8 +74,8 @@ function analyzeTokens(state: SessionState, messages: WithParts[]): TokenBreakdo
         tools: 0,
         toolCount: 0,
         prunedTokens: state.stats.totalPruneTokens,
-        prunedToolCount: state.prune.toolIds.size,
-        prunedMessageCount: state.prune.messageIds.size,
+        prunedToolCount: state.prune.tools.size,
+        prunedMessageCount: state.prune.messages.size,
         total: 0,
     }
 
@@ -129,7 +129,7 @@ function analyzeTokens(state: SessionState, messages: WithParts[]): TokenBreakdo
                     foundToolIds.add(toolPart.callID)
                 }
 
-                const isPruned = toolPart.callID && state.prune.toolIds.has(toolPart.callID)
+                const isPruned = toolPart.callID && state.prune.tools.has(toolPart.callID)
                 if (!isCompacted && !isPruned) {
                     if (toolPart.state?.input) {
                         const inputStr =
diff --git a/lib/commands/stats.ts b/lib/commands/stats.ts
@@ -51,8 +51,8 @@ export async function handleStatsCommand(ctx: StatsCommandContext): Promise<void
 
     // Session stats from in-memory state
     const sessionTokens = state.stats.totalPruneTokens
-    const sessionTools = state.prune.toolIds.size
-    const sessionMessages = state.prune.messageIds.size
+    const sessionTools = state.prune.tools.size
+    const sessionMessages = state.prune.messages.size
 
     // All-time stats from storage files
     const allTime = await loadAllSessionStats(logger)
diff --git a/lib/commands/sweep.ts b/lib/commands/sweep.ts
@@ -12,7 +12,7 @@ import type { SessionState, WithParts, ToolParameterEntry } from "../state"
 import type { PluginConfig } from "../config"
 import { sendIgnoredMessage } from "../ui/notification"
 import { formatPrunedItemsList } from "../ui/utils"
-import { getCurrentParams, calculateTokensSaved } from "../strategies/utils"
+import { getCurrentParams, getTotalToolTokens } from "../strategies/utils"
 import { buildToolIdList, isIgnoredUserMessage } from "../messages/utils"
 import { saveSessionState } from "../state/persistence"
 import { isMessageCompacted } from "../shared-utils"
@@ -164,7 +164,7 @@ export async function handleSweepCommand(ctx: SweepCommandContext): Promise<void
 
     // Filter out already-pruned tools, protected tools, and protected file paths
     const newToolIds = toolIdsToSweep.filter((id) => {
-        if (state.prune.toolIds.has(id)) {
+        if (state.prune.tools.has(id)) {
             return false
         }
         const entry = state.toolParameters.get(id)
@@ -214,13 +214,13 @@ export async function handleSweepCommand(ctx: SweepCommandContext): Promise<void
         return
     }
 
+    const tokensSaved = getTotalToolTokens(state, newToolIds)
+
     // Add to prune list
     for (const id of newToolIds) {
-        state.prune.toolIds.add(id)
+        const entry = state.toolParameters.get(id)
+        state.prune.tools.set(id, entry?.tokenCount ?? 0)
     }
-
-    // Calculate tokens saved
-    const tokensSaved = calculateTokensSaved(state, messages, newToolIds)
     state.stats.pruneTokenCounter += tokensSaved
     state.stats.totalPruneTokens += state.stats.pruneTokenCounter
     state.stats.pruneTokenCounter = 0
diff --git a/lib/messages/inject.ts b/lib/messages/inject.ts
@@ -162,7 +162,7 @@ const buildPrunableToolsList = (
     const toolIdList = state.toolIdList
 
     state.toolParameters.forEach((toolParameterEntry, toolCallId) => {
-        if (state.prune.toolIds.has(toolCallId)) {
+        if (state.prune.tools.has(toolCallId)) {
             return
         }
 
diff --git a/lib/messages/prune.ts b/lib/messages/prune.ts
@@ -38,7 +38,7 @@ const pruneFullTool = (state: SessionState, logger: Logger, messages: WithParts[
             if (part.type !== "tool") {
                 continue
             }
-            if (!state.prune.toolIds.has(part.callID)) {
+            if (!state.prune.tools.has(part.callID)) {
                 continue
             }
             if (part.tool !== "edit" && part.tool !== "write") {
@@ -79,7 +79,7 @@ const pruneToolOutputs = (state: SessionState, logger: Logger, messages: WithPar
             if (part.type !== "tool") {
                 continue
             }
-            if (!state.prune.toolIds.has(part.callID)) {
+            if (!state.prune.tools.has(part.callID)) {
                 continue
             }
             if (part.state.status !== "completed") {
@@ -105,7 +105,7 @@ const pruneToolInputs = (state: SessionState, logger: Logger, messages: WithPart
             if (part.type !== "tool") {
                 continue
             }
-            if (!state.prune.toolIds.has(part.callID)) {
+            if (!state.prune.tools.has(part.callID)) {
                 continue
             }
             if (part.state.status !== "completed") {
@@ -133,7 +133,7 @@ const pruneToolErrors = (state: SessionState, logger: Logger, messages: WithPart
             if (part.type !== "tool") {
                 continue
             }
-            if (!state.prune.toolIds.has(part.callID)) {
+            if (!state.prune.tools.has(part.callID)) {
                 continue
             }
             if (part.state.status !== "error") {
@@ -158,7 +158,7 @@ const filterCompressedRanges = (
     logger: Logger,
     messages: WithParts[],
 ): void => {
-    if (!state.prune.messageIds?.size) {
+    if (!state.prune.messages?.size) {
         return
     }
 
@@ -193,7 +193,7 @@ const filterCompressedRanges = (
         }
 
         // Skip messages that are in the prune list
-        if (state.prune.messageIds.has(msgId)) {
+        if (state.prune.messages.has(msgId)) {
             continue
         }
 
diff --git a/lib/shared-utils.ts b/lib/shared-utils.ts
@@ -5,7 +5,7 @@ export const isMessageCompacted = (state: SessionState, msg: WithParts): boolean
     if (msg.info.time.created < state.lastCompaction) {
         return true
     }
-    if (state.prune.messageIds.has(msg.info.id)) {
+    if (state.prune.messages.has(msg.info.id)) {
         return true
     }
     return false
diff --git a/lib/state/persistence.ts b/lib/state/persistence.ts
@@ -11,10 +11,14 @@ import { join } from "path"
 import type { SessionState, SessionStats, CompressSummary } from "./types"
 import type { Logger } from "../logger"
 
-/** Prune state as stored on disk (arrays for JSON compatibility) */
+/** Prune state as stored on disk */
 export interface PersistedPrune {
-    toolIds: string[]
-    messageIds: string[]
+    // New format: tool/message IDs with token counts
+    tools?: Record<string, number>
+    messages?: Record<string, number>
+    // Legacy format: plain ID arrays (backward compatibility)
+    toolIds?: string[]
+    messageIds?: string[]
 }
 
 export interface PersistedSessionState {
@@ -58,8 +62,8 @@ export async function saveSessionState(
         const state: PersistedSessionState = {
             sessionName: sessionName,
             prune: {
-                toolIds: [...sessionState.prune.toolIds],
-                messageIds: [...sessionState.prune.messageIds],
+                tools: Object.fromEntries(sessionState.prune.tools),
+                messages: Object.fromEntries(sessionState.prune.messages),
             },
             compressSummaries: sessionState.compressSummaries,
             stats: sessionState.stats,
@@ -96,7 +100,9 @@ export async function loadSessionState(
         const content = await fs.readFile(filePath, "utf-8")
         const state = JSON.parse(content) as PersistedSessionState
 
-        if (!state || !state.prune || !Array.isArray(state.prune.toolIds) || !state.stats) {
+        const hasNewFormat = state?.prune?.tools && typeof state.prune.tools === "object"
+        const hasLegacyFormat = Array.isArray(state?.prune?.toolIds)
+        if (!state || !state.prune || (!hasNewFormat && !hasLegacyFormat) || !state.stats) {
             logger.warn("Invalid session state file, ignoring", {
                 sessionId: sessionId,
             })
@@ -166,10 +172,14 @@ export async function loadAllSessionStats(logger: Logger): Promise<AggregatedSta
                 const content = await fs.readFile(filePath, "utf-8")
                 const state = JSON.parse(content) as PersistedSessionState
 
-                if (state?.stats?.totalPruneTokens && state?.prune?.toolIds) {
+                if (state?.stats?.totalPruneTokens && state?.prune) {
                     result.totalTokens += state.stats.totalPruneTokens
-                    result.totalTools += state.prune.toolIds.length
-                    result.totalMessages += state.prune.messageIds?.length || 0
+                    result.totalTools += state.prune.tools
+                        ? Object.keys(state.prune.tools).length
+                        : (state.prune.toolIds?.length ?? 0)
+                    result.totalMessages += state.prune.messages
+                        ? Object.keys(state.prune.messages).length
+                        : (state.prune.messageIds?.length ?? 0)
                     result.sessionCount++
                 }
             } catch {
diff --git a/lib/state/state.ts b/lib/state/state.ts
@@ -6,6 +6,7 @@ import {
     findLastCompactionTimestamp,
     countTurns,
     resetOnCompaction,
+    loadPruneMap,
 } from "./utils"
 import { getLastUserMessage } from "../shared-utils"
 
@@ -48,8 +49,8 @@ export function createSessionState(): SessionState {
         sessionId: null,
         isSubAgent: false,
         prune: {
-            toolIds: new Set<string>(),
-            messageIds: new Set<string>(),
+            tools: new Map<string, number>(),
+            messages: new Map<string, number>(),
         },
         compressSummaries: [],
         stats: {
@@ -71,8 +72,8 @@ export function resetSessionState(state: SessionState): void {
     state.sessionId = null
     state.isSubAgent = false
     state.prune = {
-        toolIds: new Set<string>(),
-        messageIds: new Set<string>(),
+        tools: new Map<string, number>(),
+        messages: new Map<string, number>(),
     }
     state.compressSummaries = []
     state.stats = {
@@ -118,10 +119,8 @@ export async function ensureSessionInitialized(
         return
     }
 
-    state.prune = {
-        toolIds: new Set(persisted.prune.toolIds || []),
-        messageIds: new Set(persisted.prune.messageIds || []),
-    }
+    state.prune.tools = loadPruneMap(persisted.prune.tools, persisted.prune.toolIds)
+    state.prune.messages = loadPruneMap(persisted.prune.messages, persisted.prune.messageIds)
     state.compressSummaries = persisted.compressSummaries || []
     state.stats = {
         pruneTokenCounter: persisted.stats?.pruneTokenCounter || 0,
diff --git a/lib/state/types.ts b/lib/state/types.ts
@@ -27,8 +27,8 @@ export interface CompressSummary {
 }
 
 export interface Prune {
-    toolIds: Set<string>
-    messageIds: Set<string>
+    tools: Map<string, number>
+    messages: Map<string, number>
 }
 
 export interface SessionState {
diff --git a/lib/state/utils.ts b/lib/state/utils.ts
@@ -36,10 +36,19 @@ export function countTurns(state: SessionState, messages: WithParts[]): number {
     return turnCount
 }
 
+export function loadPruneMap(
+    obj?: Record<string, number>,
+    legacyArr?: string[],
+): Map<string, number> {
+    if (obj) return new Map(Object.entries(obj))
+    if (legacyArr) return new Map(legacyArr.map((id) => [id, 0]))
+    return new Map()
+}
+
 export function resetOnCompaction(state: SessionState): void {
     state.toolParameters.clear()
-    state.prune.toolIds = new Set<string>()
-    state.prune.messageIds = new Set<string>()
+    state.prune.tools = new Map<string, number>()
+    state.prune.messages = new Map<string, number>()
     state.compressSummaries = []
     state.nudgeCounter = 0
     state.lastToolPrune = false
diff --git a/lib/strategies/deduplication.ts b/lib/strategies/deduplication.ts
@@ -2,7 +2,7 @@ import { PluginConfig } from "../config"
 import { Logger } from "../logger"
 import type { SessionState, WithParts } from "../state"
 import { getFilePathsFromParameters, isProtected } from "../protected-file-patterns"
-import { calculateTokensSaved } from "./utils"
+import { getTotalToolTokens } from "./utils"
 
 /**
  * Deduplication strategy - prunes older tool calls that have identical
@@ -25,7 +25,7 @@ export const deduplicate = (
     }
 
     // Filter out IDs already pruned
-    const unprunedIds = allToolIds.filter((id) => !state.prune.toolIds.has(id))
+    const unprunedIds = allToolIds.filter((id) => !state.prune.tools.has(id))
 
     if (unprunedIds.length === 0) {
         return
@@ -74,11 +74,12 @@ export const deduplicate = (
         }
     }
 
-    state.stats.totalPruneTokens += calculateTokensSaved(state, messages, newPruneIds)
+    state.stats.totalPruneTokens += getTotalToolTokens(state, newPruneIds)
 
     if (newPruneIds.length > 0) {
         for (const id of newPruneIds) {
-            state.prune.toolIds.add(id)
+            const entry = state.toolParameters.get(id)
+            state.prune.tools.set(id, entry?.tokenCount ?? 0)
         }
         logger.debug(`Marked ${newPruneIds.length} duplicate tool calls for pruning`)
     }
diff --git a/lib/strategies/purge-errors.ts b/lib/strategies/purge-errors.ts
@@ -2,7 +2,7 @@ import { PluginConfig } from "../config"
 import { Logger } from "../logger"
 import type { SessionState, WithParts } from "../state"
 import { getFilePathsFromParameters, isProtected } from "../protected-file-patterns"
-import { calculateTokensSaved } from "./utils"
+import { getTotalToolTokens } from "./utils"
 
 /**
  * Purge Errors strategy - prunes tool inputs for tools that errored
@@ -28,7 +28,7 @@ export const purgeErrors = (
     }
 
     // Filter out IDs already pruned
-    const unprunedIds = allToolIds.filter((id) => !state.prune.toolIds.has(id))
+    const unprunedIds = allToolIds.filter((id) => !state.prune.tools.has(id))
 
     if (unprunedIds.length === 0) {
         return
@@ -68,9 +68,10 @@ export const purgeErrors = (
     }
 
     if (newPruneIds.length > 0) {
-        state.stats.totalPruneTokens += calculateTokensSaved(state, messages, newPruneIds)
+        state.stats.totalPruneTokens += getTotalToolTokens(state, newPruneIds)
         for (const id of newPruneIds) {
-            state.prune.toolIds.add(id)
+            const entry = state.toolParameters.get(id)
+            state.prune.tools.set(id, entry?.tokenCount ?? 0)
         }
         logger.debug(
             `Marked ${newPruneIds.length} error tool calls for pruning (older than ${turnThreshold} turns)`,
diff --git a/lib/strategies/supersede-writes.ts b/lib/strategies/supersede-writes.ts
@@ -2,7 +2,7 @@ import { PluginConfig } from "../config"
 import { Logger } from "../logger"
 import type { SessionState, WithParts } from "../state"
 import { getFilePathsFromParameters, isProtected } from "../protected-file-patterns"
-import { calculateTokensSaved } from "./utils"
+import { getTotalToolTokens } from "./utils"
 
 /**
  * Supersede Writes strategy - prunes write tool inputs for files that have
@@ -28,7 +28,7 @@ export const supersedeWrites = (
     }
 
     // Filter out IDs already pruned
-    const unprunedIds = allToolIds.filter((id) => !state.prune.toolIds.has(id))
+    const unprunedIds = allToolIds.filter((id) => !state.prune.tools.has(id))
     if (unprunedIds.length === 0) {
         return
     }
@@ -88,7 +88,7 @@ export const supersedeWrites = (
         // For each write, check if there's a read that comes after it
         for (const write of writes) {
             // Skip if already pruned
-            if (state.prune.toolIds.has(write.id)) {
+            if (state.prune.tools.has(write.id)) {
                 continue
             }
 
@@ -101,9 +101,10 @@ export const supersedeWrites = (
     }
 
     if (newPruneIds.length > 0) {
-        state.stats.totalPruneTokens += calculateTokensSaved(state, messages, newPruneIds)
+        state.stats.totalPruneTokens += getTotalToolTokens(state, newPruneIds)
         for (const id of newPruneIds) {
-            state.prune.toolIds.add(id)
+            const entry = state.toolParameters.get(id)
+            state.prune.tools.set(id, entry?.tokenCount ?? 0)
         }
         logger.debug(`Marked ${newPruneIds.length} superseded write tool calls for pruning`)
     }
diff --git a/lib/strategies/utils.ts b/lib/strategies/utils.ts
diff --git a/lib/tools/compress.ts b/lib/tools/compress.ts
diff --git a/lib/tools/prune-shared.ts b/lib/tools/prune-shared.ts
diff --git a/lib/tools/utils.ts b/lib/tools/utils.ts

Original file line number	Diff line number	Diff line change
`@@ -162,7 +162,7 @@ const buildPrunableToolsList = (`
`162`	`162`	`const toolIdList = state.toolIdList`
`163`	`163`
`164`	`164`	`state.toolParameters.forEach((toolParameterEntry, toolCallId) => {`
`165`		`- if (state.prune.toolIds.has(toolCallId)) {`
	`165`	`+ if (state.prune.tools.has(toolCallId)) {`
`166`	`166`	`return`
`167`	`167`	`}`
`168`	`168`
Original file line number	Diff line number	Diff line change
`@@ -38,7 +38,7 @@ const pruneFullTool = (state: SessionState, logger: Logger, messages: WithParts[`
`38`	`38`	`if (part.type !== "tool") {`
`39`	`39`	`continue`
`40`	`40`	`}`
`41`		`- if (!state.prune.toolIds.has(part.callID)) {`
	`41`	`+ if (!state.prune.tools.has(part.callID)) {`
`42`	`42`	`continue`
`43`	`43`	`}`
`44`	`44`	`if (part.tool !== "edit" && part.tool !== "write") {`
`@@ -79,7 +79,7 @@ const pruneToolOutputs = (state: SessionState, logger: Logger, messages: WithPar`
`79`	`79`	`if (part.type !== "tool") {`
`80`	`80`	`continue`
`81`	`81`	`}`
`82`		`- if (!state.prune.toolIds.has(part.callID)) {`
	`82`	`+ if (!state.prune.tools.has(part.callID)) {`
`83`	`83`	`continue`
`84`	`84`	`}`
`85`	`85`	`if (part.state.status !== "completed") {`
`@@ -105,7 +105,7 @@ const pruneToolInputs = (state: SessionState, logger: Logger, messages: WithPart`
`105`	`105`	`if (part.type !== "tool") {`
`106`	`106`	`continue`
`107`	`107`	`}`
`108`		`- if (!state.prune.toolIds.has(part.callID)) {`
	`108`	`+ if (!state.prune.tools.has(part.callID)) {`
`109`	`109`	`continue`
`110`	`110`	`}`
`111`	`111`	`if (part.state.status !== "completed") {`
`@@ -133,7 +133,7 @@ const pruneToolErrors = (state: SessionState, logger: Logger, messages: WithPart`
`133`	`133`	`if (part.type !== "tool") {`
`134`	`134`	`continue`
`135`	`135`	`}`
`136`		`- if (!state.prune.toolIds.has(part.callID)) {`
	`136`	`+ if (!state.prune.tools.has(part.callID)) {`
`137`	`137`	`continue`
`138`	`138`	`}`
`139`	`139`	`if (part.state.status !== "error") {`
`@@ -158,7 +158,7 @@ const filterCompressedRanges = (`
`158`	`158`	`logger: Logger,`
`159`	`159`	`messages: WithParts[],`
`160`	`160`	`): void => {`
`161`		`- if (!state.prune.messageIds?.size) {`
	`161`	`+ if (!state.prune.messages?.size) {`
`162`	`162`	`return`
`163`	`163`	`}`
`164`	`164`
`@@ -193,7 +193,7 @@ const filterCompressedRanges = (`
`193`	`193`	`}`
`194`	`194`
`195`	`195`	`// Skip messages that are in the prune list`
`196`		`- if (state.prune.messageIds.has(msgId)) {`
	`196`	`+ if (state.prune.messages.has(msgId)) {`
`197`	`197`	`continue`
`198`	`198`	`}`
`199`	`199`
Original file line number	Diff line number	Diff line change
`@@ -5,7 +5,7 @@ export const isMessageCompacted = (state: SessionState, msg: WithParts): boolean`
`5`	`5`	`if (msg.info.time.created < state.lastCompaction) {`
`6`	`6`	`return true`
`7`	`7`	`}`
`8`		`- if (state.prune.messageIds.has(msg.info.id)) {`
	`8`	`+ if (state.prune.messages.has(msg.info.id)) {`
`9`	`9`	`return true`
`10`	`10`	`}`
`11`	`11`	`return false`
Original file line number	Diff line number	Diff line change
`@@ -27,8 +27,8 @@ export interface CompressSummary {`
`27`	`27`	`}`
`28`	`28`
`29`	`29`	`export interface Prune {`
`30`		`- toolIds: Set<string>`
`31`		`- messageIds: Set<string>`
	`30`	`+ tools: Map<string, number>`
	`31`	`+ messages: Map<string, number>`
`32`	`32`	`}`
`33`	`33`
`34`	`34`	`export interface SessionState {`