diff --git a/lib/core/janitor.ts b/lib/core/janitor.ts index 75f7232..55772ac 100644 --- a/lib/core/janitor.ts +++ b/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" diff --git a/lib/fetch-wrapper/tool-tracker.ts b/lib/fetch-wrapper/tool-tracker.ts index 639b99b..a195a67 100644 --- a/lib/fetch-wrapper/tool-tracker.ts +++ b/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 +} diff --git a/lib/hooks.ts b/lib/hooks.ts index 617abe1..48ca6cc 100644 --- a/lib/hooks.ts +++ b/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 { @@ -80,7 +80,7 @@ export function createChatParamsHandler( clearAllMappings() state.toolParameters.clear() if (toolTracker) { - resetToolTrackerCount(toolTracker) + clearToolTracker(toolTracker) } } diff --git a/lib/state/tool-cache.ts b/lib/state/tool-cache.ts index de1c9c3..2970720 100644 --- a/lib/state/tool-cache.ts +++ b/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++ } }