diff --git a/README.md b/README.md index 9bfd061..9d15e73 100644 --- a/README.md +++ b/README.md @@ -77,6 +77,11 @@ DCP uses its own config file: "enabled": true, // Additional tools to protect from pruning "protectedTools": [], + // Protect from pruning for message turns + "turnProtection": { + "enabled": false, + "turns": 4 + }, // Nudge the LLM to use the prune tool (every tool results) "nudge": { "enabled": true, diff --git a/lib/config.ts b/lib/config.ts index e1dbd1a..9a670fe 100644 --- a/lib/config.ts +++ b/lib/config.ts @@ -22,9 +22,15 @@ export interface PruneToolNudge { frequency: number } +export interface PruneToolTurnProtection { + enabled: boolean + turns: number +} + export interface PruneTool { enabled: boolean protectedTools: string[] + turnProtection: PruneToolTurnProtection nudge: PruneToolNudge } @@ -72,6 +78,9 @@ export const VALID_CONFIG_KEYS = new Set([ 'strategies.pruneTool', 'strategies.pruneTool.enabled', 'strategies.pruneTool.protectedTools', + 'strategies.pruneTool.turnProtection', + 'strategies.pruneTool.turnProtection.enabled', + 'strategies.pruneTool.turnProtection.turns', 'strategies.pruneTool.nudge', 'strategies.pruneTool.nudge.enabled', 'strategies.pruneTool.nudge.frequency' @@ -158,6 +167,14 @@ function validateConfigTypes(config: Record): ValidationError[] { if (strategies.pruneTool.protectedTools !== undefined && !Array.isArray(strategies.pruneTool.protectedTools)) { errors.push({ key: 'strategies.pruneTool.protectedTools', expected: 'string[]', actual: typeof strategies.pruneTool.protectedTools }) } + if (strategies.pruneTool.turnProtection) { + if (strategies.pruneTool.turnProtection.enabled !== undefined && typeof strategies.pruneTool.turnProtection.enabled !== 'boolean') { + errors.push({ key: 'strategies.pruneTool.turnProtection.enabled', expected: 'boolean', actual: typeof strategies.pruneTool.turnProtection.enabled }) + } + if (strategies.pruneTool.turnProtection.turns !== undefined && typeof strategies.pruneTool.turnProtection.turns !== 'number') { + errors.push({ key: 'strategies.pruneTool.turnProtection.turns', expected: 'number', actual: typeof strategies.pruneTool.turnProtection.turns }) + } + } if (strategies.pruneTool.nudge) { if (strategies.pruneTool.nudge.enabled !== undefined && typeof strategies.pruneTool.nudge.enabled !== 'boolean') { errors.push({ key: 'strategies.pruneTool.nudge.enabled', expected: 'boolean', actual: typeof strategies.pruneTool.nudge.enabled }) @@ -240,6 +257,10 @@ const defaultConfig: PluginConfig = { pruneTool: { enabled: true, protectedTools: [...DEFAULT_PROTECTED_TOOLS], + turnProtection: { + enabled: false, + turns: 4 + }, nudge: { enabled: true, frequency: 10 @@ -337,12 +358,17 @@ function createDefaultConfig(): void { "enabled": true }, // Exposes a prune tool to your LLM to call when it determines pruning is necessary - "pruneTool": { - "enabled": true, + \"pruneTool\": { + \"enabled\": true, // Additional tools to protect from pruning - "protectedTools": [], + \"protectedTools\": [], + // Protect from pruning for message turns + \"turnProtection\": { + \"enabled\": false, + \"turns\": 4 + }, // Nudge the LLM to use the prune tool (every tool results) - "nudge": { + \"nudge\": { "enabled": true, "frequency": 10 } @@ -426,6 +452,10 @@ function mergeStrategies( ...(override.pruneTool?.protectedTools ?? []) ]) ], + turnProtection: { + enabled: override.pruneTool?.turnProtection?.enabled ?? base.pruneTool.turnProtection.enabled, + turns: override.pruneTool?.turnProtection?.turns ?? base.pruneTool.turnProtection.turns + }, nudge: { enabled: override.pruneTool?.nudge?.enabled ?? base.pruneTool.nudge.enabled, frequency: override.pruneTool?.nudge?.frequency ?? base.pruneTool.nudge.frequency @@ -452,6 +482,7 @@ function deepCloneConfig(config: PluginConfig): PluginConfig { pruneTool: { ...config.strategies.pruneTool, protectedTools: [...config.strategies.pruneTool.protectedTools], + turnProtection: { ...config.strategies.pruneTool.turnProtection }, nudge: { ...config.strategies.pruneTool.nudge } }, supersedeWrites: { diff --git a/lib/messages/prune.ts b/lib/messages/prune.ts index 33d9a7a..8759bcf 100644 --- a/lib/messages/prune.ts +++ b/lib/messages/prune.ts @@ -10,6 +10,17 @@ const PRUNED_TOOL_INPUT_REPLACEMENT = '[Input removed to save context]' const PRUNED_TOOL_OUTPUT_REPLACEMENT = '[Output removed to save context - information superseded or no longer needed]' const NUDGE_STRING = loadPrompt("nudge") +const wrapPrunableTools = (content: string): string => ` +The following tools have been invoked and are available for pruning. This list does not mandate immediate action. Consider your current goals and the resources you need before discarding valuable tool inputs or outputs. Keep the context free of noise. +${content} +` +const PRUNABLE_TOOLS_COOLDOWN = ` +Pruning was just performed. Do not use the prune tool again. A fresh list will be available after your next tool use. +` + +const SYNTHETIC_MESSAGE_ID = "msg_01234567890123456789012345" +const SYNTHETIC_PART_ID = "prt_01234567890123456789012345" + const buildPrunableToolsList = ( state: SessionState, config: PluginConfig, @@ -41,7 +52,7 @@ const buildPrunableToolsList = ( return "" } - return `\nThe following tools have been invoked and are available for pruning. This list does not mandate immediate action. Consider your current goals and the resources you need before discarding valuable tool inputs or outputs. Keep the context free of noise.\n${lines.join('\n')}\n` + return wrapPrunableTools(lines.join('\n')) } export const insertPruneToolContext = ( @@ -59,20 +70,31 @@ export const insertPruneToolContext = ( return } - const prunableToolsList = buildPrunableToolsList(state, config, logger, messages) - if (!prunableToolsList) { - return - } + let prunableToolsContent: string + + if (state.lastToolPrune) { + logger.debug("Last tool was prune - injecting cooldown message") + prunableToolsContent = PRUNABLE_TOOLS_COOLDOWN + } else { + const prunableToolsList = buildPrunableToolsList(state, config, logger, messages) + if (!prunableToolsList) { + return + } + + logger.debug("prunable-tools: \n" + prunableToolsList) + + let nudgeString = "" + if (state.nudgeCounter >= config.strategies.pruneTool.nudge.frequency) { + logger.info("Inserting prune nudge message") + nudgeString = "\n" + NUDGE_STRING + } - let nudgeString = "" - if (state.nudgeCounter >= config.strategies.pruneTool.nudge.frequency) { - logger.info("Inserting prune nudge message") - nudgeString = "\n" + NUDGE_STRING + prunableToolsContent = prunableToolsList + nudgeString } const userMessage: WithParts = { info: { - id: "msg_01234567890123456789012345", + id: SYNTHETIC_MESSAGE_ID, sessionID: lastUserMessage.info.sessionID, role: "user", time: { created: Date.now() }, @@ -84,11 +106,11 @@ export const insertPruneToolContext = ( }, parts: [ { - id: "prt_01234567890123456789012345", + id: SYNTHETIC_PART_ID, sessionID: lastUserMessage.info.sessionID, - messageID: "msg_01234567890123456789012345", + messageID: SYNTHETIC_MESSAGE_ID, type: "text", - text: prunableToolsList + nudgeString, + text: prunableToolsContent, } ] } diff --git a/lib/prompts/synthetic.txt b/lib/prompts/synthetic.txt index 30057d5..001e9b1 100644 --- a/lib/prompts/synthetic.txt +++ b/lib/prompts/synthetic.txt @@ -28,6 +28,7 @@ Pruning that forces you to re-call the same tool later is a net loss. Only prune NOTES When in doubt, keep it. Prune often yet remain strategic about it. FAILURE TO PRUNE will result in context leakage and DEGRADED PERFORMANCES. +There may be tools in session context that do not appear in the list, this is expected, you can ONLY prune what you see in . diff --git a/lib/prompts/tool.txt b/lib/prompts/tool.txt index ccc68ff..a7982b4 100644 --- a/lib/prompts/tool.txt +++ b/lib/prompts/tool.txt @@ -1,7 +1,7 @@ Prunes tool outputs from context to manage conversation size and reduce noise. For `write` and `edit` tools, the input content is pruned instead of the output. ## IMPORTANT: The Prunable List -A `` list is injected into user messages showing available tool outputs you can prune. Each line has the format `ID: tool, parameter` (e.g., `20: read, /path/to/file.ts`). You MUST only use numeric IDs that appear in this list to select which tools to prune. +A `` list is injected into user messages showing available tool outputs you can prune when there are tools available for pruning. Each line has the format `ID: tool, parameter` (e.g., `20: read, /path/to/file.ts`). You MUST only use numeric IDs that appear in this list to select which tools to prune. **Note:** For `write` and `edit` tools, pruning removes the input content (the code being written/edited) while preserving the output confirmation. This is useful after completing a file modification when you no longer need the raw content in context. diff --git a/lib/state/state.ts b/lib/state/state.ts index caab6d9..e33f4c8 100644 --- a/lib/state/state.ts +++ b/lib/state/state.ts @@ -2,7 +2,7 @@ import type { SessionState, ToolParameterEntry, WithParts } from "./types" import type { Logger } from "../logger" import { loadSessionState } from "./persistence" import { isSubAgentSession } from "./utils" -import { getLastUserMessage } from "../shared-utils" +import { getLastUserMessage, isMessageCompacted } from "../shared-utils" export const checkSession = async ( client: any, @@ -34,6 +34,8 @@ export const checkSession = async ( state.prune.toolIds = [] logger.info("Detected compaction from messages - cleared tool cache", { timestamp: lastCompactionTimestamp }) } + + state.currentTurn = countTurns(state, messages) } export function createSessionState(): SessionState { @@ -50,7 +52,8 @@ export function createSessionState(): SessionState { toolParameters: new Map(), nudgeCounter: 0, lastToolPrune: false, - lastCompaction: 0 + lastCompaction: 0, + currentTurn: 0 } } @@ -68,6 +71,7 @@ export function resetSessionState(state: SessionState): void { state.nudgeCounter = 0 state.lastToolPrune = false state.lastCompaction = 0 + state.currentTurn = 0 } export async function ensureSessionInitialized( @@ -92,6 +96,7 @@ export async function ensureSessionInitialized( logger.info("isSubAgent = " + isSubAgent) state.lastCompaction = findLastCompactionTimestamp(messages) + state.currentTurn = countTurns(state, messages) const persisted = await loadSessionState(sessionId, logger) if (persisted === null) { @@ -116,3 +121,18 @@ function findLastCompactionTimestamp(messages: WithParts[]): number { } return 0 } + +export function countTurns(state: SessionState, messages: WithParts[]): number { + let turnCount = 0 + for (const msg of messages) { + if (isMessageCompacted(state, msg)) { + continue + } + for (const part of msg.parts) { + if (part.type === "step-start") { + turnCount++ + } + } + } + return turnCount +} diff --git a/lib/state/tool-cache.ts b/lib/state/tool-cache.ts index ee2e2dc..f8ad2b3 100644 --- a/lib/state/tool-cache.ts +++ b/lib/state/tool-cache.ts @@ -18,6 +18,7 @@ export async function syncToolCache( logger.info("Syncing tool parameters from OpenCode messages") state.nudgeCounter = 0 + let turnCounter = 0 for (const msg of messages) { if (isMessageCompacted(state, msg)) { @@ -25,19 +26,37 @@ export async function syncToolCache( } for (const part of msg.parts) { - if (part.type !== "tool" || !part.callID) { + if (part.type === "step-start") { + turnCounter++ continue } - if (state.toolParameters.has(part.callID)) { + + if (part.type !== "tool" || !part.callID) { continue } + const isProtectedByTurn = config.strategies.pruneTool.turnProtection.enabled && + config.strategies.pruneTool.turnProtection.turns > 0 && + (state.currentTurn - turnCounter) < config.strategies.pruneTool.turnProtection.turns + + state.lastToolPrune = part.tool === "prune" + if (part.tool === "prune") { state.nudgeCounter = 0 - } else if (!config.strategies.pruneTool.protectedTools.includes(part.tool)) { + } else if ( + !config.strategies.pruneTool.protectedTools.includes(part.tool) && + !isProtectedByTurn + ) { state.nudgeCounter++ } - state.lastToolPrune = part.tool === "prune" + + if (state.toolParameters.has(part.callID)) { + continue + } + + if (isProtectedByTurn) { + continue + } state.toolParameters.set( part.callID, @@ -46,12 +65,14 @@ export async function syncToolCache( parameters: part.state?.input ?? {}, status: part.state.status as ToolStatus | undefined, error: part.state.status === "error" ? part.state.error : undefined, + turn: turnCounter, } ) - logger.info("Cached tool id: " + part.callID) + logger.info(`Cached tool id: ${part.callID} (created on turn ${turnCounter})`) } } - logger.info("Synced cache - size: " + state.toolParameters.size) + + logger.info(`Synced cache - size: ${state.toolParameters.size}, currentTurn: ${state.currentTurn}, nudgeCounter: ${state.nudgeCounter}`) trimToolParametersCache(state) } catch (error) { logger.warn("Failed to sync tool parameters from OpenCode", { diff --git a/lib/state/types.ts b/lib/state/types.ts index 678bf29..04847d5 100644 --- a/lib/state/types.ts +++ b/lib/state/types.ts @@ -12,6 +12,7 @@ export interface ToolParameterEntry { parameters: any status?: ToolStatus error?: string + turn: number // Which turn (step-start count) this tool was called on } export interface SessionStats { @@ -32,4 +33,5 @@ export interface SessionState { nudgeCounter: number lastToolPrune: boolean lastCompaction: number + currentTurn: number // Current turn count derived from step-start parts } diff --git a/lib/strategies/deduplication.ts b/lib/strategies/deduplication.ts index 21c4be6..11462a8 100644 --- a/lib/strategies/deduplication.ts +++ b/lib/strategies/deduplication.ts @@ -41,7 +41,7 @@ export const deduplicate = ( for (const id of unprunedIds) { const metadata = state.toolParameters.get(id) if (!metadata) { - logger.warn(`Missing metadata for tool call ID: ${id}`) + // logger.warn(`Missing metadata for tool call ID: ${id}`) continue } diff --git a/lib/strategies/on-idle.ts b/lib/strategies/on-idle.ts index f0870c2..602298e 100644 --- a/lib/strategies/on-idle.ts +++ b/lib/strategies/on-idle.ts @@ -45,7 +45,8 @@ function parseMessages( tool: part.tool, parameters: parameters, status: part.state?.status, - error: part.state?.status === "error" ? part.state.error : undefined + error: part.state?.status === "error" ? part.state.error : undefined, + turn: cachedData?.turn ?? 0 }) } } diff --git a/lib/strategies/prune-tool.ts b/lib/strategies/prune-tool.ts index 6cf9052..e037f3f 100644 --- a/lib/strategies/prune-tool.ts +++ b/lib/strategies/prune-tool.ts @@ -83,11 +83,17 @@ export function createPruneTool( return "Invalid IDs provided. Only use numeric IDs from the list." } - // Check for protected tools (model hallucinated an ID not in the prunable list) + // Validate that all IDs exist in cache and aren't protected + // (rejects hallucinated IDs and turn-protected tools not shown in ) for (const index of numericToolIds) { const id = toolIdList[index] const metadata = state.toolParameters.get(id) - if (metadata && config.strategies.pruneTool.protectedTools.includes(metadata.tool)) { + if (!metadata) { + logger.debug("Rejecting prune request - ID not in cache (turn-protected or hallucinated)", { index, id }) + return "Invalid IDs provided. Only use numeric IDs from the list." + } + if (config.strategies.pruneTool.protectedTools.includes(metadata.tool)) { + logger.debug("Rejecting prune request - protected tool", { index, id, tool: metadata.tool }) return "Invalid IDs provided. Only use numeric IDs from the list." } }