Skip to content

Commit e3ed599

Browse files
TarquinenTarquinen
committed
refactor: split prune.ts into inject.ts and prune.ts for SRP
1 parent a25248e commit e3ed599

File tree

3 files changed

+132
-127
lines changed

3 files changed

+132
-127
lines changed

‎lib/messages/index.ts‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1 +1,2 @@
1-
export { prune, insertPruneToolContext } from "./prune"
1+
export { prune } from "./prune"
2+
export { insertPruneToolContext } from "./inject"

‎lib/messages/inject.ts‎

Lines changed: 129 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,129 @@
1+
import type { SessionState, WithParts } from "../state"
2+
import type { Logger } from "../logger"
3+
import type { PluginConfig } from "../config"
4+
import { loadPrompt } from "../prompt"
5+
import { extractParameterKey, buildToolIdList, createSyntheticUserMessage } from "./utils"
6+
import { getLastUserMessage } from "../shared-utils"
7+
8+
const getNudgeString = (config: PluginConfig): string => {
9+
const discardEnabled = config.tools.discard.enabled
10+
const extractEnabled = config.tools.extract.enabled
11+
12+
if (discardEnabled && extractEnabled) {
13+
return loadPrompt(`user/nudge/nudge-both`)
14+
} else if (discardEnabled) {
15+
return loadPrompt(`user/nudge/nudge-discard`)
16+
} else if (extractEnabled) {
17+
return loadPrompt(`user/nudge/nudge-extract`)
18+
}
19+
return ""
20+
}
21+
22+
const wrapPrunableTools = (content: string): string => `<prunable-tools>
23+
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. Consolidate your prunes for efficiency; it is rarely worth pruning a single tiny tool output. Keep the context free of noise.
24+
${content}
25+
</prunable-tools>`
26+
27+
const getCooldownMessage = (config: PluginConfig): string => {
28+
const discardEnabled = config.tools.discard.enabled
29+
const extractEnabled = config.tools.extract.enabled
30+
31+
let toolName: string
32+
if (discardEnabled && extractEnabled) {
33+
toolName = "discard or extract tools"
34+
} else if (discardEnabled) {
35+
toolName = "discard tool"
36+
} else {
37+
toolName = "extract tool"
38+
}
39+
40+
return `<prunable-tools>
41+
Context management was just performed. Do not use the ${toolName} again. A fresh list will be available after your next tool use.
42+
</prunable-tools>`
43+
}
44+
45+
const buildPrunableToolsList = (
46+
state: SessionState,
47+
config: PluginConfig,
48+
logger: Logger,
49+
messages: WithParts[],
50+
): string => {
51+
const lines: string[] = []
52+
const toolIdList: string[] = buildToolIdList(state, messages, logger)
53+
54+
state.toolParameters.forEach((toolParameterEntry, toolCallId) => {
55+
if (state.prune.toolIds.includes(toolCallId)) {
56+
return
57+
}
58+
59+
const allProtectedTools = config.tools.settings.protectedTools
60+
if (allProtectedTools.includes(toolParameterEntry.tool)) {
61+
return
62+
}
63+
64+
const numericId = toolIdList.indexOf(toolCallId)
65+
if (numericId === -1) {
66+
logger.warn(`Tool in cache but not in toolIdList - possible stale entry`, {
67+
toolCallId,
68+
tool: toolParameterEntry.tool,
69+
})
70+
return
71+
}
72+
const paramKey = extractParameterKey(toolParameterEntry.tool, toolParameterEntry.parameters)
73+
const description = paramKey
74+
? `${toolParameterEntry.tool}, ${paramKey}`
75+
: toolParameterEntry.tool
76+
lines.push(`${numericId}: ${description}`)
77+
logger.debug(
78+
`Prunable tool found - ID: ${numericId}, Tool: ${toolParameterEntry.tool}, Call ID: ${toolCallId}`,
79+
)
80+
})
81+
82+
if (lines.length === 0) {
83+
return ""
84+
}
85+
86+
return wrapPrunableTools(lines.join("\n"))
87+
}
88+
89+
export const insertPruneToolContext = (
90+
state: SessionState,
91+
config: PluginConfig,
92+
logger: Logger,
93+
messages: WithParts[],
94+
): void => {
95+
if (!config.tools.discard.enabled && !config.tools.extract.enabled) {
96+
return
97+
}
98+
99+
let prunableToolsContent: string
100+
101+
if (state.lastToolPrune) {
102+
logger.debug("Last tool was prune - injecting cooldown message")
103+
prunableToolsContent = getCooldownMessage(config)
104+
} else {
105+
const prunableToolsList = buildPrunableToolsList(state, config, logger, messages)
106+
if (!prunableToolsList) {
107+
return
108+
}
109+
110+
logger.debug("prunable-tools: \n" + prunableToolsList)
111+
112+
let nudgeString = ""
113+
if (
114+
config.tools.settings.nudgeEnabled &&
115+
state.nudgeCounter >= config.tools.settings.nudgeFrequency
116+
) {
117+
logger.info("Inserting prune nudge message")
118+
nudgeString = "\n" + getNudgeString(config)
119+
}
120+
121+
prunableToolsContent = prunableToolsList + nudgeString
122+
}
123+
124+
const lastUserMessage = getLastUserMessage(messages)
125+
if (!lastUserMessage) {
126+
return
127+
}
128+
messages.push(createSyntheticUserMessage(lastUserMessage, prunableToolsContent))
129+
}

‎lib/messages/prune.ts‎

