Merge pull request #152 from Opencode-DCP/prune-write-and-edit

Prune write and edit

Author: jorgenwh

README.md

@@ -102,7 +102,7 @@ DCP uses its own config file:
 ### Protected Tools
 
 By default, these tools are always protected from pruning across all strategies:
-`task`, `todowrite`, `todoread`, `prune`, `batch`, `write`, `edit`
+`task`, `todowrite`, `todoread`, `prune`, `batch`
 
 The `protectedTools` arrays in each strategy add to this default list.
 

lib/config.ts

@@ -40,7 +40,7 @@ export interface PluginConfig {
     }
 }
 
-const DEFAULT_PROTECTED_TOOLS = ['task', 'todowrite', 'todoread', 'prune', 'batch', 'write', 'edit']
+const DEFAULT_PROTECTED_TOOLS = ['task', 'todowrite', 'todoread', 'prune', 'batch']
 
 // Valid config keys for validation against user config
 export const VALID_CONFIG_KEYS = new Set([

lib/messages/prune.ts

@@ -6,6 +6,7 @@ import { extractParameterKey, buildToolIdList } from "./utils"
 import { getLastUserMessage } from "../shared-utils"
 import { UserMessage } from "@opencode-ai/sdk"
 
+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")
 
@@ -39,7 +40,7 @@ const buildPrunableToolsList = (
         return ""
     }
 
-    return `<prunable-tools>\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 outputs. Keep the context free of noise.\n${lines.join('\n')}\n</prunable-tools>`
+    return `<prunable-tools>\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</prunable-tools>`
 }
 
 export const insertPruneToolContext = (
@@ -101,7 +102,7 @@ export const prune = (
     messages: WithParts[]
 ): void => {
     pruneToolOutputs(state, logger, messages)
-    // more prune methods coming here
+    pruneToolInputs(state, logger, messages)
 }
 
 const pruneToolOutputs = (
@@ -117,9 +118,37 @@ const pruneToolOutputs = (
             if (!state.prune.toolIds.includes(part.callID)) {
                 continue
             }
+            // Skip write and edit tools - their inputs are pruned instead
+            if (part.tool === 'write' || part.tool === 'edit') {
+                continue
+            }
             if (part.state.status === 'completed') {
                 part.state.output = PRUNED_TOOL_OUTPUT_REPLACEMENT
             }
         }
     }
 }
+
+const pruneToolInputs = (
+    state: SessionState,
+    logger: Logger,
+    messages: WithParts[]
+): void => {
+    for (const msg of messages) {
+        for (const part of msg.parts) {
+            if (part.type !== 'tool') {
+                continue
+            }
+            if (!state.prune.toolIds.includes(part.callID)) {
+                continue
+            }
+            // Only prune inputs for write and edit tools
+            if (part.tool !== 'write' && part.tool !== 'edit') {
+                continue
+            }
+            if (part.state.input?.content !== undefined) {
+                part.state.input.content = PRUNED_TOOL_INPUT_REPLACEMENT
+            }
+        }
+    }
+}

lib/prompts/synthetic.txt

@@ -17,6 +17,7 @@ You WILL use the `prune` tool when ANY of these are true:
 - You are about to start a new phase of work
 - You have distilled enough information in your messages to prune related tools
 - Context contains tools output that are unhelpful, noise, or made obsolete by newer outputs
+- Write or edit operations are complete (pruning removes the large input content)
 
 You MUST NOT prune when:
 - The tool output will be needed for upcoming implementation work

lib/prompts/tool.txt

@@ -1,8 +1,10 @@
-Prunes tool outputs from context to manage conversation size and reduce noise.
+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 `<prunable-tools>` 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.
 
+**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.
+
 ## CRITICAL: When and How to Prune
 
 You must use this tool in three specific scenarios. The rules for distillation (summarizing findings) differ for each. **You must specify the reason as the first element of the `ids` array** to indicate which scenario applies.
@@ -62,3 +64,9 @@ The tests passed. The feature is verified.
 Assistant: [Reads 'auth.ts' to understand the login flow]
 I've understood the auth flow. I'll need to modify this file to add the new validation, so I'm keeping this read in context rather than distilling and pruning.
 </example_keep>
+
+<example_edit_completion>
+Assistant: [Edits 'auth.ts' to add validation]
+The edit was successful. I no longer need the raw edit content in context.
+[Uses prune with ids: ["completion", "15"]]
+</example_edit_completion>

lib/strategies/utils.ts

@@ -50,13 +50,23 @@ export const calculateTokensSaved = (
                 if (part.type !== 'tool' || !pruneToolIds.includes(part.callID)) {
                     continue
                 }
+                // For write and edit tools, count input content as that is all we prune for these tools
+                // (input is present in both completed and error states)
+                if (part.tool === "write" || part.tool === "edit") {
+                    const inputContent = part.state.input?.content
+                    const content = typeof inputContent === 'string'
+                        ? inputContent
+                        : JSON.stringify(inputContent ?? '')
+                    contents.push(content)
+                    continue
+                }
+                // For other tools, count output or error based on status
                 if (part.state.status === "completed") {
                     const content = typeof part.state.output === 'string'
                         ? part.state.output
                         : JSON.stringify(part.state.output)
                     contents.push(content)
-                }
-                if (part.state.status === "error") {
+                } else if (part.state.status === "error") {
                     const content = typeof part.state.error === 'string'
                         ? part.state.error
                         : JSON.stringify(part.state.error)