diff --git a/lib/config.ts b/lib/config.ts index 74e173e..ad2856c 100644 --- a/lib/config.ts +++ b/lib/config.ts @@ -123,7 +123,7 @@ function createDefaultConfig(): void { "pruning_summary": "detailed", // How often to nudge the AI to prune (every N tool results, 0 = disabled) "nudge_freq": 10 - // Additional tools to protect from pruning (merged with built-in: task, todowrite, todoread, prune) + // Additional tools to protect from pruning (merged with built-in: task, todowrite, todoread, prune, batch) // "protectedTools": ["bash"] } ` diff --git a/lib/core/janitor.ts b/lib/core/janitor.ts index 726170c..70005d8 100644 --- a/lib/core/janitor.ts +++ b/lib/core/janitor.ts @@ -16,6 +16,7 @@ export interface SessionStats { totalToolsPruned: number totalTokensSaved: number totalGCTokens: number + totalGCTools: number } export interface GCStats { @@ -138,10 +139,8 @@ async function runWithStrategies( const alreadyPrunedIds = state.prunedIds.get(sessionID) ?? [] const unprunedToolCallIds = toolCallIds.filter(id => !alreadyPrunedIds.includes(id)) - // Get pending GC stats (accumulated since last notification) const gcPending = state.gcPending.get(sessionID) ?? null - // If nothing to analyze and no GC activity, exit early if (unprunedToolCallIds.length === 0 && !gcPending) { return null } @@ -169,7 +168,6 @@ async function runWithStrategies( const finalNewlyPrunedIds = llmPrunedIds.filter(id => !alreadyPrunedIds.includes(id)) - // If AI pruned nothing and no GC activity, nothing to report if (finalNewlyPrunedIds.length === 0 && !gcPending) { return null } @@ -177,22 +175,21 @@ async function runWithStrategies( // PHASE 2: CALCULATE STATS & NOTIFICATION const tokensSaved = await calculateTokensSaved(finalNewlyPrunedIds, toolOutputs) - // Get current session stats, initializing with proper defaults const currentStats = state.stats.get(sessionID) ?? { totalToolsPruned: 0, totalTokensSaved: 0, - totalGCTokens: 0 + totalGCTokens: 0, + totalGCTools: 0 } - // Update session stats including GC contribution const sessionStats: SessionStats = { totalToolsPruned: currentStats.totalToolsPruned + finalNewlyPrunedIds.length, totalTokensSaved: currentStats.totalTokensSaved + tokensSaved, - totalGCTokens: currentStats.totalGCTokens + (gcPending?.tokensCollected ?? 0) + totalGCTokens: currentStats.totalGCTokens + (gcPending?.tokensCollected ?? 0), + totalGCTools: currentStats.totalGCTools + (gcPending?.toolsDeduped ?? 0) } state.stats.set(sessionID, sessionStats) - // Send unified notification (handles all scenarios) const notificationSent = await sendUnifiedNotification( ctx.notificationCtx, sessionID, @@ -207,12 +204,10 @@ async function runWithStrategies( currentAgent ) - // Clear pending GC stats after notification (whether sent or not - we've consumed them) if (gcPending) { state.gcPending.delete(sessionID) } - // If we only had GC activity (no AI pruning), return null but notification was sent if (finalNewlyPrunedIds.length === 0) { if (notificationSent) { logger.info("janitor", `GC-only notification: ~${formatTokenCount(gcPending?.tokensCollected ?? 0)} tokens from ${gcPending?.toolsDeduped ?? 0} deduped tools`, { @@ -466,7 +461,9 @@ async function calculateTokensSaved(prunedIds: string[], toolOutputs: Map { const metadata = toolMetadata.get(id) - return !metadata || !protectedTools.includes(metadata.tool) + const protectedToolsLower = protectedTools.map(t => t.toLowerCase()) + return !metadata || !protectedToolsLower.includes(metadata.tool.toLowerCase()) }) for (const id of deduplicatableIds) { diff --git a/lib/fetch-wrapper/index.ts b/lib/fetch-wrapper/index.ts index 450e99c..868c9b8 100644 --- a/lib/fetch-wrapper/index.ts +++ b/lib/fetch-wrapper/index.ts @@ -8,6 +8,7 @@ import { handleGemini } from "./gemini" import { handleOpenAIResponses } from "./openai-responses" import { runStrategies } from "../core/strategies" import { accumulateGCStats } from "./gc-tracker" +import { trimToolParametersCache } from "../state/tool-cache" export type { FetchHandlerContext, FetchHandlerResult, SynthPrompts } from "./types" @@ -41,7 +42,6 @@ export function installFetchWrapper( } globalThis.fetch = async (input: any, init?: any) => { - // Skip all DCP processing for subagent sessions if (state.lastSeenSessionId && state.subagentSessions.has(state.lastSeenSessionId)) { logger.debug("fetch-wrapper", "Skipping DCP processing for subagent session", { sessionId: state.lastSeenSessionId.substring(0, 8) @@ -55,6 +55,9 @@ export function installFetchWrapper( const inputUrl = typeof input === 'string' ? input : 'URL object' let modified = false + // Capture tool IDs before handlers run to track what gets cached this request + const toolIdsBefore = new Set(state.toolParameters.keys()) + // Try each format handler in order // OpenAI Chat Completions & Anthropic style (body.messages) if (body.messages && Array.isArray(body.messages)) { @@ -80,9 +83,11 @@ export function installFetchWrapper( } } - // Run strategies after handlers have populated toolParameters cache const sessionId = state.lastSeenSessionId - if (sessionId && state.toolParameters.size > 0) { + const toolIdsAfter = Array.from(state.toolParameters.keys()) + const newToolsCached = toolIdsAfter.filter(id => !toolIdsBefore.has(id)).length > 0 + + if (sessionId && newToolsCached && state.toolParameters.size > 0) { const toolIds = Array.from(state.toolParameters.keys()) const alreadyPruned = state.prunedIds.get(sessionId) ?? [] const alreadyPrunedLower = new Set(alreadyPruned.map(id => id.toLowerCase())) @@ -94,14 +99,13 @@ export function installFetchWrapper( config.protectedTools ) if (result.prunedIds.length > 0) { - // Normalize to lowercase to match janitor's ID normalization const normalizedIds = result.prunedIds.map(id => id.toLowerCase()) state.prunedIds.set(sessionId, [...new Set([...alreadyPruned, ...normalizedIds])]) - - // Track GC activity for the next notification accumulateGCStats(state, sessionId, result.prunedIds, body, logger) } } + + trimToolParametersCache(state) } if (modified) { diff --git a/lib/state/index.ts b/lib/state/index.ts index b48c656..2808cb4 100644 --- a/lib/state/index.ts +++ b/lib/state/index.ts @@ -64,7 +64,8 @@ export async function ensureSessionRestored( const stats: SessionStats = { totalToolsPruned: persisted.stats.totalToolsPruned, totalTokensSaved: persisted.stats.totalTokensSaved, - totalGCTokens: persisted.stats.totalGCTokens ?? 0 + totalGCTokens: persisted.stats.totalGCTokens ?? 0, + totalGCTools: persisted.stats.totalGCTools ?? 0 } state.stats.set(sessionId, stats) } diff --git a/lib/state/tool-cache.ts b/lib/state/tool-cache.ts index aa57b4b..27549ea 100644 --- a/lib/state/tool-cache.ts +++ b/lib/state/tool-cache.ts @@ -59,3 +59,23 @@ export function cacheToolParametersFromInput( } } } + +/** Maximum number of entries to keep in the tool parameters cache */ +const MAX_TOOL_CACHE_SIZE = 500 + +/** + * Trim the tool parameters cache to prevent unbounded memory growth. + * Uses FIFO eviction - removes oldest entries first. + */ +export function trimToolParametersCache(state: PluginState): void { + if (state.toolParameters.size <= MAX_TOOL_CACHE_SIZE) { + return + } + + const keysToRemove = Array.from(state.toolParameters.keys()) + .slice(0, state.toolParameters.size - MAX_TOOL_CACHE_SIZE) + + for (const key of keysToRemove) { + state.toolParameters.delete(key) + } +} diff --git a/lib/tokenizer.ts b/lib/tokenizer.ts index 4013564..711a449 100644 --- a/lib/tokenizer.ts +++ b/lib/tokenizer.ts @@ -9,7 +9,7 @@ export async function estimateTokensBatch(texts: string[]): Promise { export function formatTokenCount(tokens: number): string { if (tokens >= 1000) { - return `${(tokens / 1000).toFixed(1)}K`.replace('.0K', 'K') + return `${(tokens / 1000).toFixed(1)}K`.replace('.0K', 'K') + ' tokens' } - return tokens.toString() + return tokens.toString() + ' tokens' } diff --git a/lib/ui/notification.ts b/lib/ui/notification.ts index 7ea5772..47bf8dd 100644 --- a/lib/ui/notification.ts +++ b/lib/ui/notification.ts @@ -75,105 +75,78 @@ export async function sendUnifiedNotification( } function buildMinimalMessage(data: NotificationData): string { - const hasAiPruning = data.aiPrunedCount > 0 - const hasGcActivity = data.gcPending && data.gcPending.toolsDeduped > 0 + const { justNowTokens, totalTokens } = calculateStats(data) - if (hasAiPruning) { - const gcTokens = hasGcActivity ? data.gcPending!.tokensCollected : 0 - const totalSaved = formatTokenCount(data.aiTokensSaved + gcTokens) - const toolText = data.aiPrunedCount === 1 ? 'tool' : 'tools' - - let cycleStats = `${data.aiPrunedCount} ${toolText}` - if (hasGcActivity) { - cycleStats += `, ~${formatTokenCount(data.gcPending!.tokensCollected)} ๐Ÿ—‘๏ธ` - } + return formatStatsHeader(totalTokens, justNowTokens) +} - let message = `๐Ÿงน DCP: ~${totalSaved} saved (${cycleStats})` - message += buildSessionSuffix(data.sessionStats, data.aiPrunedCount) +function calculateStats(data: NotificationData): { + justNowTokens: number + totalTokens: number +} { + const justNowTokens = data.aiTokensSaved + (data.gcPending?.tokensCollected ?? 0) - return message - } else { - const tokensCollected = formatTokenCount(data.gcPending!.tokensCollected) + const totalTokens = data.sessionStats + ? data.sessionStats.totalTokensSaved + data.sessionStats.totalGCTokens + : justNowTokens - let message = `๐Ÿ—‘๏ธ DCP: ~${tokensCollected} collected` - message += buildSessionSuffix(data.sessionStats, 0) + return { justNowTokens, totalTokens } +} - return message - } +function formatStatsHeader( + totalTokens: number, + justNowTokens: number +): string { + const totalTokensStr = `~${formatTokenCount(totalTokens)}` + const justNowTokensStr = `~${formatTokenCount(justNowTokens)}` + + const maxTokenLen = Math.max(totalTokensStr.length, justNowTokensStr.length) + const totalTokensPadded = totalTokensStr.padStart(maxTokenLen) + const justNowTokensPadded = justNowTokensStr.padStart(maxTokenLen) + + return [ + `โ–ฃ DCP Stats`, + ` Total saved โ”‚ ${totalTokensPadded}`, + ` Just now โ”‚ ${justNowTokensPadded}`, + ].join('\n') } function buildDetailedMessage(data: NotificationData, workingDirectory?: string): string { - const hasAiPruning = data.aiPrunedCount > 0 - const hasGcActivity = data.gcPending && data.gcPending.toolsDeduped > 0 + const { justNowTokens, totalTokens } = calculateStats(data) - let message: string - - if (hasAiPruning) { - const gcTokens = hasGcActivity ? data.gcPending!.tokensCollected : 0 - const totalSaved = formatTokenCount(data.aiTokensSaved + gcTokens) - const toolText = data.aiPrunedCount === 1 ? 'tool' : 'tools' - - let cycleStats = `${data.aiPrunedCount} ${toolText}` - if (hasGcActivity) { - cycleStats += `, ~${formatTokenCount(data.gcPending!.tokensCollected)} ๐Ÿ—‘๏ธ` - } + let message = formatStatsHeader(totalTokens, justNowTokens) - message = `๐Ÿงน DCP: ~${totalSaved} saved (${cycleStats})` - message += buildSessionSuffix(data.sessionStats, data.aiPrunedCount) - message += '\n' + if (data.aiPrunedCount > 0) { + message += '\n\nโ–ฃ Pruned tools:' - message += `\n๐Ÿค– LLM analysis (${data.aiPrunedIds.length}):\n` - const toolsSummary = buildToolsSummary(data.aiPrunedIds, data.toolMetadata, workingDirectory) + for (const prunedId of data.aiPrunedIds) { + const normalizedId = prunedId.toLowerCase() + const metadata = data.toolMetadata.get(normalizedId) - for (const [toolName, params] of toolsSummary.entries()) { - if (params.length > 0) { - message += ` ${toolName} (${params.length}):\n` - for (const param of params) { - message += ` ${param}\n` + if (metadata) { + const paramKey = extractParameterKey(metadata) + if (paramKey) { + const displayKey = truncate(shortenPath(paramKey, workingDirectory), 60) + message += `\nโ†’ ${metadata.tool}: ${displayKey}` + } else { + message += `\nโ†’ ${metadata.tool}` } } } - const foundToolNames = new Set(toolsSummary.keys()) - const missingTools = data.aiPrunedIds.filter(id => { - const normalizedId = id.toLowerCase() - const metadata = data.toolMetadata.get(normalizedId) - return !metadata || !foundToolNames.has(metadata.tool) - }) + const knownCount = data.aiPrunedIds.filter(id => + data.toolMetadata.has(id.toLowerCase()) + ).length + const unknownCount = data.aiPrunedIds.length - knownCount - if (missingTools.length > 0) { - message += ` (${missingTools.length} tool${missingTools.length > 1 ? 's' : ''} with unknown metadata)\n` + if (unknownCount > 0) { + message += `\nโ†’ (${unknownCount} tool${unknownCount > 1 ? 's' : ''} with unknown metadata)` } - } else { - const tokensCollected = formatTokenCount(data.gcPending!.tokensCollected) - - message = `๐Ÿ—‘๏ธ DCP: ~${tokensCollected} collected` - message += buildSessionSuffix(data.sessionStats, 0) } return message.trim() } -function buildSessionSuffix(sessionStats: SessionStats | null, currentAiPruned: number): string { - if (!sessionStats) { - return '' - } - - if (sessionStats.totalToolsPruned <= currentAiPruned) { - return '' - } - - const totalSaved = sessionStats.totalTokensSaved + sessionStats.totalGCTokens - let suffix = ` โ”‚ Session: ~${formatTokenCount(totalSaved)} (${sessionStats.totalToolsPruned} tools` - - if (sessionStats.totalGCTokens > 0) { - suffix += `, ~${formatTokenCount(sessionStats.totalGCTokens)} ๐Ÿ—‘๏ธ` - } - - suffix += ')' - return suffix -} - export function formatPruningResultForTool( result: PruningResult, workingDirectory?: string