Merge pull request #113 from Tarquinen/fix/tool-count-on-restart

Fix tool result counting on session restart

Author: Tarquinen

lib/core/janitor.ts

@@ -2,7 +2,7 @@ import { z } from "zod"
 import type { Logger } from "../logger"
 import type { PruningStrategy } from "../config"
 import type { PluginState } from "../state"
-import type { ToolMetadata, PruneReason, SessionStats, GCStats, PruningResult } from "../fetch-wrapper/types"
+import type { ToolMetadata, SessionStats, GCStats, PruningResult } from "../fetch-wrapper/types"
 import { findCurrentAgent } from "../hooks"
 import { buildAnalysisPrompt } from "./prompt"
 import { selectModel, extractModelFromSession } from "../model-selector"

lib/fetch-wrapper/tool-tracker.ts

@@ -11,3 +11,9 @@ export function createToolTracker(): ToolTracker {
 export function resetToolTrackerCount(tracker: ToolTracker): void {
     tracker.toolResultCount = 0
 }
+
+export function clearToolTracker(tracker: ToolTracker): void {
+    tracker.seenToolResultIds.clear()
+    tracker.toolResultCount = 0
+    tracker.skipNextIdle = false
+}

lib/hooks.ts

@@ -4,7 +4,7 @@ import type { JanitorContext } from "./core/janitor"
 import { runOnIdle } from "./core/janitor"
 import type { PluginConfig, PruningStrategy } from "./config"
 import type { ToolTracker } from "./fetch-wrapper/tool-tracker"
-import { resetToolTrackerCount } from "./fetch-wrapper/tool-tracker"
+import { resetToolTrackerCount, clearToolTracker } from "./fetch-wrapper/tool-tracker"
 import { clearAllMappings } from "./state/id-mapping"
 
 export async function isSubagentSession(client: any, sessionID: string): Promise<boolean> {
@@ -80,7 +80,7 @@ export function createChatParamsHandler(
             clearAllMappings()
             state.toolParameters.clear()
             if (toolTracker) {
-                resetToolTrackerCount(toolTracker)
+                clearToolTracker(toolTracker)
             }
         }
 

lib/state/tool-cache.ts

@@ -30,6 +30,10 @@ export async function syncToolCache(
         }
 
         let synced = 0
+        // Build lowercase set of pruned IDs for comparison (IDs in state may be mixed case)
+        const prunedIdsLower = tracker
+            ? new Set((state.prunedIds.get(sessionId) ?? []).map(id => id.toLowerCase()))
+            : null
 
         for (const msg of messages) {
             if (!msg.parts) continue
@@ -43,7 +47,8 @@ export async function syncToolCache(
                 if (tracker && !tracker.seenToolResultIds.has(id)) {
                     tracker.seenToolResultIds.add(id)
                     // Only count non-protected tools toward nudge threshold
-                    if (!part.tool || !protectedTools?.has(part.tool)) {
+                    // Also skip already-pruned tools to avoid re-counting on restart
+                    if ((!part.tool || !protectedTools?.has(part.tool)) && !prunedIdsLower?.has(id)) {
                         tracker.toolResultCount++
                     }
                 }