Skip to content

Commit 0e3a28d

Browse files
TarquinenTarquinen
committed
Update config and state management for discard and extract tools
1 parent 60dffe3 commit 0e3a28d

File tree

5 files changed

+337
-173
lines changed

5 files changed

+337
-173
lines changed

‎lib/config.ts‎

Lines changed: 153 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,20 @@ export interface PruneTool {
3434
nudge: PruneToolNudge
3535
}
3636

37+
export interface DiscardTool {
38+
enabled: boolean
39+
protectedTools: string[]
40+
turnProtection: PruneToolTurnProtection
41+
nudge: PruneToolNudge
42+
}
43+
44+
export interface ExtractTool {
45+
enabled: boolean
46+
protectedTools: string[]
47+
turnProtection: PruneToolTurnProtection
48+
nudge: PruneToolNudge
49+
}
50+
3751
export interface SupersedeWrites {
3852
enabled: boolean
3953
}
@@ -45,12 +59,13 @@ export interface PluginConfig {
4559
strategies: {
4660
deduplication: Deduplication
4761
onIdle: OnIdle
48-
pruneTool: PruneTool
62+
discardTool: DiscardTool
63+
extractTool: ExtractTool
4964
supersedeWrites: SupersedeWrites
5065
}
5166
}
5267

53-
const DEFAULT_PROTECTED_TOOLS = ['task', 'todowrite', 'todoread', 'prune', 'batch']
68+
const DEFAULT_PROTECTED_TOOLS = ['task', 'todowrite', 'todoread', 'discard', 'extract', 'batch']
5469

