From e3ed5996e8e4341433e1f2a464a89e3bdbc87c5d Mon Sep 17 00:00:00 2001 From: Daniel Smolsky Date: Mon, 29 Dec 2025 17:44:28 -0500 Subject: [PATCH] refactor: split prune.ts into inject.ts and prune.ts for SRP --- lib/messages/index.ts | 3 +- lib/messages/inject.ts | 129 +++++++++++++++++++++++++++++++++++++++++ lib/messages/prune.ts | 127 +--------------------------------------- 3 files changed, 132 insertions(+), 127 deletions(-) create mode 100644 lib/messages/inject.ts diff --git a/lib/messages/index.ts b/lib/messages/index.ts index e8540038..7c9cc89c 100644 --- a/lib/messages/index.ts +++ b/lib/messages/index.ts @@ -1 +1,2 @@ -export { prune, insertPruneToolContext } from "./prune" +export { prune } from "./prune" +export { insertPruneToolContext } from "./inject" diff --git a/lib/messages/inject.ts b/lib/messages/inject.ts new file mode 100644 index 00000000..03b75ec2 --- /dev/null +++ b/lib/messages/inject.ts @@ -0,0 +1,129 @@ +import type { SessionState, WithParts } from "../state" +import type { Logger } from "../logger" +import type { PluginConfig } from "../config" +import { loadPrompt } from "../prompt" +import { extractParameterKey, buildToolIdList, createSyntheticUserMessage } from "./utils" +import { getLastUserMessage } from "../shared-utils" + +const getNudgeString = (config: PluginConfig): string => { + const discardEnabled = config.tools.discard.enabled + const extractEnabled = config.tools.extract.enabled + + if (discardEnabled && extractEnabled) { + return loadPrompt(`user/nudge/nudge-both`) + } else if (discardEnabled) { + return loadPrompt(`user/nudge/nudge-discard`) + } else if (extractEnabled) { + return loadPrompt(`user/nudge/nudge-extract`) + } + return "" +} + +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. Consolidate your prunes for efficiency; it is rarely worth pruning a single tiny tool output. Keep the context free of noise. +${content} +` + +const getCooldownMessage = (config: PluginConfig): string => { + const discardEnabled = config.tools.discard.enabled + const extractEnabled = config.tools.extract.enabled + + let toolName: string + if (discardEnabled && extractEnabled) { + toolName = "discard or extract tools" + } else if (discardEnabled) { + toolName = "discard tool" + } else { + toolName = "extract tool" + } + + return ` +Context management was just performed. Do not use the ${toolName} again. A fresh list will be available after your next tool use. +` +} + +const buildPrunableToolsList = ( + state: SessionState, + config: PluginConfig, + logger: Logger, + messages: WithParts[], +): string => { + const lines: string[] = [] + const toolIdList: string[] = buildToolIdList(state, messages, logger) + + state.toolParameters.forEach((toolParameterEntry, toolCallId) => { + if (state.prune.toolIds.includes(toolCallId)) { + return + } + + const allProtectedTools = config.tools.settings.protectedTools + if (allProtectedTools.includes(toolParameterEntry.tool)) { + return + } + + const numericId = toolIdList.indexOf(toolCallId) + if (numericId === -1) { + logger.warn(`Tool in cache but not in toolIdList - possible stale entry`, { + toolCallId, + tool: toolParameterEntry.tool, + }) + return + } + const paramKey = extractParameterKey(toolParameterEntry.tool, toolParameterEntry.parameters) + const description = paramKey + ? `${toolParameterEntry.tool}, ${paramKey}` + : toolParameterEntry.tool + lines.push(`${numericId}: ${description}`) + logger.debug( + `Prunable tool found - ID: ${numericId}, Tool: ${toolParameterEntry.tool}, Call ID: ${toolCallId}`, + ) + }) + + if (lines.length === 0) { + return "" + } + + return wrapPrunableTools(lines.join("\n")) +} + +export const insertPruneToolContext = ( + state: SessionState, + config: PluginConfig, + logger: Logger, + messages: WithParts[], +): void => { + if (!config.tools.discard.enabled && !config.tools.extract.enabled) { + return + } + + let prunableToolsContent: string + + if (state.lastToolPrune) { + logger.debug("Last tool was prune - injecting cooldown message") + prunableToolsContent = getCooldownMessage(config) + } else { + const prunableToolsList = buildPrunableToolsList(state, config, logger, messages) + if (!prunableToolsList) { + return + } + + logger.debug("prunable-tools: \n" + prunableToolsList) + + let nudgeString = "" + if ( + config.tools.settings.nudgeEnabled && + state.nudgeCounter >= config.tools.settings.nudgeFrequency + ) { + logger.info("Inserting prune nudge message") + nudgeString = "\n" + getNudgeString(config) + } + + prunableToolsContent = prunableToolsList + nudgeString + } + + const lastUserMessage = getLastUserMessage(messages) + if (!lastUserMessage) { + return + } + messages.push(createSyntheticUserMessage(lastUserMessage, prunableToolsContent)) +} diff --git a/lib/messages/prune.ts b/lib/messages/prune.ts index 351be092..c8c1d7d4 100644 --- a/lib/messages/prune.ts +++ b/lib/messages/prune.ts @@ -1,9 +1,7 @@ import type { SessionState, WithParts } from "../state" import type { Logger } from "../logger" import type { PluginConfig } from "../config" -import { loadPrompt } from "../prompt" -import { extractParameterKey, buildToolIdList, createSyntheticUserMessage } from "./utils" -import { getLastUserMessage, isMessageCompacted } from "../shared-utils" +import { isMessageCompacted } from "../shared-utils" const PRUNED_TOOL_INPUT_REPLACEMENT = "[content removed to save context, this is not what was written to the file, but a placeholder]" @@ -11,129 +9,6 @@ const PRUNED_TOOL_OUTPUT_REPLACEMENT = "[Output removed to save context - information superseded or no longer needed]" const PRUNED_TOOL_ERROR_INPUT_REPLACEMENT = "[input removed due to failed tool call]" -const getNudgeString = (config: PluginConfig): string => { - const discardEnabled = config.tools.discard.enabled - const extractEnabled = config.tools.extract.enabled - - if (discardEnabled && extractEnabled) { - return loadPrompt(`user/nudge/nudge-both`) - } else if (discardEnabled) { - return loadPrompt(`user/nudge/nudge-discard`) - } else if (extractEnabled) { - return loadPrompt(`user/nudge/nudge-extract`) - } - return "" -} - -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. Consolidate your prunes for efficiency; it is rarely worth pruning a single tiny tool output. Keep the context free of noise. -${content} -` - -const getCooldownMessage = (config: PluginConfig): string => { - const discardEnabled = config.tools.discard.enabled - const extractEnabled = config.tools.extract.enabled - - let toolName: string - if (discardEnabled && extractEnabled) { - toolName = "discard or extract tools" - } else if (discardEnabled) { - toolName = "discard tool" - } else { - toolName = "extract tool" - } - - return ` -Context management was just performed. Do not use the ${toolName} again. A fresh list will be available after your next tool use. -` -} - -const buildPrunableToolsList = ( - state: SessionState, - config: PluginConfig, - logger: Logger, - messages: WithParts[], -): string => { - const lines: string[] = [] - const toolIdList: string[] = buildToolIdList(state, messages, logger) - - state.toolParameters.forEach((toolParameterEntry, toolCallId) => { - if (state.prune.toolIds.includes(toolCallId)) { - return - } - - const allProtectedTools = config.tools.settings.protectedTools - if (allProtectedTools.includes(toolParameterEntry.tool)) { - return - } - - const numericId = toolIdList.indexOf(toolCallId) - if (numericId === -1) { - logger.warn(`Tool in cache but not in toolIdList - possible stale entry`, { - toolCallId, - tool: toolParameterEntry.tool, - }) - return - } - const paramKey = extractParameterKey(toolParameterEntry.tool, toolParameterEntry.parameters) - const description = paramKey - ? `${toolParameterEntry.tool}, ${paramKey}` - : toolParameterEntry.tool - lines.push(`${numericId}: ${description}`) - logger.debug( - `Prunable tool found - ID: ${numericId}, Tool: ${toolParameterEntry.tool}, Call ID: ${toolCallId}`, - ) - }) - - if (lines.length === 0) { - return "" - } - - return wrapPrunableTools(lines.join("\n")) -} - -export const insertPruneToolContext = ( - state: SessionState, - config: PluginConfig, - logger: Logger, - messages: WithParts[], -): void => { - if (!config.tools.discard.enabled && !config.tools.extract.enabled) { - return - } - - let prunableToolsContent: string - - if (state.lastToolPrune) { - logger.debug("Last tool was prune - injecting cooldown message") - prunableToolsContent = getCooldownMessage(config) - } else { - const prunableToolsList = buildPrunableToolsList(state, config, logger, messages) - if (!prunableToolsList) { - return - } - - logger.debug("prunable-tools: \n" + prunableToolsList) - - let nudgeString = "" - if ( - config.tools.settings.nudgeEnabled && - state.nudgeCounter >= config.tools.settings.nudgeFrequency - ) { - logger.info("Inserting prune nudge message") - nudgeString = "\n" + getNudgeString(config) - } - - prunableToolsContent = prunableToolsList + nudgeString - } - - const lastUserMessage = getLastUserMessage(messages) - if (!lastUserMessage) { - return - } - messages.push(createSyntheticUserMessage(lastUserMessage, prunableToolsContent)) -} - export const prune = ( state: SessionState, logger: Logger,