Skip to content

Commit 431e9be

Browse files
TarquinenTarquinen
authored
Merge pull request #58 from Tarquinen/fix/pr-57-bugs
fix: case-sensitive protected tools, cross-session dedup, unbounded cache
2 parents 1b1fc0c + 296ed70 commit 431e9be

File tree

8 files changed

+94
-98
lines changed

8 files changed

+94
-98
lines changed

‎lib/config.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -123,7 +123,7 @@ function createDefaultConfig(): void {
123123
"pruning_summary": "detailed",
124124
// How often to nudge the AI to prune (every N tool results, 0 = disabled)
125125
"nudge_freq": 10
126-
// Additional tools to protect from pruning (merged with built-in: task, todowrite, todoread, prune)
126+
// Additional tools to protect from pruning (merged with built-in: task, todowrite, todoread, prune, batch)
127127
// "protectedTools": ["bash"]
128128
}
129129
`

‎lib/core/janitor.ts‎

Lines changed: 8 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ export interface SessionStats {
1616
totalToolsPruned: number
1717
totalTokensSaved: number
1818
totalGCTokens: number
19+
totalGCTools: number
1920
}
2021

2122
export interface GCStats {
@@ -138,10 +139,8 @@ async function runWithStrategies(
138139
const alreadyPrunedIds = state.prunedIds.get(sessionID) ?? []
139140
const unprunedToolCallIds = toolCallIds.filter(id => !alreadyPrunedIds.includes(id))
140141

141-
// Get pending GC stats (accumulated since last notification)
142142
const gcPending = state.gcPending.get(sessionID) ?? null
143143

144-
// If nothing to analyze and no GC activity, exit early
145144
if (unprunedToolCallIds.length === 0 && !gcPending) {
146145
return null
147146
}
@@ -169,30 +168,28 @@ async function runWithStrategies(
169168

170169
const finalNewlyPrunedIds = llmPrunedIds.filter(id => !alreadyPrunedIds.includes(id))
171170

172-
// If AI pruned nothing and no GC activity, nothing to report
173171
if (finalNewlyPrunedIds.length === 0 && !gcPending) {
174172
return null
175173
}
176174

177175
// PHASE 2: CALCULATE STATS & NOTIFICATION
178176
const tokensSaved = await calculateTokensSaved(finalNewlyPrunedIds, toolOutputs)
179177

180-
// Get current session stats, initializing with proper defaults
181178
const currentStats = state.stats.get(sessionID) ?? {
182179
totalToolsPruned: 0,
183180
totalTokensSaved: 0,
184-
totalGCTokens: 0
181+
totalGCTokens: 0,
182+
totalGCTools: 0
185183
}
186184

187-
// Update session stats including GC contribution
188185
const sessionStats: SessionStats = {
189186
totalToolsPruned: currentStats.totalToolsPruned + finalNewlyPrunedIds.length,
190187
totalTokensSaved: currentStats.totalTokensSaved + tokensSaved,
191-
totalGCTokens: currentStats.totalGCTokens + (gcPending?.tokensCollected ?? 0)
188+
totalGCTokens: currentStats.totalGCTokens + (gcPending?.tokensCollected ?? 0),
189+
totalGCTools: currentStats.totalGCTools + (gcPending?.toolsDeduped ?? 0)
192190
}
193191
state.stats.set(sessionID, sessionStats)
194192

195-
// Send unified notification (handles all scenarios)
196193
const notificationSent = await sendUnifiedNotification(
197194
ctx.notificationCtx,
198195
sessionID,
@@ -207,12 +204,10 @@ async function runWithStrategies(
207204
currentAgent
208205
)
209206

210-
// Clear pending GC stats after notification (whether sent or not - we've consumed them)
211207
if (gcPending) {
212208
state.gcPending.delete(sessionID)
213209
}
214210

215-
// If we only had GC activity (no AI pruning), return null but notification was sent
216211
if (finalNewlyPrunedIds.length === 0) {
217212
if (notificationSent) {
218213
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<string
466461
const outputsToTokenize: string[] = []
467462

468463
for (const prunedId of prunedIds) {
469-
const output = toolOutputs.get(prunedId)
464+
// toolOutputs uses lowercase keys, so normalize the lookup
465+
const normalizedId = prunedId.toLowerCase()
466+
const output = toolOutputs.get(normalizedId)
470467
if (output) {
471468
outputsToTokenize.push(output)
472469
}

‎lib/core/strategies/deduplication.ts‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,8 @@ export const deduplicationStrategy: PruningStrategy = {
1717

1818
const deduplicatableIds = unprunedIds.filter(id => {
1919
const metadata = toolMetadata.get(id)
20-
return !metadata || !protectedTools.includes(metadata.tool)
20+
const protectedToolsLower = protectedTools.map(t => t.toLowerCase())
21+
return !metadata || !protectedToolsLower.includes(metadata.tool.toLowerCase())
2122
})
2223

2324
for (const id of deduplicatableIds) {

‎lib/fetch-wrapper/index.ts‎

Lines changed: 10 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import { handleGemini } from "./gemini"
88
import { handleOpenAIResponses } from "./openai-responses"
99
import { runStrategies } from "../core/strategies"
1010
import { accumulateGCStats } from "./gc-tracker"
11+
import { trimToolParametersCache } from "../state/tool-cache"
1112

1213
export type { FetchHandlerContext, FetchHandlerResult, SynthPrompts } from "./types"
1314

@@ -41,7 +42,6 @@ export function installFetchWrapper(
4142
}
4243

4344
globalThis.fetch = async (input: any, init?: any) => {
44-
// Skip all DCP processing for subagent sessions
4545
if (state.lastSeenSessionId && state.subagentSessions.has(state.lastSeenSessionId)) {
4646
logger.debug("fetch-wrapper", "Skipping DCP processing for subagent session", {
4747
sessionId: state.lastSeenSessionId.substring(0, 8)
@@ -55,6 +55,9 @@ export function installFetchWrapper(
5555
const inputUrl = typeof input === 'string' ? input : 'URL object'
5656
let modified = false
5757

58+
// Capture tool IDs before handlers run to track what gets cached this request
59+
const toolIdsBefore = new Set(state.toolParameters.keys())
60+
5861
// Try each format handler in order
5962
// OpenAI Chat Completions & Anthropic style (body.messages)
6063
if (body.messages && Array.isArray(body.messages)) {
@@ -80,9 +83,11 @@ export function installFetchWrapper(
8083
}
8184
}
8285

83-
// Run strategies after handlers have populated toolParameters cache
8486
const sessionId = state.lastSeenSessionId
85-
if (sessionId && state.toolParameters.size > 0) {
87+
const toolIdsAfter = Array.from(state.toolParameters.keys())
88+
const newToolsCached = toolIdsAfter.filter(id => !toolIdsBefore.has(id)).length > 0
89+
90+
if (sessionId && newToolsCached && state.toolParameters.size > 0) {
8691
const toolIds = Array.from(state.toolParameters.keys())
8792
const alreadyPruned = state.prunedIds.get(sessionId) ?? []
8893
const alreadyPrunedLower = new Set(alreadyPruned.map(id => id.toLowerCase()))
@@ -94,14 +99,13 @@ export function installFetchWrapper(
9499
config.protectedTools
95100
)
96101
if (result.prunedIds.length > 0) {
97-
// Normalize to lowercase to match janitor's ID normalization
98102
const normalizedIds = result.prunedIds.map(id => id.toLowerCase())
99103
state.prunedIds.set(sessionId, [...new Set([...alreadyPruned, ...normalizedIds])])
100-
101-
// Track GC activity for the next notification
102104
accumulateGCStats(state, sessionId, result.prunedIds, body, logger)
103105
}
104106
}
107+
108+
trimToolParametersCache(state)
105109
}
106110

107111
if (modified) {

‎lib/state/index.ts‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -64,7 +64,8 @@ export async function ensureSessionRestored(
6464
const stats: SessionStats = {
6565
totalToolsPruned: persisted.stats.totalToolsPruned,
6666
totalTokensSaved: persisted.stats.totalTokensSaved,
67-
totalGCTokens: persisted.stats.totalGCTokens ?? 0
67+
totalGCTokens: persisted.stats.totalGCTokens ?? 0,
68+
totalGCTools: persisted.stats.totalGCTools ?? 0
6869
}
6970
state.stats.set(sessionId, stats)
7071
}

‎lib/state/tool-cache.ts‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -59,3 +59,23 @@ export function cacheToolParametersFromInput(
5959
}
6060
}
6161
}
62+
63+
/** Maximum number of entries to keep in the tool parameters cache */
64+
const MAX_TOOL_CACHE_SIZE = 500
65+
66+
/**
67+
* Trim the tool parameters cache to prevent unbounded memory growth.
68+
* Uses FIFO eviction - removes oldest entries first.
69+
*/
70+
export function trimToolParametersCache(state: PluginState): void {
71+
if (state.toolParameters.size <= MAX_TOOL_CACHE_SIZE) {
72+
return
73+
}
74+
75+
const keysToRemove = Array.from(state.toolParameters.keys())
76+
.slice(0, state.toolParameters.size - MAX_TOOL_CACHE_SIZE)
77+
78+
for (const key of keysToRemove) {
79+
state.toolParameters.delete(key)
80+
}
81+
}

‎lib/tokenizer.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ export async function estimateTokensBatch(texts: string[]): Promise<number[]> {
99

1010
export function formatTokenCount(tokens: number): string {
1111
if (tokens >= 1000) {
12-
return `${(tokens / 1000).toFixed(1)}K`.replace('.0K', 'K')
12+
return `${(tokens / 1000).toFixed(1)}K`.replace('.0K', 'K') + ' tokens'
1313
}
14-
return tokens.toString()
14+
return tokens.toString() + ' tokens'
1515
}

‎lib/ui/notification.ts‎

Lines changed: 49 additions & 76 deletions
Original file line numberDiff line numberDiff line change
@@ -75,105 +75,78 @@ export async function sendUnifiedNotification(
7575
}
7676

7777
function buildMinimalMessage(data: NotificationData): string {
78-
const hasAiPruning = data.aiPrunedCount > 0
79-
const hasGcActivity = data.gcPending && data.gcPending.toolsDeduped > 0
78+
const { justNowTokens, totalTokens } = calculateStats(data)
8079

81-
if (hasAiPruning) {
82-
const gcTokens = hasGcActivity ? data.gcPending!.tokensCollected : 0
83-
const totalSaved = formatTokenCount(data.aiTokensSaved + gcTokens)
84-
const toolText = data.aiPrunedCount === 1 ? 'tool' : 'tools'
85-
86-
let cycleStats = `${data.aiPrunedCount} ${toolText}`
87-
if (hasGcActivity) {
88-
cycleStats += `, ~${formatTokenCount(data.gcPending!.tokensCollected)} 🗑️`
89-
}
80+
return formatStatsHeader(totalTokens, justNowTokens)
81+
}
9082

91-
let message = `🧹 DCP: ~${totalSaved} saved (${cycleStats})`
92-
message += buildSessionSuffix(data.sessionStats, data.aiPrunedCount)
83+
function calculateStats(data: NotificationData): {
84+
justNowTokens: number
85+
totalTokens: number
86+
} {
87+
const justNowTokens = data.aiTokensSaved + (data.gcPending?.tokensCollected ?? 0)
9388

94-
return message
95-
} else {
96-
const tokensCollected = formatTokenCount(data.gcPending!.tokensCollected)
89+
const totalTokens = data.sessionStats
90+
? data.sessionStats.totalTokensSaved + data.sessionStats.totalGCTokens
91+
: justNowTokens
9792

98-
let message = `🗑️ DCP: ~${tokensCollected} collected`
99-
message += buildSessionSuffix(data.sessionStats, 0)
93+
return { justNowTokens, totalTokens }
94+
}
10095

101-
return message
102-
}
96+
function formatStatsHeader(
97+
totalTokens: number,
98+
justNowTokens: number
99+
): string {
100+
const totalTokensStr = `~${formatTokenCount(totalTokens)}`
101+
const justNowTokensStr = `~${formatTokenCount(justNowTokens)}`
102+
103+
const maxTokenLen = Math.max(totalTokensStr.length, justNowTokensStr.length)
104+
const totalTokensPadded = totalTokensStr.padStart(maxTokenLen)
105+
const justNowTokensPadded = justNowTokensStr.padStart(maxTokenLen)
106+
107+
return [
108+
`▣ DCP Stats`,
109+
` Total saved │ ${totalTokensPadded}`,
110+
` Just now │ ${justNowTokensPadded}`,
111+
].join('\n')
103112
}
104113

105114
function buildDetailedMessage(data: NotificationData, workingDirectory?: string): string {
106-
const hasAiPruning = data.aiPrunedCount > 0
107-
const hasGcActivity = data.gcPending && data.gcPending.toolsDeduped > 0
115+
const { justNowTokens, totalTokens } = calculateStats(data)
108116

109-
let message: string
110-
111-
if (hasAiPruning) {
112-
const gcTokens = hasGcActivity ? data.gcPending!.tokensCollected : 0
113-
const totalSaved = formatTokenCount(data.aiTokensSaved + gcTokens)
114-
const toolText = data.aiPrunedCount === 1 ? 'tool' : 'tools'
115-
116-
let cycleStats = `${data.aiPrunedCount} ${toolText}`
117-
if (hasGcActivity) {
118-
cycleStats += `, ~${formatTokenCount(data.gcPending!.tokensCollected)} 🗑️`
119-
}
117+
let message = formatStatsHeader(totalTokens, justNowTokens)
120118

121-
message = `🧹 DCP: ~${totalSaved} saved (${cycleStats})`
122-
message += buildSessionSuffix(data.sessionStats, data.aiPrunedCount)
123-
message += '\n'
119+
if (data.aiPrunedCount > 0) {
120+
message += '\n\n▣ Pruned tools:'
124121

125-
message += `\n🤖 LLM analysis (${data.aiPrunedIds.length}):\n`
126-
const toolsSummary = buildToolsSummary(data.aiPrunedIds, data.toolMetadata, workingDirectory)
122+
for (const prunedId of data.aiPrunedIds) {
123+
const normalizedId = prunedId.toLowerCase()
124+
const metadata = data.toolMetadata.get(normalizedId)
127125

128-
for (const [toolName, params] of toolsSummary.entries()) {
129-
if (params.length > 0) {
130-
message += ` ${toolName} (${params.length}):\n`
131-
for (const param of params) {
132-
message += ` ${param}\n`
126+
if (metadata) {
127+
const paramKey = extractParameterKey(metadata)
128+
if (paramKey) {
129+
const displayKey = truncate(shortenPath(paramKey, workingDirectory), 60)
130+
message += `\n→ ${metadata.tool}: ${displayKey}`
131+
} else {
132+
message += `\n→ ${metadata.tool}`
133133
}
134134
}
135135
}
136136

137-
const foundToolNames = new Set(toolsSummary.keys())
138-
const missingTools = data.aiPrunedIds.filter(id => {
139-
const normalizedId = id.toLowerCase()
140-
const metadata = data.toolMetadata.get(normalizedId)
141-
return !metadata || !foundToolNames.has(metadata.tool)
142-
})
137+
const knownCount = data.aiPrunedIds.filter(id =>
138+
data.toolMetadata.has(id.toLowerCase())
139+
).length
140+
const unknownCount = data.aiPrunedIds.length - knownCount
143141

144-
if (missingTools.length > 0) {
145-
message += ` (${missingTools.length} tool${missingTools.length > 1 ? 's' : ''} with unknown metadata)\n`
142+
if (unknownCount > 0) {
143+
message += `\n→ (${unknownCount} tool${unknownCount > 1 ? 's' : ''} with unknown metadata)`
146144
}
147-
} else {
148-
const tokensCollected = formatTokenCount(data.gcPending!.tokensCollected)
149-
150-
message = `🗑️ DCP: ~${tokensCollected} collected`
151-
message += buildSessionSuffix(data.sessionStats, 0)
152145
}
153146

154147
return message.trim()
155148
}
156149

157-
function buildSessionSuffix(sessionStats: SessionStats | null, currentAiPruned: number): string {
158-
if (!sessionStats) {
159-
return ''
160-
}
161-
162-
if (sessionStats.totalToolsPruned <= currentAiPruned) {
163-
return ''
164-
}
165-
166-
const totalSaved = sessionStats.totalTokensSaved + sessionStats.totalGCTokens
167-
let suffix = ` │ Session: ~${formatTokenCount(totalSaved)} (${sessionStats.totalToolsPruned} tools`
168-
169-
if (sessionStats.totalGCTokens > 0) {
170-
suffix += `, ~${formatTokenCount(sessionStats.totalGCTokens)} 🗑️`
171-
}
172-
173-
suffix += ')'
174-
return suffix
175-
}
176-
177150
export function formatPruningResultForTool(
178151
result: PruningResult,
179152
workingDirectory?: string

0 commit comments

Comments
 (0)