Skip to content

Commit 3d926eb

Browse files
TarquinenTarquinen
authored
Merge pull request #167 from Opencode-DCP/feat/turn-based-tool-protection
Feat/turn based tool protection
2 parents a6c28e2 + f4ae605 commit 3d926eb

File tree

11 files changed

+139
-30
lines changed

11 files changed

+139
-30
lines changed

‎README.md‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,11 @@ DCP uses its own config file:
7777
"enabled": true,
7878
// Additional tools to protect from pruning
7979
"protectedTools": [],
80+
// Protect from pruning for <turn protection> message turns
81+
"turnProtection": {
82+
"enabled": false,
83+
"turns": 4
84+
},
8085
// Nudge the LLM to use the prune tool (every <frequency> tool results)
8186
"nudge": {
8287
"enabled": true,

‎lib/config.ts‎

Lines changed: 35 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -22,9 +22,15 @@ export interface PruneToolNudge {
2222
frequency: number
2323
}
2424

25+
export interface PruneToolTurnProtection {
26+
enabled: boolean
27+
turns: number
28+
}
29+
2530
export interface PruneTool {
2631
enabled: boolean
2732
protectedTools: string[]
33+
turnProtection: PruneToolTurnProtection
2834
nudge: PruneToolNudge
2935
}
3036

@@ -72,6 +78,9 @@ export const VALID_CONFIG_KEYS = new Set([
7278
'strategies.pruneTool',
7379
'strategies.pruneTool.enabled',
7480
'strategies.pruneTool.protectedTools',
81+
'strategies.pruneTool.turnProtection',
82+
'strategies.pruneTool.turnProtection.enabled',
83+
'strategies.pruneTool.turnProtection.turns',
7584
'strategies.pruneTool.nudge',
7685
'strategies.pruneTool.nudge.enabled',
7786
'strategies.pruneTool.nudge.frequency'
@@ -158,6 +167,14 @@ function validateConfigTypes(config: Record<string, any>): ValidationError[] {
158167
if (strategies.pruneTool.protectedTools !== undefined && !Array.isArray(strategies.pruneTool.protectedTools)) {
159168
errors.push({ key: 'strategies.pruneTool.protectedTools', expected: 'string[]', actual: typeof strategies.pruneTool.protectedTools })
160169
}
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 })
173+
}
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 })
176+
}
177+
}
161178
if (strategies.pruneTool.nudge) {
162179
if (strategies.pruneTool.nudge.enabled !== undefined && typeof strategies.pruneTool.nudge.enabled !== 'boolean') {
163180
errors.push({ key: 'strategies.pruneTool.nudge.enabled', expected: 'boolean', actual: typeof strategies.pruneTool.nudge.enabled })
@@ -240,6 +257,10 @@ const defaultConfig: PluginConfig = {
240257
pruneTool: {
241258
enabled: true,
242259
protectedTools: [...DEFAULT_PROTECTED_TOOLS],
260+
turnProtection: {
261+
enabled: false,
262+
turns: 4
263+
},
243264
nudge: {
244265
enabled: true,
245266
frequency: 10
@@ -337,12 +358,17 @@ function createDefaultConfig(): void {
337358
"enabled": true
338359
},
339360
// Exposes a prune tool to your LLM to call when it determines pruning is necessary
340-
"pruneTool": {
341-
"enabled": true,
361+
\"pruneTool\": {
362+
\"enabled\": true,
342363
// Additional tools to protect from pruning
343-
"protectedTools": [],
364+
\"protectedTools\": [],
365+
// Protect from pruning for <turn protection> message turns
366+
\"turnProtection\": {
367+
\"enabled\": false,
368+
\"turns\": 4
369+
},
344370
// Nudge the LLM to use the prune tool (every <frequency> tool results)
345-
"nudge": {
371+
\"nudge\": {
346372
"enabled": true,
347373
"frequency": 10
348374
}
@@ -426,6 +452,10 @@ function mergeStrategies(
426452
...(override.pruneTool?.protectedTools ?? [])
427453
])
428454
],
455+
turnProtection: {
456+
enabled: override.pruneTool?.turnProtection?.enabled ?? base.pruneTool.turnProtection.enabled,
457+
turns: override.pruneTool?.turnProtection?.turns ?? base.pruneTool.turnProtection.turns
458+
},
429459
nudge: {
430460
enabled: override.pruneTool?.nudge?.enabled ?? base.pruneTool.nudge.enabled,
431461
frequency: override.pruneTool?.nudge?.frequency ?? base.pruneTool.nudge.frequency
@@ -452,6 +482,7 @@ function deepCloneConfig(config: PluginConfig): PluginConfig {
452482
pruneTool: {
453483
...config.strategies.pruneTool,
454484
protectedTools: [...config.strategies.pruneTool.protectedTools],
485+
turnProtection: { ...config.strategies.pruneTool.turnProtection },
455486
nudge: { ...config.strategies.pruneTool.nudge }
456487
},
457488
supersedeWrites: {

‎lib/messages/prune.ts‎

Lines changed: 35 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,17 @@ const PRUNED_TOOL_INPUT_REPLACEMENT = '[Input removed to save context]'
1010
const PRUNED_TOOL_OUTPUT_REPLACEMENT = '[Output removed to save context - information superseded or no longer needed]'
1111
const NUDGE_STRING = loadPrompt("nudge")
1212

13+
const wrapPrunableTools = (content: string): string => `<prunable-tools>
14+
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.
15+
${content}
16+
</prunable-tools>`
17+
const PRUNABLE_TOOLS_COOLDOWN = `<prunable-tools>
18+
Pruning was just performed. Do not use the prune tool again. A fresh list will be available after your next tool use.
19+
</prunable-tools>`
20+
21+
const SYNTHETIC_MESSAGE_ID = "msg_01234567890123456789012345"
22+
const SYNTHETIC_PART_ID = "prt_01234567890123456789012345"
23+
1324
const buildPrunableToolsList = (
1425
state: SessionState,
1526
config: PluginConfig,
@@ -41,7 +52,7 @@ const buildPrunableToolsList = (
4152
return ""
4253
}
4354

44-
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>`
55+
return wrapPrunableTools(lines.join('\n'))
4556
}
4657

4758
export const insertPruneToolContext = (
@@ -59,20 +70,31 @@ export const insertPruneToolContext = (
5970
return
6071
}
6172

62-
const prunableToolsList = buildPrunableToolsList(state, config, logger, messages)
63-
if (!prunableToolsList) {
64-
return
65-
}
73+
let prunableToolsContent: string
74+
75+
if (state.lastToolPrune) {
76+
logger.debug("Last tool was prune - injecting cooldown message")
77+
prunableToolsContent = PRUNABLE_TOOLS_COOLDOWN
78+
} else {
79+
const prunableToolsList = buildPrunableToolsList(state, config, logger, messages)
80+
if (!prunableToolsList) {
81+
return
82+
}
83+
84+
logger.debug("prunable-tools: \n" + prunableToolsList)
85+
86+
let nudgeString = ""
87+
if (state.nudgeCounter >= config.strategies.pruneTool.nudge.frequency) {
88+
logger.info("Inserting prune nudge message")
89+
nudgeString = "\n" + NUDGE_STRING
90+
}
6691

67-
let nudgeString = ""
68-
if (state.nudgeCounter >= config.strategies.pruneTool.nudge.frequency) {
69-
logger.info("Inserting prune nudge message")
70-
nudgeString = "\n" + NUDGE_STRING
92+
prunableToolsContent = prunableToolsList + nudgeString
7193
}
7294

7395
const userMessage: WithParts = {
7496
info: {
75-
id: "msg_01234567890123456789012345",
97+
id: SYNTHETIC_MESSAGE_ID,
7698
sessionID: lastUserMessage.info.sessionID,
7799
role: "user",
78100
time: { created: Date.now() },
@@ -84,11 +106,11 @@ export const insertPruneToolContext = (
84106
},
85107
parts: [
86108
{
87-
id: "prt_01234567890123456789012345",
109+
id: SYNTHETIC_PART_ID,
88110
sessionID: lastUserMessage.info.sessionID,
89-
messageID: "msg_01234567890123456789012345",
111+
messageID: SYNTHETIC_MESSAGE_ID,
90112
type: "text",
91-
text: prunableToolsList + nudgeString,
113+
text: prunableToolsContent,
92114
}
93115
]
94116
}

‎lib/prompts/synthetic.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@ Pruning that forces you to re-call the same tool later is a net loss. Only prune
2828
NOTES
2929
When in doubt, keep it. Prune often yet remain strategic about it.
3030
FAILURE TO PRUNE will result in context leakage and DEGRADED PERFORMANCES.
31+
There may be tools in session context that do not appear in the <prunable-tools> list, this is expected, you can ONLY prune what you see in <prunable-tools>.
3132

3233
</instruction>
3334

‎lib/prompts/tool.txt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
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.
22

33
## IMPORTANT: The Prunable List
4-
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.
4+
A `<prunable-tools>` 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.
55

66
**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.
77

‎lib/state/state.ts‎

Lines changed: 22 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ import type { SessionState, ToolParameterEntry, WithParts } from "./types"
22
import type { Logger } from "../logger"
33
import { loadSessionState } from "./persistence"
44
import { isSubAgentSession } from "./utils"
5-
import { getLastUserMessage } from "../shared-utils"
5+
import { getLastUserMessage, isMessageCompacted } from "../shared-utils"
66

77
export const checkSession = async (
88
client: any,
@@ -34,6 +34,8 @@ export const checkSession = async (
3434
state.prune.toolIds = []
3535
logger.info("Detected compaction from messages - cleared tool cache", { timestamp: lastCompactionTimestamp })
3636
}
37+
38+
state.currentTurn = countTurns(state, messages)
3739
}
3840

3941
export function createSessionState(): SessionState {
@@ -50,7 +52,8 @@ export function createSessionState(): SessionState {
5052
toolParameters: new Map<string, ToolParameterEntry>(),
5153
nudgeCounter: 0,
5254
lastToolPrune: false,
53-
lastCompaction: 0
55+
lastCompaction: 0,
56+
currentTurn: 0
5457
}
5558
}
5659

@@ -68,6 +71,7 @@ export function resetSessionState(state: SessionState): void {
6871
state.nudgeCounter = 0
6972
state.lastToolPrune = false
7073
state.lastCompaction = 0
74+
state.currentTurn = 0
7175
}
7276

7377
export async function ensureSessionInitialized(
@@ -92,6 +96,7 @@ export async function ensureSessionInitialized(
9296
logger.info("isSubAgent = " + isSubAgent)
9397

9498
state.lastCompaction = findLastCompactionTimestamp(messages)
99+
state.currentTurn = countTurns(state, messages)
95100

96101
const persisted = await loadSessionState(sessionId, logger)
97102
if (persisted === null) {
@@ -116,3 +121,18 @@ function findLastCompactionTimestamp(messages: WithParts[]): number {
116121
}
117122
return 0
118123
}
124+
125+
export function countTurns(state: SessionState, messages: WithParts[]): number {
126+
let turnCount = 0
127+
for (const msg of messages) {
128+
if (isMessageCompacted(state, msg)) {
129+
continue
130+
}
131+
for (const part of msg.parts) {
132+
if (part.type === "step-start") {
133+
turnCount++
134+
}
135+
}
136+
}
137+
return turnCount
138+
}

‎lib/state/tool-cache.ts‎

Lines changed: 27 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -18,26 +18,45 @@ export async function syncToolCache(
1818
logger.info("Syncing tool parameters from OpenCode messages")
1919

2020
state.nudgeCounter = 0
21+
let turnCounter = 0
2122

2223
for (const msg of messages) {
2324
if (isMessageCompacted(state, msg)) {
2425
continue
2526
}
2627

2728
for (const part of msg.parts) {
28-
if (part.type !== "tool" || !part.callID) {
29+
if (part.type === "step-start") {
30+
turnCounter++
2931
continue
3032
}
31-
if (state.toolParameters.has(part.callID)) {
33+
34+
if (part.type !== "tool" || !part.callID) {
3235
continue
3336
}
3437

38+
const isProtectedByTurn = config.strategies.pruneTool.turnProtection.enabled &&
39+
config.strategies.pruneTool.turnProtection.turns > 0 &&
40+
(state.currentTurn - turnCounter) < config.strategies.pruneTool.turnProtection.turns
41+
42+
state.lastToolPrune = part.tool === "prune"
43+
3544
if (part.tool === "prune") {
3645
state.nudgeCounter = 0
37-
} else if (!config.strategies.pruneTool.protectedTools.includes(part.tool)) {
46+
} else if (
47+
!config.strategies.pruneTool.protectedTools.includes(part.tool) &&
48+
!isProtectedByTurn
49+
) {
3850
state.nudgeCounter++
3951
}
40-
state.lastToolPrune = part.tool === "prune"
52+
53+
if (state.toolParameters.has(part.callID)) {
54+
continue
55+
}
56+
57+
if (isProtectedByTurn) {
58+
continue
59+
}
4160

4261
state.toolParameters.set(
4362
part.callID,
@@ -46,12 +65,14 @@ export async function syncToolCache(
4665
parameters: part.state?.input ?? {},
4766
status: part.state.status as ToolStatus | undefined,
4867
error: part.state.status === "error" ? part.state.error : undefined,
68+
turn: turnCounter,
4969
}
5070
)
51-
logger.info("Cached tool id: " + part.callID)
71+
logger.info(`Cached tool id: ${part.callID} (created on turn ${turnCounter})`)
5272
}
5373
}
54-
logger.info("Synced cache - size: " + state.toolParameters.size)
74+
75+
logger.info(`Synced cache - size: ${state.toolParameters.size}, currentTurn: ${state.currentTurn}, nudgeCounter: ${state.nudgeCounter}`)
5576
trimToolParametersCache(state)
5677
} catch (error) {
5778
logger.warn("Failed to sync tool parameters from OpenCode", {

‎lib/state/types.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ export interface ToolParameterEntry {
1212
parameters: any
1313
status?: ToolStatus
1414
error?: string
15+
turn: number // Which turn (step-start count) this tool was called on
1516
}
1617

1718
export interface SessionStats {
@@ -32,4 +33,5 @@ export interface SessionState {
3233
nudgeCounter: number
3334
lastToolPrune: boolean
3435
lastCompaction: number
36+
currentTurn: number // Current turn count derived from step-start parts
3537
}

‎lib/strategies/deduplication.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,7 @@ export const deduplicate = (
4141
for (const id of unprunedIds) {
4242
const metadata = state.toolParameters.get(id)
4343
if (!metadata) {
44-
logger.warn(`Missing metadata for tool call ID: ${id}`)
44+
// logger.warn(`Missing metadata for tool call ID: ${id}`)
4545
continue
4646
}
4747

‎lib/strategies/on-idle.ts‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,8 @@ function parseMessages(
4545
tool: part.tool,
4646
parameters: parameters,
4747
status: part.state?.status,
48-
error: part.state?.status === "error" ? part.state.error : undefined
48+
error: part.state?.status === "error" ? part.state.error : undefined,
49+
turn: cachedData?.turn ?? 0
4950
})
5051
}
5152
}

0 commit comments

Comments
 (0)