Lines changed: 1 addition & 126 deletions
Original file line numberDiff line numberDiff line change
@@ -1,139 +1,14 @@
11
import type { SessionState, WithParts } from "../state"
22
import type { Logger } from "../logger"
33
import type { PluginConfig } from "../config"
4-
import { loadPrompt } from "../prompt"
5-
import { extractParameterKey, buildToolIdList, createSyntheticUserMessage } from "./utils"
6-
import { getLastUserMessage, isMessageCompacted } from "../shared-utils"
4+
import { isMessageCompacted } from "../shared-utils"
75

86
const PRUNED_TOOL_INPUT_REPLACEMENT =
97
"[content removed to save context, this is not what was written to the file, but a placeholder]"
108
const PRUNED_TOOL_OUTPUT_REPLACEMENT =
119
"[Output removed to save context - information superseded or no longer needed]"
1210
const PRUNED_TOOL_ERROR_INPUT_REPLACEMENT = "[input removed due to failed tool call]"
1311

14-
const getNudgeString = (config: PluginConfig): string => {
15-
const discardEnabled = config.tools.discard.enabled
16-
const extractEnabled = config.tools.extract.enabled
17-
18-
if (discardEnabled && extractEnabled) {
19-
return loadPrompt(`user/nudge/nudge-both`)
20-
} else if (discardEnabled) {
21-
return loadPrompt(`user/nudge/nudge-discard`)
22-
} else if (extractEnabled) {
23-
return loadPrompt(`user/nudge/nudge-extract`)
24-
}
25-
return ""
26-
}
27-
28-
const wrapPrunableTools = (content: string): string => `<prunable-tools>
29-
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. Consolidate your prunes for efficiency; it is rarely worth pruning a single tiny tool output. Keep the context free of noise.
30-
${content}
31-
</prunable-tools>`
32-
33-
const getCooldownMessage = (config: PluginConfig): string => {
34-
const discardEnabled = config.tools.discard.enabled
35-
const extractEnabled = config.tools.extract.enabled
36-
37-
let toolName: string
38-
if (discardEnabled && extractEnabled) {
39-
toolName = "discard or extract tools"
40-
} else if (discardEnabled) {
41-
toolName = "discard tool"
42-
} else {
43-
toolName = "extract tool"
44-
}
45-
46-
return `<prunable-tools>
47-
Context management was just performed. Do not use the ${toolName} again. A fresh list will be available after your next tool use.
48-
</prunable-tools>`
49-
}
50-
51-
const buildPrunableToolsList = (
52-
state: SessionState,
53-
config: PluginConfig,
54-
logger: Logger,
55-
messages: WithParts[],
56-
): string => {
57-
const lines: string[] = []
58-
const toolIdList: string[] = buildToolIdList(state, messages, logger)
59-
60-
state.toolParameters.forEach((toolParameterEntry, toolCallId) => {
61-
if (state.prune.toolIds.includes(toolCallId)) {
62-
return
63-
}
64-
65-
const allProtectedTools = config.tools.settings.protectedTools
66-
if (allProtectedTools.includes(toolParameterEntry.tool)) {
67-
return
68-
}
69-
70-
const numericId = toolIdList.indexOf(toolCallId)
71-
if (numericId === -1) {
72-
logger.warn(`Tool in cache but not in toolIdList - possible stale entry`, {
73-
toolCallId,
74-
tool: toolParameterEntry.tool,
75-
})
76-
return
77-
}
78-
const paramKey = extractParameterKey(toolParameterEntry.tool, toolParameterEntry.parameters)
79-
const description = paramKey
80-
? `${toolParameterEntry.tool}, ${paramKey}`
81-
: toolParameterEntry.tool
82-
lines.push(`${numericId}: ${description}`)
83-
logger.debug(
84-
`Prunable tool found - ID: ${numericId}, Tool: ${toolParameterEntry.tool}, Call ID: ${toolCallId}`,
85-
)
86-
})
87-
88-
if (lines.length === 0) {
89-
return ""
90-
}
91-
92-
return wrapPrunableTools(lines.join("\n"))
93-
}
94-
95-
export const insertPruneToolContext = (
96-
state: SessionState,
97-
config: PluginConfig,
98-
logger: Logger,
99-
messages: WithParts[],
100-
): void => {
101-
if (!config.tools.discard.enabled && !config.tools.extract.enabled) {
102-
return
103-
}
104-
105-
let prunableToolsContent: string
106-
107-
if (state.lastToolPrune) {
108-
logger.debug("Last tool was prune - injecting cooldown message")
109-
prunableToolsContent = getCooldownMessage(config)
110-
} else {
111-
const prunableToolsList = buildPrunableToolsList(state, config, logger, messages)
112-
if (!prunableToolsList) {
113-
return
114-
}
115-
116-
logger.debug("prunable-tools: \n" + prunableToolsList)
117-
118-
let nudgeString = ""
119-
if (
120-
config.tools.settings.nudgeEnabled &&
121-
state.nudgeCounter >= config.tools.settings.nudgeFrequency
122-
) {
123-
logger.info("Inserting prune nudge message")
124-
nudgeString = "\n" + getNudgeString(config)
125-
}
126-
127-
prunableToolsContent = prunableToolsList + nudgeString
128-
}
129-
130-
const lastUserMessage = getLastUserMessage(messages)
131-
if (!lastUserMessage) {
132-
return
133-
}
134-
messages.push(createSyntheticUserMessage(lastUserMessage, prunableToolsContent))
135-
}
136-
13712
export const prune = (
13813
state: SessionState,
13914
logger: Logger,

0 commit comments

Comments
 (0)