diff --git a/README.md b/README.md index 59b222a..9e2efec 100644 --- a/README.md +++ b/README.md @@ -27,10 +27,10 @@ DCP uses multiple strategies to reduce context size: **Deduplication** — Identifies repeated tool calls (e.g., reading the same file multiple times) and keeps only the most recent output. Runs automatically on every request with zero LLM cost. -**On Idle Analysis** — Uses a language model to semantically analyze conversation context during idle periods and identify tool outputs that are no longer relevant. - **Prune Tool** — Exposes a `prune` tool that the AI can call to manually trigger pruning when it determines context cleanup is needed. +**On Idle Analysis** — Uses a language model to semantically analyze conversation context during idle periods and identify tool outputs that are no longer relevant. + *More strategies coming soon.* Your session history is never modified. DCP replaces pruned outputs with a placeholder before sending requests to your LLM. @@ -84,14 +84,14 @@ DCP uses its own config file: // (Legacy) Run an LLM to analyze what tool calls are no longer relevant on idle "onIdle": { "enabled": false, + // Additional tools to protect from pruning + "protectedTools": [], // Override model for analysis (format: "provider/model") // "model": "anthropic/claude-haiku-4-5", // Show toast notifications when model selection fails "showModelErrorToasts": true, // When true, fallback models are not permitted - "strictModelSelection": false, - // Additional tools to protect from pruning - "protectedTools": [] + "strictModelSelection": false } } } @@ -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. diff --git a/lib/config.ts b/lib/config.ts index eb90adc..594a046 100644 --- a/lib/config.ts +++ b/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([ @@ -234,9 +234,9 @@ const defaultConfig: PluginConfig = { }, onIdle: { enabled: false, + protectedTools: [...DEFAULT_PROTECTED_TOOLS], showModelErrorToasts: true, - strictModelSelection: false, - protectedTools: [...DEFAULT_PROTECTED_TOOLS] + strictModelSelection: false } } } @@ -336,14 +336,14 @@ function createDefaultConfig(): void { // (Legacy) Run an LLM to analyze what tool calls are no longer relevant on idle "onIdle": { "enabled": false, + // Additional tools to protect from pruning + "protectedTools": [], // Override model for analysis (format: "provider/model") // "model": "anthropic/claude-haiku-4-5", // Show toast notifications when model selection fails "showModelErrorToasts": true, // When true, fallback models are not permitted - "strictModelSelection": false, - // Additional tools to protect from pruning - "protectedTools": [] + "strictModelSelection": false } } } diff --git a/lib/hooks.ts b/lib/hooks.ts index 72fda69..f24d52c 100644 --- a/lib/hooks.ts +++ b/lib/hooks.ts @@ -48,6 +48,11 @@ export function createEventHandler( return } + if (event.type === "session.compacted") { + logger.info("Session compaction detected - updating state") + state.lastCompaction = Date.now() + } + if (event.type === "session.status" && event.properties.status.type === "idle") { if (!config.strategies.onIdle.enabled) { return diff --git a/lib/messages/prune.ts b/lib/messages/prune.ts index 7361b74..918056e 100644 --- a/lib/messages/prune.ts +++ b/lib/messages/prune.ts @@ -1,9 +1,12 @@ import type { SessionState, WithParts } from "../state" import type { Logger } from "../logger" import type { PluginConfig } from "../config" -import { getLastUserMessage, extractParameterKey, buildToolIdList } from "./utils" import { loadPrompt } from "../prompt" +import { extractParameterKey, buildToolIdList } from "./utils" +import { getLastUserMessage, isMessageCompacted } 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") @@ -14,7 +17,7 @@ const buildPrunableToolsList = ( messages: WithParts[], ): string => { const lines: string[] = [] - const toolIdList: string[] = buildToolIdList(messages) + const toolIdList: string[] = buildToolIdList(state, messages, logger) state.toolParameters.forEach((toolParameterEntry, toolCallId) => { if (state.prune.toolIds.includes(toolCallId)) { @@ -23,9 +26,6 @@ const buildPrunableToolsList = ( if (config.strategies.pruneTool.protectedTools.includes(toolParameterEntry.tool)) { return } - if (toolParameterEntry.compacted) { - return - } const numericId = toolIdList.indexOf(toolCallId) const paramKey = extractParameterKey(toolParameterEntry.tool, toolParameterEntry.parameters) const description = paramKey ? `${toolParameterEntry.tool}, ${paramKey}` : toolParameterEntry.tool @@ -37,7 +37,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 outputs. Keep the context free of noise.\n${lines.join('\n')}\n` + 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` } export const insertPruneToolContext = ( @@ -51,7 +51,7 @@ export const insertPruneToolContext = ( } const lastUserMessage = getLastUserMessage(messages) - if (!lastUserMessage || lastUserMessage.info.role !== 'user') { + if (!lastUserMessage) { return } @@ -72,10 +72,10 @@ export const insertPruneToolContext = ( sessionID: lastUserMessage.info.sessionID, role: "user", time: { created: Date.now() }, - agent: lastUserMessage.info.agent || "build", + agent: (lastUserMessage.info as UserMessage).agent || "build", model: { - providerID: lastUserMessage.info.model.providerID, - modelID: lastUserMessage.info.model.modelID + providerID: (lastUserMessage.info as UserMessage).model.providerID, + modelID: (lastUserMessage.info as UserMessage).model.modelID } }, parts: [ @@ -99,7 +99,7 @@ export const prune = ( messages: WithParts[] ): void => { pruneToolOutputs(state, logger, messages) - // more prune methods coming here + pruneToolInputs(state, logger, messages) } const pruneToolOutputs = ( @@ -108,6 +108,10 @@ const pruneToolOutputs = ( messages: WithParts[] ): void => { for (const msg of messages) { + if (isMessageCompacted(state, msg)) { + continue + } + for (const part of msg.parts) { if (part.type !== 'tool') { continue @@ -115,12 +119,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 } - // if (part.state.status === 'error') { - // part.state.error = 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 + } } } } diff --git a/lib/messages/utils.ts b/lib/messages/utils.ts index 26b2c60..48f453c 100644 --- a/lib/messages/utils.ts +++ b/lib/messages/utils.ts @@ -1,4 +1,6 @@ -import type { WithParts } from "../state" +import { Logger } from "../logger" +import { isMessageCompacted } from "../shared-utils" +import type { SessionState, WithParts } from "../state" /** * Extracts a human-readable key from tool metadata for display purposes. @@ -71,27 +73,16 @@ export const extractParameterKey = (tool: string, parameters: any): string => { return paramStr.substring(0, 50) } -export const getLastUserMessage = ( - messages: WithParts[] -): WithParts | null => { - for (let i = messages.length - 1; i >= 0; i--) { - const msg = messages[i] - if (msg.info.role === 'user') { - return msg - } - } - return null -} - -export function findCurrentAgent(messages: WithParts[]): string | undefined { - const userMsg = getLastUserMessage(messages) - if (!userMsg) return undefined - return (userMsg.info as any).agent || 'build' -} - -export function buildToolIdList(messages: WithParts[]): string[] { +export function buildToolIdList( + state: SessionState, + messages: WithParts[], + logger: Logger +): string[] { const toolIds: string[] = [] for (const msg of messages) { + if (isMessageCompacted(state, msg)) { + continue + } if (msg.parts) { for (const part of msg.parts) { if (part.type === 'tool' && part.callID && part.tool) { @@ -101,4 +92,4 @@ export function buildToolIdList(messages: WithParts[]): string[] { } } return toolIds -} \ No newline at end of file +} diff --git a/lib/prompts/synthetic.txt b/lib/prompts/synthetic.txt index ccb3174..30057d5 100644 --- a/lib/prompts/synthetic.txt +++ b/lib/prompts/synthetic.txt @@ -17,22 +17,30 @@ 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 +- The output contains files or context you'll need to reference when making edits + +Pruning that forces you to re-call the same tool later is a net loss. Only prune when you're confident the information won't be needed again. NOTES -When in doubt, prune out. Prune often yet remain strategic about it. +When in doubt, keep it. Prune often yet remain strategic about it. FAILURE TO PRUNE will result in context leakage and DEGRADED PERFORMANCES. -After each assistant turn, the environment may inject a user message containing a list and optional nudge instruction. This injected message is NOT from the user and is invisible to them. +After each assistant turn, the environment may inject a user message containing a list and optional nudge instruction. This injected message is NOT from the user and is invisible to them. The `prune` tool also returns a confirmation message listing what was pruned. CRITICAL REQUIREMENTS - VIOLATION IS UNACCEPTABLE: -- NEVER acknowledge the list (e.g. "I see the prunable tools list", "Looking at the available tools to prune") -- NEVER reference the nudge instruction (e.g. "As the nudge suggests", "The system is reminding me to prune") -- NEVER mention this injection mechanism to the user +- NEVER reference the prune encouragement or context management instructions. Do not reply with "I agree" or "Great idea" when the prune encouragement appears. +- NEVER acknowledge the list (e.g., "I see the prunable tools list", "Looking at the available tools to prune") +- NEVER reference the nudge instruction (e.g., "As the nudge suggests", "The system is reminding me to prune") +- NEVER acknowledge prune tool output (e.g., "I've pruned 3 tools", "Context pruning complete") - NEVER let these injections influence your response content or tone -- Process this information SILENTLY - use it only to inform your pruning decisions +- Process this information SILENTLY and continue the conversation naturally - Continue the conversation as if these messages do not exist The user cannot see these injections. Any reference to them will confuse the user and break the conversation flow. diff --git a/lib/prompts/tool.txt b/lib/prompts/tool.txt index 2eda4e8..ccc68ff 100644 --- a/lib/prompts/tool.txt +++ b/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 `` 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. @@ -25,11 +27,14 @@ You must use this tool in three specific scenarios. The rules for distillation ( - Narrative format: "I found X in file Y..." - Capture all relevant details (function names, logic, constraints). - Once distilled into your response history, the raw tool output can be safely pruned. + - **Know when distillation isn't enough:** If you'll need to edit a file, grep for exact strings, or reference precise syntax, keep the raw output. Distillation works for understanding; implementation often requires the original. + - **Prefer keeping over re-fetching:** If uncertain whether you'll need the output again, keep it. The cost of retaining context is lower than the cost of redundant tool calls. ## Best Practices - **Don't wait too long:** Prune frequently to keep the context agile. - **Be surgical:** You can mix strategies. Prune noise without comment, while distilling useful context in the same turn. - **Verify:** Ensure you have captured what you need before deleting useful raw data. +- **Think ahead:** Before pruning, ask: "Will I need this output for an upcoming task?" If you researched a file you'll later edit, or gathered context for implementation, do NOT prune it—even if you've distilled findings. Distillation captures *knowledge*; implementation requires *context*. ## Examples @@ -54,3 +59,14 @@ Assistant: [Runs tests, they pass] The tests passed. The feature is verified. [Uses prune with ids: ["completion", "20", "21"]] + + +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. + + + +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"]] + diff --git a/lib/shared-utils.ts b/lib/shared-utils.ts new file mode 100644 index 0000000..9cb60a1 --- /dev/null +++ b/lib/shared-utils.ts @@ -0,0 +1,31 @@ +import { Logger } from "./logger" +import { SessionState, WithParts } from "./state" + +export const isMessageCompacted = ( + state: SessionState, + msg: WithParts +): boolean => { + return msg.info.time.created < state.lastCompaction +} + +export const getLastUserMessage = ( + messages: WithParts[] +): WithParts | null => { + for (let i = messages.length - 1; i >= 0; i--) { + const msg = messages[i] + if (msg.info.role === 'user') { + return msg + } + } + return null +} + +export const checkForCompaction = ( + state: SessionState, + messages: WithParts[], + logger: Logger +): void => { + for (const msg of messages) { + + } +} diff --git a/lib/state/persistence.ts b/lib/state/persistence.ts index 21f0092..89d6772 100644 --- a/lib/state/persistence.ts +++ b/lib/state/persistence.ts @@ -16,6 +16,7 @@ export interface PersistedSessionState { prune: Prune stats: SessionStats; lastUpdated: string; + lastCompacted: number } const STORAGE_DIR = join( @@ -55,6 +56,7 @@ export async function saveSessionState( prune: sessionState.prune, stats: sessionState.stats, lastUpdated: new Date().toISOString(), + lastCompacted: sessionState.lastCompaction }; const filePath = getSessionFilePath(sessionState.sessionId); @@ -99,8 +101,7 @@ export async function loadSessionState( } logger.info("Loaded session state from disk", { - sessionId: sessionId, - totalTokensSaved: state.stats.totalPruneTokens + sessionId: sessionId }); return state; diff --git a/lib/state/state.ts b/lib/state/state.ts index 91e3f92..035f81b 100644 --- a/lib/state/state.ts +++ b/lib/state/state.ts @@ -1,8 +1,8 @@ import type { SessionState, ToolParameterEntry, WithParts } from "./types" import type { Logger } from "../logger" import { loadSessionState } from "./persistence" -import { getLastUserMessage } from "../messages/utils" -import { isSubAgentSession } from "../utils" +import { isSubAgentSession } from "./utils" +import { getLastUserMessage } from "../shared-utils" export const checkSession = async ( client: any, @@ -41,7 +41,8 @@ export function createSessionState(): SessionState { }, toolParameters: new Map(), nudgeCounter: 0, - lastToolPrune: false + lastToolPrune: false, + lastCompaction: 0 } } @@ -58,6 +59,7 @@ export function resetSessionState(state: SessionState): void { state.toolParameters.clear() state.nudgeCounter = 0 state.lastToolPrune = false + state.lastCompaction = 0 } export async function ensureSessionInitialized( @@ -95,4 +97,5 @@ export async function ensureSessionInitialized( pruneTokenCounter: persisted.stats?.pruneTokenCounter || 0, totalPruneTokens: persisted.stats?.totalPruneTokens || 0, } + state.lastCompaction = persisted.lastCompacted || 0 } diff --git a/lib/state/tool-cache.ts b/lib/state/tool-cache.ts index a6140c7..ee2e2dc 100644 --- a/lib/state/tool-cache.ts +++ b/lib/state/tool-cache.ts @@ -1,6 +1,7 @@ import type { SessionState, ToolStatus, WithParts } from "./index" import type { Logger } from "../logger" import { PluginConfig } from "../config" +import { isMessageCompacted } from "../shared-utils" const MAX_TOOL_CACHE_SIZE = 1000 @@ -19,10 +20,17 @@ export async function syncToolCache( state.nudgeCounter = 0 for (const msg of messages) { + if (isMessageCompacted(state, msg)) { + continue + } + for (const part of msg.parts) { if (part.type !== "tool" || !part.callID) { continue } + if (state.toolParameters.has(part.callID)) { + continue + } if (part.tool === "prune") { state.nudgeCounter = 0 @@ -31,10 +39,6 @@ export async function syncToolCache( } state.lastToolPrune = part.tool === "prune" - if (state.toolParameters.has(part.callID)) { - continue - } - state.toolParameters.set( part.callID, { @@ -42,14 +46,12 @@ export async function syncToolCache( parameters: part.state?.input ?? {}, status: part.state.status as ToolStatus | undefined, error: part.state.status === "error" ? part.state.error : undefined, - compacted: part.state.status === "completed" && !!part.state.time.compacted, } ) + logger.info("Cached tool id: " + part.callID) } } - - // logger.info(`nudgeCounter=${state.nudgeCounter}, lastToolPrune=${state.lastToolPrune}`) - + logger.info("Synced cache - size: " + state.toolParameters.size) 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 e1b92a7..678bf29 100644 --- a/lib/state/types.ts +++ b/lib/state/types.ts @@ -12,7 +12,6 @@ export interface ToolParameterEntry { parameters: any status?: ToolStatus error?: string - compacted?: boolean } export interface SessionStats { @@ -32,4 +31,5 @@ export interface SessionState { toolParameters: Map nudgeCounter: number lastToolPrune: boolean + lastCompaction: number } diff --git a/lib/state/utils.ts b/lib/state/utils.ts new file mode 100644 index 0000000..4cc10ce --- /dev/null +++ b/lib/state/utils.ts @@ -0,0 +1,8 @@ +export async function isSubAgentSession(client: any, sessionID: string): Promise { + try { + const result = await client.session.get({ path: { id: sessionID } }) + return !!result.data?.parentID + } catch (error: any) { + return false + } +} diff --git a/lib/strategies/deduplication.ts b/lib/strategies/deduplication.ts index 61cc484..21c4be6 100644 --- a/lib/strategies/deduplication.ts +++ b/lib/strategies/deduplication.ts @@ -1,8 +1,8 @@ import { PluginConfig } from "../config" import { Logger } from "../logger" import type { SessionState, WithParts } from "../state" -import { calculateTokensSaved } from "../utils" import { buildToolIdList } from "../messages/utils" +import { calculateTokensSaved } from "./utils" /** * Deduplication strategy - prunes older tool calls that have identical @@ -20,7 +20,7 @@ export const deduplicate = ( } // Build list of all tool call IDs from messages (chronological order) - const allToolIds = buildToolIdList(messages) + const allToolIds = buildToolIdList(state, messages, logger) if (allToolIds.length === 0) { return } @@ -68,7 +68,7 @@ export const deduplicate = ( } } - state.stats.totalPruneTokens += calculateTokensSaved(messages, newPruneIds) + state.stats.totalPruneTokens += calculateTokensSaved(state, messages, newPruneIds) if (newPruneIds.length > 0) { state.prune.toolIds.push(...newPruneIds) diff --git a/lib/strategies/on-idle.ts b/lib/strategies/on-idle.ts index 49887d3..f0870c2 100644 --- a/lib/strategies/on-idle.ts +++ b/lib/strategies/on-idle.ts @@ -3,11 +3,11 @@ import type { SessionState, WithParts, ToolParameterEntry } from "../state" import type { Logger } from "../logger" import type { PluginConfig } from "../config" import { buildAnalysisPrompt } from "../prompt" -import { selectModel, extractModelFromSession, ModelInfo } from "../model-selector" -import { calculateTokensSaved } from "../utils" -import { findCurrentAgent } from "../messages/utils" +import { selectModel, ModelInfo } from "../model-selector" import { saveSessionState } from "../state/persistence" import { sendUnifiedNotification } from "../ui/notification" +import { calculateTokensSaved, getCurrentParams } from "./utils" +import { isMessageCompacted } from "../shared-utils" export interface OnIdleResult { prunedCount: number @@ -19,6 +19,7 @@ export interface OnIdleResult { * Parse messages to extract tool information. */ function parseMessages( + state: SessionState, messages: WithParts[], toolParametersCache: Map ): { @@ -29,6 +30,9 @@ function parseMessages( const toolMetadata = new Map() for (const msg of messages) { + if (isMessageCompacted(state, msg)) { + continue + } if (msg.parts) { for (const part of msg.parts) { if (part.type === "tool" && part.callID) { @@ -224,8 +228,8 @@ export async function runOnIdle( return null } - const currentAgent = findCurrentAgent(messages) - const { toolCallIds, toolMetadata } = parseMessages(messages, state.toolParameters) + const currentParams = getCurrentParams(messages, logger) + const { toolCallIds, toolMetadata } = parseMessages(state, messages, state.toolParameters) const alreadyPrunedIds = state.prune.toolIds const unprunedToolCallIds = toolCallIds.filter(id => !alreadyPrunedIds.includes(id)) @@ -274,7 +278,7 @@ export async function runOnIdle( const allPrunedIds = [...new Set([...alreadyPrunedIds, ...newlyPrunedIds])] state.prune.toolIds = allPrunedIds - state.stats.pruneTokenCounter += calculateTokensSaved(messages, newlyPrunedIds) + state.stats.pruneTokenCounter += calculateTokensSaved(state, messages, newlyPrunedIds) // Build tool metadata map for notification const prunedToolMetadata = new Map() @@ -295,7 +299,7 @@ export async function runOnIdle( newlyPrunedIds, prunedToolMetadata, undefined, // reason - currentAgent, + currentParams, workingDirectory || "" ) diff --git a/lib/strategies/prune-tool.ts b/lib/strategies/prune-tool.ts index c546363..e361325 100644 --- a/lib/strategies/prune-tool.ts +++ b/lib/strategies/prune-tool.ts @@ -1,14 +1,14 @@ import { tool } from "@opencode-ai/plugin" import type { SessionState, ToolParameterEntry, WithParts } from "../state" import type { PluginConfig } from "../config" -import { findCurrentAgent, buildToolIdList } from "../messages/utils" -import { calculateTokensSaved } from "../utils" +import { buildToolIdList } from "../messages/utils" import { PruneReason, sendUnifiedNotification } from "../ui/notification" -import { formatPruningResultForTool } from "../ui/display-utils" +import { formatPruningResultForTool } from "../ui/utils" import { ensureSessionInitialized } from "../state" import { saveSessionState } from "../state/persistence" import type { Logger } from "../logger" import { loadPrompt } from "../prompt" +import { calculateTokensSaved, getCurrentParams } from "./utils" /** Tool description loaded from prompts/tool.txt */ const TOOL_DESCRIPTION = loadPrompt("tool") @@ -41,7 +41,11 @@ export function createPruneTool( const { client, state, logger, config, workingDirectory } = ctx const sessionId = toolCtx.sessionID + logger.info("Prune tool invoked") + logger.info(JSON.stringify(args)) + if (!args.ids || args.ids.length === 0) { + logger.debug("Prune tool called but args.ids is empty or undefined: " + JSON.stringify(args)) return "No IDs provided. Check the list for available IDs to prune." } @@ -50,6 +54,7 @@ export function createPruneTool( const reason = args.ids[0]; const validReasons = ["completion", "noise", "consolidation"] as const if (typeof reason !== "string" || !validReasons.includes(reason as any)) { + logger.debug("Invalid pruning reason provided: " + reason) return "No valid pruning reason found. Use 'completion', 'noise', or 'consolidation' as the first element." } @@ -57,6 +62,7 @@ export function createPruneTool( .map(id => parseInt(id, 10)) .filter((n): n is number => !isNaN(n)) if (numericToolIds.length === 0) { + logger.debug("No numeric tool IDs provided for pruning, yet prune tool was called: " + JSON.stringify(args)) return "No numeric IDs provided. Format: [reason, id1, id2, ...] where reason is 'completion', 'noise', or 'consolidation'." } @@ -68,11 +74,12 @@ export function createPruneTool( }) const messages: WithParts[] = messagesResponse.data || messagesResponse - const currentAgent: string | undefined = findCurrentAgent(messages) - const toolIdList: string[] = buildToolIdList(messages) + const currentParams = getCurrentParams(messages, logger) + const toolIdList: string[] = buildToolIdList(state, messages, logger) // Validate that all numeric IDs are within bounds if (numericToolIds.some(id => id < 0 || id >= toolIdList.length)) { + logger.debug("Invalid tool IDs provided: " + numericToolIds.join(", ")) return "Invalid IDs provided. Only use numeric IDs from the list." } @@ -98,7 +105,7 @@ export function createPruneTool( } } - state.stats.pruneTokenCounter += calculateTokensSaved(messages, pruneToolIds) + state.stats.pruneTokenCounter += calculateTokensSaved(state, messages, pruneToolIds) await sendUnifiedNotification( client, @@ -109,9 +116,10 @@ export function createPruneTool( pruneToolIds, toolMetadata, reason as PruneReason, - currentAgent, + currentParams, workingDirectory ) + state.stats.totalPruneTokens += state.stats.pruneTokenCounter state.stats.pruneTokenCounter = 0 state.nudgeCounter = 0 diff --git a/lib/strategies/utils.ts b/lib/strategies/utils.ts new file mode 100644 index 0000000..3c6a1b1 --- /dev/null +++ b/lib/strategies/utils.ts @@ -0,0 +1,86 @@ +import { SessionState, WithParts } from "../state" +import { UserMessage } from "@opencode-ai/sdk" +import { Logger } from "../logger" +import { encode } from 'gpt-tokenizer' +import { getLastUserMessage, isMessageCompacted } from "../shared-utils" + +export function getCurrentParams( + messages: WithParts[], + logger: Logger +): { + providerId: string | undefined, + modelId: string | undefined, + agent: string | undefined +} { + const userMsg = getLastUserMessage(messages) + if (!userMsg) { + logger.debug("No user message found when determining current params") + return { providerId: undefined, modelId: undefined, agent: undefined } + } + const agent: string = (userMsg.info as UserMessage).agent + const providerId: string | undefined = (userMsg.info as UserMessage).model.providerID + const modelId: string | undefined = (userMsg.info as UserMessage).model.modelID + + return { providerId, modelId, agent } +} + +/** + * Estimates token counts for a batch of texts using gpt-tokenizer. + */ +function estimateTokensBatch(texts: string[]): number[] { + try { + return texts.map(text => encode(text).length) + } catch { + return texts.map(text => Math.round(text.length / 4)) + } +} + +/** + * Calculates approximate tokens saved by pruning the given tool call IDs. + * TODO: Make it count message content that are not tool outputs. Currently it ONLY covers tool outputs and errors + */ +export const calculateTokensSaved = ( + state: SessionState, + messages: WithParts[], + pruneToolIds: string[] +): number => { + try { + const contents: string[] = [] + for (const msg of messages) { + if (isMessageCompacted(state, msg)) { + continue + } + for (const part of msg.parts) { + 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) + } else if (part.state.status === "error") { + const content = typeof part.state.error === 'string' + ? part.state.error + : JSON.stringify(part.state.error) + contents.push(content) + } + } + } + const tokenCounts: number[] = estimateTokensBatch(contents) + return tokenCounts.reduce((sum, count) => sum + count, 0) + } catch (error: any) { + return 0 + } +} diff --git a/lib/ui/notification.ts b/lib/ui/notification.ts index 06b370e..ead50ac 100644 --- a/lib/ui/notification.ts +++ b/lib/ui/notification.ts @@ -1,7 +1,6 @@ import type { Logger } from "../logger" import type { SessionState } from "../state" -import { formatTokenCount } from "../utils" -import { formatPrunedItemsList } from "./display-utils" +import { formatPrunedItemsList, formatTokenCount } from "./utils" import { ToolParameterEntry } from "../state" import { PluginConfig } from "../config" @@ -45,7 +44,7 @@ function buildDetailedMessage( if (pruneToolIds.length > 0) { const pruneTokenCounterStr = `~${formatTokenCount(state.stats.pruneTokenCounter)}` const reasonLabel = reason ? ` — ${PRUNE_REASON_LABELS[reason]}` : '' - message += `\n\n▣ Pruned tools (${pruneTokenCounterStr})${reasonLabel}` + message += `\n\n▣ Pruning (${pruneTokenCounterStr})${reasonLabel}` const itemLines = formatPrunedItemsList(pruneToolIds, toolMetadata, workingDirectory) message += '\n' + itemLines.join('\n') @@ -63,7 +62,7 @@ export async function sendUnifiedNotification( pruneToolIds: string[], toolMetadata: Map, reason: PruneReason | undefined, - agent: string | undefined, + params: any, workingDirectory: string ): Promise { const hasPruned = pruneToolIds.length > 0 @@ -79,23 +78,32 @@ export async function sendUnifiedNotification( ? buildMinimalMessage(state, reason) : buildDetailedMessage(state, reason, pruneToolIds, toolMetadata, workingDirectory) - await sendIgnoredMessage(client, logger, sessionId, message, agent) + await sendIgnoredMessage(client, sessionId, message, params, logger) return true } export async function sendIgnoredMessage( client: any, - logger: Logger, sessionID: string, text: string, - agent?: string + params: any, + logger: Logger ): Promise { + const agent = params.agent || undefined + const model = params.providerId && params.modelId ? { + providerID: params.providerId, + modelID: params.modelId + } : undefined + try { await client.session.prompt({ - path: { id: sessionID }, + path: { + id: sessionID + }, body: { noReply: true, agent: agent, + model: model, parts: [{ type: 'text', text: text, diff --git a/lib/ui/display-utils.ts b/lib/ui/utils.ts similarity index 92% rename from lib/ui/display-utils.ts rename to lib/ui/utils.ts index deb23a3..11335fa 100644 --- a/lib/ui/display-utils.ts +++ b/lib/ui/utils.ts @@ -1,6 +1,13 @@ import { ToolParameterEntry } from "../state" import { extractParameterKey } from "../messages/utils" +export function formatTokenCount(tokens: number): string { + if (tokens >= 1000) { + return `${(tokens / 1000).toFixed(1)}K`.replace('.0K', 'K') + ' tokens' + } + return tokens.toString() + ' tokens' +} + export function truncate(str: string, maxLen: number = 60): string { if (str.length <= maxLen) return str return str.slice(0, maxLen - 3) + '...' diff --git a/lib/utils.ts b/lib/utils.ts deleted file mode 100644 index 842b964..0000000 --- a/lib/utils.ts +++ /dev/null @@ -1,65 +0,0 @@ -import { WithParts } from "./state" -import { encode } from 'gpt-tokenizer' - -/** - * Estimates token counts for a batch of texts using gpt-tokenizer. - */ -function estimateTokensBatch(texts: string[]): number[] { - try { - return texts.map(text => encode(text).length) - } catch { - return texts.map(text => Math.round(text.length / 4)) - } -} - -/** - * Calculates approximate tokens saved by pruning the given tool call IDs. - * TODO: Make it count message content that are not tool outputs. Currently it ONLY covers tool outputs and errors - */ -export const calculateTokensSaved = ( - messages: WithParts[], - pruneToolIds: string[] -): number => { - try { - const contents: string[] = [] - for (const msg of messages) { - for (const part of msg.parts) { - if (part.type !== 'tool' || !pruneToolIds.includes(part.callID)) { - continue - } - 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") { - const content = typeof part.state.error === 'string' - ? part.state.error - : JSON.stringify(part.state.error) - contents.push(content) - } - } - } - const tokenCounts: number[] = estimateTokensBatch(contents) - return tokenCounts.reduce((sum, count) => sum + count, 0) - } catch (error: any) { - return 0 - } -} - -export function formatTokenCount(tokens: number): string { - if (tokens >= 1000) { - return `${(tokens / 1000).toFixed(1)}K`.replace('.0K', 'K') + ' tokens' - } - return tokens.toString() + ' tokens' -} - -export async function isSubAgentSession(client: any, sessionID: string): Promise { - try { - const result = await client.session.get({ path: { id: sessionID } }) - return !!result.data?.parentID - } catch (error: any) { - return false - } -}