5570
// Valid config keys for validation against user config
5671
export const VALID_CONFIG_KEYS = new Set([
@@ -74,16 +89,26 @@ export const VALID_CONFIG_KEYS = new Set([
7489
'strategies.onIdle.showModelErrorToasts',
7590
'strategies.onIdle.strictModelSelection',
7691
'strategies.onIdle.protectedTools',
77-
// strategies.pruneTool
78-
'strategies.pruneTool',
79-
'strategies.pruneTool.enabled',
80-
'strategies.pruneTool.protectedTools',
81-
'strategies.pruneTool.turnProtection',
82-
'strategies.pruneTool.turnProtection.enabled',
83-
'strategies.pruneTool.turnProtection.turns',
84-
'strategies.pruneTool.nudge',
85-
'strategies.pruneTool.nudge.enabled',
86-
'strategies.pruneTool.nudge.frequency'
92+
// strategies.discardTool
93+
'strategies.discardTool',
94+
'strategies.discardTool.enabled',
95+
'strategies.discardTool.protectedTools',
96+
'strategies.discardTool.turnProtection',
97+
'strategies.discardTool.turnProtection.enabled',
98+
'strategies.discardTool.turnProtection.turns',
99+
'strategies.discardTool.nudge',
100+
'strategies.discardTool.nudge.enabled',
101+
'strategies.discardTool.nudge.frequency',
102+
// strategies.extractTool
103+
'strategies.extractTool',
104+
'strategies.extractTool.enabled',
105+
'strategies.extractTool.protectedTools',
106+
'strategies.extractTool.turnProtection',
107+
'strategies.extractTool.turnProtection.enabled',
108+
'strategies.extractTool.turnProtection.turns',
109+
'strategies.extractTool.nudge',
110+
'strategies.extractTool.nudge.enabled',
111+
'strategies.extractTool.nudge.frequency'
87112
])
88113

89114
// Extract all key paths from a config object for validation
@@ -159,28 +184,54 @@ function validateConfigTypes(config: Record<string, any>): ValidationError[] {
159184
}
160185
}
161186

162-
// pruneTool
163-
if (strategies.pruneTool) {
164-
if (strategies.pruneTool.enabled !== undefined && typeof strategies.pruneTool.enabled !== 'boolean') {
165-
errors.push({ key: 'strategies.pruneTool.enabled', expected: 'boolean', actual: typeof strategies.pruneTool.enabled })
187+
// discardTool
188+
if (strategies.discardTool) {
189+
if (strategies.discardTool.enabled !== undefined && typeof strategies.discardTool.enabled !== 'boolean') {
190+
errors.push({ key: 'strategies.discardTool.enabled', expected: 'boolean', actual: typeof strategies.discardTool.enabled })
166191
}
167-
if (strategies.pruneTool.protectedTools !== undefined && !Array.isArray(strategies.pruneTool.protectedTools)) {
168-
errors.push({ key: 'strategies.pruneTool.protectedTools', expected: 'string[]', actual: typeof strategies.pruneTool.protectedTools })
192+
if (strategies.discardTool.protectedTools !== undefined && !Array.isArray(strategies.discardTool.protectedTools)) {
193+
errors.push({ key: 'strategies.discardTool.protectedTools', expected: 'string[]', actual: typeof strategies.discardTool.protectedTools })
169194
}
170-
if (strategies.pruneTool.turnProtection) {
171-
if (strategies.pruneTool.turnProtection.enabled !== undefined && typeof strategies.pruneTool.turnProtection.enabled !== 'boolean') {
172-
errors.push({ key: 'strategies.pruneTool.turnProtection.enabled', expected: 'boolean', actual: typeof strategies.pruneTool.turnProtection.enabled })
195+
if (strategies.discardTool.turnProtection) {
196+
if (strategies.discardTool.turnProtection.enabled !== undefined && typeof strategies.discardTool.turnProtection.enabled !== 'boolean') {
197+
errors.push({ key: 'strategies.discardTool.turnProtection.enabled', expected: 'boolean', actual: typeof strategies.discardTool.turnProtection.enabled })
173198
}
174-
if (strategies.pruneTool.turnProtection.turns !== undefined && typeof strategies.pruneTool.turnProtection.turns !== 'number') {
175-
errors.push({ key: 'strategies.pruneTool.turnProtection.turns', expected: 'number', actual: typeof strategies.pruneTool.turnProtection.turns })
199+
if (strategies.discardTool.turnProtection.turns !== undefined && typeof strategies.discardTool.turnProtection.turns !== 'number') {
200+
errors.push({ key: 'strategies.discardTool.turnProtection.turns', expected: 'number', actual: typeof strategies.discardTool.turnProtection.turns })
176201
}
177202
}
178-
if (strategies.pruneTool.nudge) {
179-
if (strategies.pruneTool.nudge.enabled !== undefined && typeof strategies.pruneTool.nudge.enabled !== 'boolean') {
180-
errors.push({ key: 'strategies.pruneTool.nudge.enabled', expected: 'boolean', actual: typeof strategies.pruneTool.nudge.enabled })
203+
if (strategies.discardTool.nudge) {
204+
if (strategies.discardTool.nudge.enabled !== undefined && typeof strategies.discardTool.nudge.enabled !== 'boolean') {
205+
errors.push({ key: 'strategies.discardTool.nudge.enabled', expected: 'boolean', actual: typeof strategies.discardTool.nudge.enabled })
181206
}
182-
if (strategies.pruneTool.nudge.frequency !== undefined && typeof strategies.pruneTool.nudge.frequency !== 'number') {
183-
errors.push({ key: 'strategies.pruneTool.nudge.frequency', expected: 'number', actual: typeof strategies.pruneTool.nudge.frequency })
207+
if (strategies.discardTool.nudge.frequency !== undefined && typeof strategies.discardTool.nudge.frequency !== 'number') {
208+
errors.push({ key: 'strategies.discardTool.nudge.frequency', expected: 'number', actual: typeof strategies.discardTool.nudge.frequency })
209+
}
210+
}
211+
}
212+
213+
// extractTool
214+
if (strategies.extractTool) {
215+
if (strategies.extractTool.enabled !== undefined && typeof strategies.extractTool.enabled !== 'boolean') {
216+
errors.push({ key: 'strategies.extractTool.enabled', expected: 'boolean', actual: typeof strategies.extractTool.enabled })
217+
}
218+
if (strategies.extractTool.protectedTools !== undefined && !Array.isArray(strategies.extractTool.protectedTools)) {
219+
errors.push({ key: 'strategies.extractTool.protectedTools', expected: 'string[]', actual: typeof strategies.extractTool.protectedTools })
220+
}
221+
if (strategies.extractTool.turnProtection) {
222+
if (strategies.extractTool.turnProtection.enabled !== undefined && typeof strategies.extractTool.turnProtection.enabled !== 'boolean') {
223+
errors.push({ key: 'strategies.extractTool.turnProtection.enabled', expected: 'boolean', actual: typeof strategies.extractTool.turnProtection.enabled })
224+
}
225+
if (strategies.extractTool.turnProtection.turns !== undefined && typeof strategies.extractTool.turnProtection.turns !== 'number') {
226+
errors.push({ key: 'strategies.extractTool.turnProtection.turns', expected: 'number', actual: typeof strategies.extractTool.turnProtection.turns })
227+
}
228+
}
229+
if (strategies.extractTool.nudge) {
230+
if (strategies.extractTool.nudge.enabled !== undefined && typeof strategies.extractTool.nudge.enabled !== 'boolean') {
231+
errors.push({ key: 'strategies.extractTool.nudge.enabled', expected: 'boolean', actual: typeof strategies.extractTool.nudge.enabled })
232+
}
233+
if (strategies.extractTool.nudge.frequency !== undefined && typeof strategies.extractTool.nudge.frequency !== 'number') {
234+
errors.push({ key: 'strategies.extractTool.nudge.frequency', expected: 'number', actual: typeof strategies.extractTool.nudge.frequency })
184235
}
185236
}
186237
}
@@ -254,7 +305,19 @@ const defaultConfig: PluginConfig = {
254305
supersedeWrites: {
255306
enabled: true
256307
},
257-
pruneTool: {
308+
discardTool: {
309+
enabled: true,
310+
protectedTools: [...DEFAULT_PROTECTED_TOOLS],
311+
turnProtection: {
312+
enabled: false,
313+
turns: 4
314+
},
315+
nudge: {
316+
enabled: true,
317+
frequency: 10
318+
}
319+
},
320+
extractTool: {
258321
enabled: true,
259322
protectedTools: [...DEFAULT_PROTECTED_TOOLS],
260323
turnProtection: {
@@ -357,18 +420,34 @@ function createDefaultConfig(): void {
357420
"supersedeWrites": {
358421
"enabled": true
359422
},
360-
// Exposes a prune tool to your LLM to call when it determines pruning is necessary
361-
\"pruneTool\": {
362-
\"enabled\": true,
423+
// Exposes a discard tool to your LLM to call when it determines pruning is necessary
424+
"discardTool": {
425+
"enabled": true,
363426
// Additional tools to protect from pruning
364-
\"protectedTools\": [],
427+
"protectedTools": [],
365428
// Protect from pruning for <turn protection> message turns
366-
\"turnProtection\": {
367-
\"enabled\": false,
368-
\"turns\": 4
429+
"turnProtection": {
430+
"enabled": false,
431+
"turns": 4
369432
},
370-
// Nudge the LLM to use the prune tool (every <frequency> tool results)
371-
\"nudge\": {
433+
// Nudge the LLM to use the discard tool (every <frequency> tool results)
434+
"nudge": {
435+
"enabled": true,
436+
"frequency": 10
437+
}
438+
},
439+
// Exposes an extract tool to your LLM to call when it determines pruning is necessary
440+
"extractTool": {
441+
"enabled": true,
442+
// Additional tools to protect from pruning
443+
"protectedTools": [],
444+
// Protect from pruning for <turn protection> message turns
445+
"turnProtection": {
446+
"enabled": false,
447+
"turns": 4
448+
},
449+
// Nudge the LLM to use the extract tool (every <frequency> tool results)
450+
"nudge": {
372451
"enabled": true,
373452
"frequency": 10
374453
}
@@ -444,21 +523,38 @@ function mergeStrategies(
444523
])
445524
]
446525
},
447-
pruneTool: {
448-
enabled: override.pruneTool?.enabled ?? base.pruneTool.enabled,
526+
discardTool: {
527+
enabled: override.discardTool?.enabled ?? base.discardTool.enabled,
449528
protectedTools: [
450529
...new Set([
451-
...base.pruneTool.protectedTools,
452-
...(override.pruneTool?.protectedTools ?? [])
530+
...base.discardTool.protectedTools,
531+
...(override.discardTool?.protectedTools ?? [])
453532
])
454533
],
455534
turnProtection: {
456-
enabled: override.pruneTool?.turnProtection?.enabled ?? base.pruneTool.turnProtection.enabled,
457-
turns: override.pruneTool?.turnProtection?.turns ?? base.pruneTool.turnProtection.turns
535+
enabled: override.discardTool?.turnProtection?.enabled ?? base.discardTool.turnProtection.enabled,
536+
turns: override.discardTool?.turnProtection?.turns ?? base.discardTool.turnProtection.turns
458537
},
459538
nudge: {
460-
enabled: override.pruneTool?.nudge?.enabled ?? base.pruneTool.nudge.enabled,
461-
frequency: override.pruneTool?.nudge?.frequency ?? base.pruneTool.nudge.frequency
539+
enabled: override.discardTool?.nudge?.enabled ?? base.discardTool.nudge.enabled,
540+
frequency: override.discardTool?.nudge?.frequency ?? base.discardTool.nudge.frequency
541+
}
542+
},
543+
extractTool: {
544+
enabled: override.extractTool?.enabled ?? base.extractTool.enabled,
545+
protectedTools: [
546+
...new Set([
547+
...base.extractTool.protectedTools,
548+
...(override.extractTool?.protectedTools ?? [])
549+
])
550+
],
551+
turnProtection: {
552+
enabled: override.extractTool?.turnProtection?.enabled ?? base.extractTool.turnProtection.enabled,
553+
turns: override.extractTool?.turnProtection?.turns ?? base.extractTool.turnProtection.turns
554+
},
555+
nudge: {
556+
enabled: override.extractTool?.nudge?.enabled ?? base.extractTool.nudge.enabled,
557+
frequency: override.extractTool?.nudge?.frequency ?? base.extractTool.nudge.frequency
462558
}
463559
},
464560
supersedeWrites: {
@@ -479,11 +575,17 @@ function deepCloneConfig(config: PluginConfig): PluginConfig {
479575
...config.strategies.onIdle,
480576
protectedTools: [...config.strategies.onIdle.protectedTools]
481577
},
482-
pruneTool: {
483-
...config.strategies.pruneTool,
484-
protectedTools: [...config.strategies.pruneTool.protectedTools],
485-
turnProtection: { ...config.strategies.pruneTool.turnProtection },
486-
nudge: { ...config.strategies.pruneTool.nudge }
578+
discardTool: {
579+
...config.strategies.discardTool,
580+
protectedTools: [...config.strategies.discardTool.protectedTools],
581+
turnProtection: { ...config.strategies.discardTool.turnProtection },
582+
nudge: { ...config.strategies.discardTool.nudge }
583+
},
584+
extractTool: {
585+
...config.strategies.extractTool,
586+
protectedTools: [...config.strategies.extractTool.protectedTools],
587+
turnProtection: { ...config.strategies.extractTool.turnProtection },
588+
nudge: { ...config.strategies.extractTool.nudge }
487589
},
488590
supersedeWrites: {
489591
...config.strategies.supersedeWrites

‎lib/messages/prune.ts‎

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,11 @@ const buildPrunableToolsList = (
3434
if (state.prune.toolIds.includes(toolCallId)) {
3535
return
3636
}
37-
if (config.strategies.pruneTool.protectedTools.includes(toolParameterEntry.tool)) {
37+
const allProtectedTools = [
38+
...config.strategies.discardTool.protectedTools,
39+
...config.strategies.extractTool.protectedTools
40+
]
41+
if (allProtectedTools.includes(toolParameterEntry.tool)) {
3842
return
3943
}
4044
const numericId = toolIdList.indexOf(toolCallId)
@@ -61,7 +65,7 @@ export const insertPruneToolContext = (
6165
logger: Logger,
6266
messages: WithParts[]
6367
): void => {
64-
if (!config.strategies.pruneTool.enabled) {
68+
if (!config.strategies.discardTool.enabled && !config.strategies.extractTool.enabled) {
6569
return
6670
}
6771

@@ -84,7 +88,11 @@ export const insertPruneToolContext = (
8488
logger.debug("prunable-tools: \n" + prunableToolsList)
8589

8690
let nudgeString = ""
87-
if (state.nudgeCounter >= config.strategies.pruneTool.nudge.frequency) {
91+
const nudgeFrequency = Math.min(
92+
config.strategies.discardTool.nudge.frequency,
93+
config.strategies.extractTool.nudge.frequency
94+
)
95+
if (state.nudgeCounter >= nudgeFrequency) {
8896
logger.info("Inserting prune nudge message")
8997
nudgeString = "\n" + NUDGE_STRING
9098
}

‎lib/state/tool-cache.ts‎

Lines changed: 17 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -35,16 +35,27 @@ export async function syncToolCache(
3535
continue
3636
}
3737

38-
const isProtectedByTurn = config.strategies.pruneTool.turnProtection.enabled &&
39-
config.strategies.pruneTool.turnProtection.turns > 0 &&
40-
(state.currentTurn - turnCounter) < config.strategies.pruneTool.turnProtection.turns
38+
const turnProtectionEnabled = config.strategies.discardTool.turnProtection.enabled ||
39+
config.strategies.extractTool.turnProtection.enabled
40+
const turnProtectionTurns = Math.max(
41+
config.strategies.discardTool.turnProtection.turns,
42+
config.strategies.extractTool.turnProtection.turns
43+
)
44+
const isProtectedByTurn = turnProtectionEnabled &&
45+
turnProtectionTurns > 0 &&
46+
(state.currentTurn - turnCounter) < turnProtectionTurns
47+
48+
state.lastToolPrune = part.tool === "discard" || part.tool === "extract"
4149

42-
state.lastToolPrune = part.tool === "prune"
50+
const allProtectedTools = [
51+
...config.strategies.discardTool.protectedTools,
52+
...config.strategies.extractTool.protectedTools
53+
]
4354

44-
if (part.tool === "prune") {
55+
if (part.tool === "discard" || part.tool === "extract") {
4556
state.nudgeCounter = 0
4657
} else if (
47-
!config.strategies.pruneTool.protectedTools.includes(part.tool) &&
58+
!allProtectedTools.includes(part.tool) &&
4859
!isProtectedByTurn
4960
) {
5061
state.nudgeCounter++

‎lib/strategies/index.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
11
export { deduplicate } from "./deduplication"
22
export { runOnIdle } from "./on-idle"
3-
export { createPruneTool } from "./prune-tool"
3+
export { createDiscardTool, createExtractTool } from "./prune-tool"
44
export { supersedeWrites } from "./supersede-writes"

0 commit comments

Comments
 (0)