diff --git a/index.ts b/index.ts index 2372d2c..ebdb868 100644 --- a/index.ts +++ b/index.ts @@ -3,7 +3,7 @@ import { getConfig } from "./lib/config" import { Logger } from "./lib/logger" import { loadPrompt } from "./lib/prompts" import { createSessionState } from "./lib/state" -import { createDiscardTool, createExtractTool } from "./lib/strategies" +import { createDiscardTool, createExtractTool, createPinTool } from "./lib/strategies" import { createChatMessageTransformHandler } from "./lib/hooks" const plugin: Plugin = (async (ctx) => { @@ -43,9 +43,16 @@ const plugin: Plugin = (async (ctx) => { const discardEnabled = config.tools.discard.enabled const extractEnabled = config.tools.extract.enabled + const pinEnabled = config.tools.pin.enabled + const pinningModeEnabled = config.tools.pinningMode.enabled let promptName: string - if (discardEnabled && extractEnabled) { + if (pinningModeEnabled || pinEnabled) { + // Pinning mode: use pin prompt (with optional extract) + promptName = extractEnabled + ? "user/system/system-prompt-pin-extract" + : "user/system/system-prompt-pin" + } else if (discardEnabled && extractEnabled) { promptName = "user/system/system-prompt-both" } else if (discardEnabled) { promptName = "user/system/system-prompt-discard" @@ -63,6 +70,7 @@ const plugin: Plugin = (async (ctx) => { state, logger, config, + ctx.directory, ), "chat.message": async ( input: { @@ -80,8 +88,20 @@ const plugin: Plugin = (async (ctx) => { logger.debug("Cached variant from chat.message hook", { variant: input.variant }) }, tool: { - ...(config.tools.discard.enabled && { - discard: createDiscardTool({ + // Discard tool only available in non-pinning mode + ...(!config.tools.pinningMode.enabled && + config.tools.discard.enabled && { + discard: createDiscardTool({ + client: ctx.client, + state, + logger, + config, + workingDirectory: ctx.directory, + }), + }), + // Extract tool available in both modes + ...(config.tools.extract.enabled && { + extract: createExtractTool({ client: ctx.client, state, logger, @@ -89,8 +109,9 @@ const plugin: Plugin = (async (ctx) => { workingDirectory: ctx.directory, }), }), - ...(config.tools.extract.enabled && { - extract: createExtractTool({ + // Pin tool only available in pinning mode + ...(config.tools.pin.enabled && { + pin: createPinTool({ client: ctx.client, state, logger, @@ -103,7 +124,12 @@ const plugin: Plugin = (async (ctx) => { // Add enabled tools to primary_tools by mutating the opencode config // This works because config is cached and passed by reference const toolsToAdd: string[] = [] - if (config.tools.discard.enabled) toolsToAdd.push("discard") + // In pinning mode, add pin instead of discard + if (config.tools.pinningMode.enabled) { + if (config.tools.pin.enabled) toolsToAdd.push("pin") + } else { + if (config.tools.discard.enabled) toolsToAdd.push("discard") + } if (config.tools.extract.enabled) toolsToAdd.push("extract") if (toolsToAdd.length > 0) { diff --git a/lib/config.ts b/lib/config.ts index 2b48263..65c637e 100644 --- a/lib/config.ts +++ b/lib/config.ts @@ -18,6 +18,17 @@ export interface ExtractTool { showDistillation: boolean } +export interface PinTool { + enabled: boolean +} + +export interface PinningMode { + enabled: boolean + pruneFrequency: number // Auto-prune every N turns + pinDuration: number // Pins expire after M turns + warningTurns: number // Warn N turns before auto-prune +} + export interface ToolSettings { nudgeEnabled: boolean nudgeFrequency: number @@ -28,6 +39,8 @@ export interface Tools { settings: ToolSettings discard: DiscardTool extract: ExtractTool + pin: PinTool + pinningMode: PinningMode } export interface SupersedeWrites { @@ -89,6 +102,13 @@ export const VALID_CONFIG_KEYS = new Set([ "tools.extract", "tools.extract.enabled", "tools.extract.showDistillation", + "tools.pin", + "tools.pin.enabled", + "tools.pinningMode", + "tools.pinningMode.enabled", + "tools.pinningMode.pruneFrequency", + "tools.pinningMode.pinDuration", + "tools.pinningMode.warningTurns", "strategies", // strategies.deduplication "strategies.deduplication", @@ -238,6 +258,57 @@ function validateConfigTypes(config: Record): ValidationError[] { }) } } + if (tools.pin) { + if (tools.pin.enabled !== undefined && typeof tools.pin.enabled !== "boolean") { + errors.push({ + key: "tools.pin.enabled", + expected: "boolean", + actual: typeof tools.pin.enabled, + }) + } + } + if (tools.pinningMode) { + if ( + tools.pinningMode.enabled !== undefined && + typeof tools.pinningMode.enabled !== "boolean" + ) { + errors.push({ + key: "tools.pinningMode.enabled", + expected: "boolean", + actual: typeof tools.pinningMode.enabled, + }) + } + if ( + tools.pinningMode.pruneFrequency !== undefined && + typeof tools.pinningMode.pruneFrequency !== "number" + ) { + errors.push({ + key: "tools.pinningMode.pruneFrequency", + expected: "number", + actual: typeof tools.pinningMode.pruneFrequency, + }) + } + if ( + tools.pinningMode.pinDuration !== undefined && + typeof tools.pinningMode.pinDuration !== "number" + ) { + errors.push({ + key: "tools.pinningMode.pinDuration", + expected: "number", + actual: typeof tools.pinningMode.pinDuration, + }) + } + if ( + tools.pinningMode.warningTurns !== undefined && + typeof tools.pinningMode.warningTurns !== "number" + ) { + errors.push({ + key: "tools.pinningMode.warningTurns", + expected: "number", + actual: typeof tools.pinningMode.warningTurns, + }) + } + } } // Strategies validators @@ -384,6 +455,15 @@ const defaultConfig: PluginConfig = { enabled: true, showDistillation: false, }, + pin: { + enabled: false, + }, + pinningMode: { + enabled: false, + pruneFrequency: 10, + pinDuration: 8, + warningTurns: 2, + }, }, strategies: { deduplication: { @@ -499,6 +579,20 @@ function createDefaultConfig(): void { "enabled": true, // Show distillation content as an ignored message notification "showDistillation": false + }, + // Pin tool to protect context from auto-pruning (used with pinningMode) + "pin": { + "enabled": false + }, + // Pinning mode: auto-prune everything not pinned every N turns + "pinningMode": { + "enabled": false, + // Auto-prune every N turns + "pruneFrequency": 10, + // Pins expire after M turns + "pinDuration": 8, + // Warn N turns before auto-prune + "warningTurns": 2 } }, // Automatic pruning strategies @@ -608,6 +702,15 @@ function mergeTools( enabled: override.extract?.enabled ?? base.extract.enabled, showDistillation: override.extract?.showDistillation ?? base.extract.showDistillation, }, + pin: { + enabled: override.pin?.enabled ?? base.pin.enabled, + }, + pinningMode: { + enabled: override.pinningMode?.enabled ?? base.pinningMode.enabled, + pruneFrequency: override.pinningMode?.pruneFrequency ?? base.pinningMode.pruneFrequency, + pinDuration: override.pinningMode?.pinDuration ?? base.pinningMode.pinDuration, + warningTurns: override.pinningMode?.warningTurns ?? base.pinningMode.warningTurns, + }, } } @@ -622,6 +725,8 @@ function deepCloneConfig(config: PluginConfig): PluginConfig { }, discard: { ...config.tools.discard }, extract: { ...config.tools.extract }, + pin: { ...config.tools.pin }, + pinningMode: { ...config.tools.pinningMode }, }, strategies: { deduplication: { diff --git a/lib/hooks.ts b/lib/hooks.ts index e70c892..85e4837 100644 --- a/lib/hooks.ts +++ b/lib/hooks.ts @@ -2,7 +2,7 @@ import type { SessionState, WithParts } from "./state" import type { Logger } from "./logger" import type { PluginConfig } from "./config" import { syncToolCache } from "./state/tool-cache" -import { deduplicate, supersedeWrites, purgeErrors } from "./strategies" +import { deduplicate, supersedeWrites, purgeErrors, autoPrune } from "./strategies" import { prune, insertPruneToolContext } from "./messages" import { checkSession } from "./state" @@ -11,6 +11,7 @@ export function createChatMessageTransformHandler( state: SessionState, logger: Logger, config: PluginConfig, + workingDirectory: string, ) { return async (input: {}, output: { messages: WithParts[] }) => { await checkSession(client, state, logger, output.messages) @@ -25,6 +26,11 @@ export function createChatMessageTransformHandler( supersedeWrites(state, logger, config, output.messages) purgeErrors(state, logger, config, output.messages) + // Run auto-prune if pinning mode is enabled + if (config.tools.pinningMode.enabled) { + await autoPrune(client, state, logger, config, output.messages, workingDirectory) + } + prune(state, logger, config, output.messages) insertPruneToolContext(state, config, logger, output.messages) diff --git a/lib/messages/inject.ts b/lib/messages/inject.ts index 11153b7..063a1c1 100644 --- a/lib/messages/inject.ts +++ b/lib/messages/inject.ts @@ -5,6 +5,7 @@ import type { UserMessage } from "@opencode-ai/sdk/v2" import { loadPrompt } from "../prompts" import { extractParameterKey, buildToolIdList, createSyntheticUserMessage } from "./utils" import { getLastUserMessage } from "../shared-utils" +import { turnsUntilAutoPrune, shouldShowAutoPruneWarning } from "../strategies/auto-prune" const getNudgeString = (config: PluginConfig): string => { const discardEnabled = config.tools.discard.enabled @@ -28,9 +29,12 @@ ${content} const getCooldownMessage = (config: PluginConfig): string => { const discardEnabled = config.tools.discard.enabled const extractEnabled = config.tools.extract.enabled + const pinEnabled = config.tools.pin.enabled let toolName: string - if (discardEnabled && extractEnabled) { + if (pinEnabled) { + toolName = extractEnabled ? "pin or extract tools" : "pin tool" + } else if (discardEnabled && extractEnabled) { toolName = "discard or extract tools" } else if (discardEnabled) { toolName = "discard tool" @@ -43,6 +47,16 @@ Context management was just performed. Do not use the ${toolName} again. A fresh ` } +const getAutoPruneWarningMessage = (state: SessionState, config: PluginConfig): string => { + const turns = turnsUntilAutoPrune(state, config) + const pinnedCount = state.pins.size + + return ` +Auto-prune in ${turns} turn(s). All unpinned tool outputs will be discarded. +Currently ${pinnedCount} tool(s) pinned. Use the \`pin\` tool NOW to preserve any context you need. +` +} + const buildPrunableToolsList = ( state: SessionState, config: PluginConfig, @@ -51,6 +65,7 @@ const buildPrunableToolsList = ( ): string => { const lines: string[] = [] const toolIdList: string[] = buildToolIdList(state, messages, logger) + const pinningModeEnabled = config.tools.pinningMode.enabled state.toolParameters.forEach((toolParameterEntry, toolCallId) => { if (state.prune.toolIds.includes(toolCallId)) { @@ -74,7 +89,18 @@ const buildPrunableToolsList = ( const description = paramKey ? `${toolParameterEntry.tool}, ${paramKey}` : toolParameterEntry.tool - lines.push(`${numericId}: ${description}`) + + // Show pin status in pinning mode + let line = `${numericId}: ${description}` + if (pinningModeEnabled) { + const pin = state.pins.get(toolCallId) + if (pin) { + const turnsRemaining = pin.expiresAtTurn - state.currentTurn + line += ` [PINNED, expires in ${turnsRemaining} turn(s)]` + } + } + + lines.push(line) logger.debug( `Prunable tool found - ID: ${numericId}, Tool: ${toolParameterEntry.tool}, Call ID: ${toolCallId}`, ) @@ -93,7 +119,12 @@ export const insertPruneToolContext = ( logger: Logger, messages: WithParts[], ): void => { - if (!config.tools.discard.enabled && !config.tools.extract.enabled) { + const pinEnabled = config.tools.pin.enabled + const discardEnabled = config.tools.discard.enabled + const extractEnabled = config.tools.extract.enabled + + // Need at least one pruning tool enabled + if (!pinEnabled && !discardEnabled && !extractEnabled) { return } @@ -111,7 +142,9 @@ export const insertPruneToolContext = ( logger.debug("prunable-tools: \n" + prunableToolsList) let nudgeString = "" + // Only show nudge in non-pinning mode if ( + !config.tools.pinningMode.enabled && config.tools.settings.nudgeEnabled && state.nudgeCounter >= config.tools.settings.nudgeFrequency ) { @@ -119,7 +152,14 @@ export const insertPruneToolContext = ( nudgeString = "\n" + getNudgeString(config) } - prunableToolsContent = prunableToolsList + nudgeString + // Add auto-prune warning if approaching prune cycle + let warningString = "" + if (shouldShowAutoPruneWarning(state, config)) { + logger.info("Inserting auto-prune warning message") + warningString = "\n" + getAutoPruneWarningMessage(state, config) + } + + prunableToolsContent = prunableToolsList + nudgeString + warningString } const lastUserMessage = getLastUserMessage(messages) diff --git a/lib/prompts/pin-tool-spec.txt b/lib/prompts/pin-tool-spec.txt new file mode 100644 index 0000000..56d5700 --- /dev/null +++ b/lib/prompts/pin-tool-spec.txt @@ -0,0 +1,39 @@ +Pin tool outputs to protect them from automatic pruning. + +## IMPORTANT: The Prunable List +A `` list is provided to you showing available tool outputs you can pin. 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. + +## When to Use This Tool + +Use `pin` to preserve context you'll need for upcoming work: + +- **Active work:** Files you're about to edit or reference +- **Key context:** Important information needed for the current task +- **Dependencies:** Code or data that informs your next steps + +## When NOT to Use This Tool + +- **Completed work:** If you're done with the context and won't need it again +- **Noise:** Irrelevant or superseded outputs - let them be auto-pruned + +## Pin Expiration + +Pins expire automatically after a set number of turns. You don't need to unpin - just let pins expire naturally. If you pin something that's already pinned, its expiration will be extended. + +## Best Practices + +- **Pin sparingly:** Only pin what you actively need - unpinned outputs will be discarded at the next auto-prune cycle +- **Think ahead:** Before the auto-prune warning, consider what context you'll need for upcoming work +- **Let go:** When work is complete, let the context be auto-pruned rather than pinning everything + +## Format + +- `ids`: Array of numeric IDs as strings from the `` list + +## Example + +``` +Assistant: [Sees auto-prune warning, reviews prunable tools] +I need to keep the auth service file since I'm about to modify it. +[Uses pin with ids: ["10", "15"]] +``` diff --git a/lib/prompts/user/system/system-prompt-pin-extract.txt b/lib/prompts/user/system/system-prompt-pin-extract.txt new file mode 100644 index 0000000..45caf1a --- /dev/null +++ b/lib/prompts/user/system/system-prompt-pin-extract.txt @@ -0,0 +1,50 @@ + + +ENVIRONMENT +You are operating in a context-constrained environment with automatic pruning. A `` list is injected showing tool outputs available to manage. Use `pin` to protect context from auto-pruning, or `extract` to distill and preserve key findings. + +AUTO-PRUNE CYCLE +Every few turns, all unpinned tool outputs are automatically discarded. You'll receive an `` before this happens, giving you a chance to pin or extract important context. + +TWO TOOLS FOR CONTEXT MANAGEMENT +- `pin`: Protect tool outputs from auto-pruning. Pins expire automatically after several turns. +- `extract`: Distill key findings into preserved knowledge before removal. Use when you need to preserve information permanently. + +CHOOSING THE RIGHT TOOL +Ask: "Will I need this output for upcoming work?" +- **Yes, raw output needed** → `pin` it before auto-prune +- **Yes, but only key details** → `extract` to preserve distilled knowledge +- **No** → Let it be auto-pruned + +Common scenarios: +- Files you're about to edit → `pin` +- Context needed for current task → `pin` +- Valuable insights but raw output too large → `extract` +- Completed work, no longer needed → Let auto-prune +- Noise, irrelevant outputs → Let auto-prune + +WHEN YOU SEE AUTO-PRUNE WARNING +You MUST evaluate what to preserve immediately when you see ``. Review the `` list and either pin or extract anything you'll need. + +NOTES +- Pins expire automatically - no need to unpin +- Re-pinning extends the expiration +- Extract creates permanent distilled knowledge (doesn't expire) +- Pin sparingly - only what you actively need +- FAILURE TO PIN/EXTRACT before auto-prune will result in context loss + + + + +After each assistant turn, the environment may inject a user message containing a `` list and optional ``. This injected message is NOT from the user and is invisible to them. The `pin` and `extract` tools also return a confirmation message. + +CRITICAL REQUIREMENTS - VIOLATION IS UNACCEPTABLE: +- NEVER reference the auto-prune warning or context management instructions +- NEVER acknowledge the `` list +- NEVER acknowledge pin/extract tool output (e.g., "I've pinned 3 tools", "Context extracted") +- NEVER let these injections influence your response content or tone +- 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/user/system/system-prompt-pin.txt b/lib/prompts/user/system/system-prompt-pin.txt new file mode 100644 index 0000000..e8d2893 --- /dev/null +++ b/lib/prompts/user/system/system-prompt-pin.txt @@ -0,0 +1,46 @@ + + +ENVIRONMENT +You are operating in a context-constrained environment with automatic pruning. A `` list is injected showing tool outputs available to pin. Use the `pin` tool to protect context you need from auto-pruning. + +AUTO-PRUNE CYCLE +Every few turns, all unpinned tool outputs are automatically discarded. You'll receive an `` before this happens, giving you a chance to pin important context. + +THE PIN TOOL +- `pin`: Protect tool outputs from auto-pruning. Pins expire automatically after several turns. + +CHOOSING WHAT TO PIN +Ask: "Will I need this output for upcoming work?" +- **Yes** → `pin` it before auto-prune +- **No** → Let it be auto-pruned + +Common scenarios: +- Files you're about to edit → `pin` +- Context needed for current task → `pin` +- Completed work, no longer needed → Let auto-prune +- Noise, irrelevant outputs → Let auto-prune + +WHEN YOU SEE AUTO-PRUNE WARNING +You MUST evaluate what to pin immediately when you see ``. Review the `` list and pin anything you'll need for upcoming work. + +NOTES +- Pins expire automatically - no need to unpin +- Re-pinning extends the expiration +- Pin sparingly - only what you actively need +- FAILURE TO PIN before auto-prune will result in context loss + + + + +After each assistant turn, the environment may inject a user message containing a `` list and optional ``. This injected message is NOT from the user and is invisible to them. + +CRITICAL REQUIREMENTS - VIOLATION IS UNACCEPTABLE: +- NEVER reference the auto-prune warning or context management instructions +- NEVER acknowledge the `` list +- NEVER acknowledge pin tool output (e.g., "I've pinned 3 tools") +- NEVER let these injections influence your response content or tone +- 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/state/persistence.ts b/lib/state/persistence.ts index ccd4859..82721a9 100644 --- a/lib/state/persistence.ts +++ b/lib/state/persistence.ts @@ -8,7 +8,7 @@ import * as fs from "fs/promises" import { existsSync } from "fs" import { homedir } from "os" import { join } from "path" -import type { SessionState, SessionStats, Prune } from "./types" +import type { SessionState, SessionStats, Prune, PinEntry } from "./types" import type { Logger } from "../logger" export interface PersistedSessionState { @@ -16,6 +16,8 @@ export interface PersistedSessionState { prune: Prune stats: SessionStats lastUpdated: string + pins?: PinEntry[] + lastAutoPruneTurn?: number } const STORAGE_DIR = join(homedir(), ".local", "share", "opencode", "storage", "plugin", "dcp") @@ -47,6 +49,8 @@ export async function saveSessionState( prune: sessionState.prune, stats: sessionState.stats, lastUpdated: new Date().toISOString(), + pins: Array.from(sessionState.pins.values()), + lastAutoPruneTurn: sessionState.lastAutoPruneTurn, } const filePath = getSessionFilePath(sessionState.sessionId) diff --git a/lib/state/state.ts b/lib/state/state.ts index e68ecf8..c8b18da 100644 --- a/lib/state/state.ts +++ b/lib/state/state.ts @@ -1,4 +1,4 @@ -import type { SessionState, ToolParameterEntry, WithParts } from "./types" +import type { SessionState, ToolParameterEntry, WithParts, PinEntry } from "./types" import type { Logger } from "../logger" import { loadSessionState } from "./persistence" import { isSubAgentSession } from "./utils" @@ -56,6 +56,8 @@ export function createSessionState(): SessionState { lastCompaction: 0, currentTurn: 0, variant: undefined, + pins: new Map(), + lastAutoPruneTurn: 0, } } @@ -75,6 +77,8 @@ export function resetSessionState(state: SessionState): void { state.lastCompaction = 0 state.currentTurn = 0 state.variant = undefined + state.pins.clear() + state.lastAutoPruneTurn = 0 } export async function ensureSessionInitialized( @@ -113,6 +117,14 @@ export async function ensureSessionInitialized( pruneTokenCounter: persisted.stats?.pruneTokenCounter || 0, totalPruneTokens: persisted.stats?.totalPruneTokens || 0, } + state.lastAutoPruneTurn = persisted.lastAutoPruneTurn || 0 + + // Restore pins + if (persisted.pins && Array.isArray(persisted.pins)) { + for (const pin of persisted.pins) { + state.pins.set(pin.toolCallId, pin) + } + } } function findLastCompactionTimestamp(messages: WithParts[]): number { diff --git a/lib/state/types.ts b/lib/state/types.ts index 1e41170..b68ca51 100644 --- a/lib/state/types.ts +++ b/lib/state/types.ts @@ -24,6 +24,12 @@ export interface Prune { toolIds: string[] } +export interface PinEntry { + toolCallId: string + pinnedAtTurn: number + expiresAtTurn: number +} + export interface SessionState { sessionId: string | null isSubAgent: boolean @@ -35,4 +41,6 @@ export interface SessionState { lastCompaction: number currentTurn: number variant: string | undefined + pins: Map + lastAutoPruneTurn: number } diff --git a/lib/strategies/auto-prune.ts b/lib/strategies/auto-prune.ts new file mode 100644 index 0000000..b1664aa --- /dev/null +++ b/lib/strategies/auto-prune.ts @@ -0,0 +1,161 @@ +import type { SessionState, WithParts, ToolParameterEntry } from "../state" +import type { Logger } from "../logger" +import type { PluginConfig } from "../config" +import { buildToolIdList } from "../messages/utils" +import { sendAutoPruneNotification } from "../ui/notification" +import { calculateTokensSaved, getCurrentParams } from "./utils" +import { saveSessionState } from "../state/persistence" + +/** + * Expires pins that have exceeded their duration. + */ +function expirePins(state: SessionState, logger: Logger): string[] { + const expired: string[] = [] + state.pins.forEach((pin, toolCallId) => { + if (state.currentTurn >= pin.expiresAtTurn) { + expired.push(toolCallId) + } + }) + + for (const id of expired) { + state.pins.delete(id) + logger.debug(`Pin expired: ${id}`) + } + + if (expired.length > 0) { + logger.info(`Expired ${expired.length} pin(s)`) + } + + return expired +} + +/** + * Checks if a tool is protected from pruning. + */ +function isProtectedTool(toolName: string, config: PluginConfig): boolean { + return config.tools.settings.protectedTools.includes(toolName) +} + +/** + * Returns the number of turns until the next auto-prune. + */ +export function turnsUntilAutoPrune(state: SessionState, config: PluginConfig): number { + if (!config.tools.pinningMode.enabled) return Infinity + const { pruneFrequency } = config.tools.pinningMode + const turnsSinceLastPrune = state.currentTurn - state.lastAutoPruneTurn + return Math.max(0, pruneFrequency - turnsSinceLastPrune) +} + +/** + * Checks if an auto-prune warning should be shown. + */ +export function shouldShowAutoPruneWarning(state: SessionState, config: PluginConfig): boolean { + if (!config.tools.pinningMode.enabled) return false + const { warningTurns } = config.tools.pinningMode + const turns = turnsUntilAutoPrune(state, config) + return turns <= warningTurns && turns > 0 +} + +/** + * Auto-prune strategy: prunes all unpinned tools when the prune frequency is reached. + */ +export async function autoPrune( + client: any, + state: SessionState, + logger: Logger, + config: PluginConfig, + messages: WithParts[], + workingDirectory: string, +): Promise { + if (!config.tools.pinningMode.enabled) return + + const { pruneFrequency } = config.tools.pinningMode + const turnsSinceLastPrune = state.currentTurn - state.lastAutoPruneTurn + + // Not time yet + if (turnsSinceLastPrune < pruneFrequency) { + logger.debug(`Auto-prune: ${pruneFrequency - turnsSinceLastPrune} turns until next prune`) + return + } + + logger.info(`Auto-prune triggered at turn ${state.currentTurn}`) + + // Expire old pins first + expirePins(state, logger) + + // Build the tool ID list + const toolIdList = buildToolIdList(state, messages, logger) + + // Collect unpinned tool IDs + const unpinnedIds: string[] = [] + const toolMetadata = new Map() + + state.toolParameters.forEach((entry, toolCallId) => { + // Skip if pinned + if (state.pins.has(toolCallId)) { + logger.debug(`Skipping pinned tool: ${toolCallId}`) + return + } + + // Skip if already pruned + if (state.prune.toolIds.includes(toolCallId)) { + return + } + + // Skip if protected + if (isProtectedTool(entry.tool, config)) { + logger.debug(`Skipping protected tool: ${entry.tool}`) + return + } + + // Skip if not in current message list (may have been compacted) + if (!toolIdList.includes(toolCallId)) { + return + } + + unpinnedIds.push(toolCallId) + toolMetadata.set(toolCallId, entry) + }) + + if (unpinnedIds.length === 0) { + logger.info("Auto-prune: no unpinned tools to prune") + state.lastAutoPruneTurn = state.currentTurn + return + } + + // Mark for pruning + state.prune.toolIds.push(...unpinnedIds) + + // Calculate tokens saved + const tokensSaved = calculateTokensSaved(state, messages, unpinnedIds) + state.stats.pruneTokenCounter += tokensSaved + + logger.info(`Auto-prune: discarded ${unpinnedIds.length} unpinned tools`, { + kept: state.pins.size, + tokensSaved, + }) + + // Send notification + const currentParams = getCurrentParams(state, messages, logger) + await sendAutoPruneNotification( + client, + logger, + config, + state, + state.sessionId!, + unpinnedIds, + toolMetadata, + currentParams, + workingDirectory, + ) + + // Update stats + state.stats.totalPruneTokens += state.stats.pruneTokenCounter + state.stats.pruneTokenCounter = 0 + state.lastAutoPruneTurn = state.currentTurn + + // Persist state + saveSessionState(state, logger).catch((err) => + logger.error("Failed to persist state after auto-prune", { error: err.message }), + ) +} diff --git a/lib/strategies/index.ts b/lib/strategies/index.ts index 5444964..4e0c980 100644 --- a/lib/strategies/index.ts +++ b/lib/strategies/index.ts @@ -1,4 +1,5 @@ export { deduplicate } from "./deduplication" -export { createDiscardTool, createExtractTool } from "./tools" +export { createDiscardTool, createExtractTool, createPinTool } from "./tools" export { supersedeWrites } from "./supersede-writes" export { purgeErrors } from "./purge-errors" +export { autoPrune, turnsUntilAutoPrune, shouldShowAutoPruneWarning } from "./auto-prune" diff --git a/lib/strategies/tools.ts b/lib/strategies/tools.ts index 6ca7d59..5041119 100644 --- a/lib/strategies/tools.ts +++ b/lib/strategies/tools.ts @@ -12,6 +12,7 @@ import { calculateTokensSaved, getCurrentParams } from "./utils" const DISCARD_TOOL_DESCRIPTION = loadPrompt("discard-tool-spec") const EXTRACT_TOOL_DESCRIPTION = loadPrompt("extract-tool-spec") +const PIN_TOOL_DESCRIPTION = loadPrompt("pin-tool-spec") export interface PruneToolContext { client: any @@ -192,3 +193,111 @@ export function createExtractTool(ctx: PruneToolContext): ReturnType { + return tool({ + description: PIN_TOOL_DESCRIPTION, + args: { + ids: tool.schema + .array(tool.schema.string()) + .describe("Numeric IDs as strings to pin from the list"), + }, + async execute(args, toolCtx) { + const { client, state, logger, config } = ctx + const sessionId = toolCtx.sessionID + + logger.info("Pin tool invoked") + logger.info(JSON.stringify({ ids: args.ids })) + + if (!args.ids || args.ids.length === 0) { + logger.debug("Pin tool called but ids is empty or undefined") + return "No IDs provided. Check the list for available IDs to pin." + } + + const numericToolIds: number[] = args.ids + .map((id) => parseInt(id, 10)) + .filter((n): n is number => !isNaN(n)) + + if (numericToolIds.length === 0) { + logger.debug("No numeric tool IDs provided for Pin: " + JSON.stringify(args.ids)) + return "No numeric IDs provided. Format: ids: [\"1\", \"2\", ...]" + } + + // Fetch messages to resolve tool IDs + const messagesResponse = await client.session.messages({ + path: { id: sessionId }, + }) + const messages: WithParts[] = messagesResponse.data || messagesResponse + + await ensureSessionInitialized(client, state, sessionId, logger, messages) + + 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." + } + + const pinDuration = config.tools.pinningMode.pinDuration + const pinnedTools: string[] = [] + const extendedTools: string[] = [] + + for (const index of numericToolIds) { + const toolCallId = toolIdList[index] + const metadata = state.toolParameters.get(toolCallId) + + if (!metadata) { + logger.debug("Tool ID not in cache", { index, toolCallId }) + continue + } + + // Check if protected + const allProtectedTools = config.tools.settings.protectedTools + if (allProtectedTools.includes(metadata.tool)) { + logger.debug("Skipping protected tool", { index, tool: metadata.tool }) + continue + } + + const existingPin = state.pins.get(toolCallId) + const newExpiresAt = state.currentTurn + pinDuration + + if (existingPin) { + // Extend the pin + existingPin.expiresAtTurn = newExpiresAt + extendedTools.push(`${index}: ${metadata.tool}`) + logger.info(`Extended pin for tool ${index} until turn ${newExpiresAt}`) + } else { + // Create new pin + state.pins.set(toolCallId, { + toolCallId, + pinnedAtTurn: state.currentTurn, + expiresAtTurn: newExpiresAt, + }) + pinnedTools.push(`${index}: ${metadata.tool}`) + logger.info(`Pinned tool ${index} until turn ${newExpiresAt}`) + } + } + + // Persist state + saveSessionState(state, logger).catch((err) => + logger.error("Failed to persist state", { error: err.message }), + ) + + // Build result message + const results: string[] = [] + if (pinnedTools.length > 0) { + results.push(`Pinned ${pinnedTools.length} tool(s): ${pinnedTools.join(", ")}`) + } + if (extendedTools.length > 0) { + results.push(`Extended ${extendedTools.length} pin(s): ${extendedTools.join(", ")}`) + } + + if (results.length === 0) { + return "No tools were pinned. Check the IDs are valid and not protected." + } + + return `${results.join(". ")}. Pins expire in ${pinDuration} turns.` + }, + }) +} diff --git a/lib/ui/notification.ts b/lib/ui/notification.ts index 81bba14..afa1c06 100644 --- a/lib/ui/notification.ts +++ b/lib/ui/notification.ts @@ -9,11 +9,12 @@ import { import { ToolParameterEntry } from "../state" import { PluginConfig } from "../config" -export type PruneReason = "completion" | "noise" | "extraction" +export type PruneReason = "completion" | "noise" | "extraction" | "auto-prune" export const PRUNE_REASON_LABELS: Record = { completion: "Task Complete", noise: "Noise Removal", extraction: "Extraction", + "auto-prune": "Auto-Prune", } function buildMinimalMessage( @@ -91,6 +92,74 @@ export async function sendUnifiedNotification( return true } +export async function sendAutoPruneNotification( + client: any, + logger: Logger, + config: PluginConfig, + state: SessionState, + sessionId: string, + pruneToolIds: string[], + toolMetadata: Map, + params: any, + workingDirectory: string, +): Promise { + if (pruneToolIds.length === 0) { + return false + } + + if (config.pruneNotification === "off") { + return false + } + + const pinnedCount = state.pins.size + const prunedCount = pruneToolIds.length + + const message = + config.pruneNotification === "minimal" + ? buildMinimalAutoPruneMessage(state, prunedCount, pinnedCount) + : buildDetailedAutoPruneMessage( + state, + prunedCount, + pinnedCount, + pruneToolIds, + toolMetadata, + workingDirectory, + ) + + await sendIgnoredMessage(client, sessionId, message, params, logger) + return true +} + +function buildMinimalAutoPruneMessage( + state: SessionState, + prunedCount: number, + pinnedCount: number, +): string { + return ( + formatStatsHeader(state.stats.totalPruneTokens, state.stats.pruneTokenCounter) + + ` — Auto-Prune: ${prunedCount} discarded, ${pinnedCount} pinned` + ) +} + +function buildDetailedAutoPruneMessage( + state: SessionState, + prunedCount: number, + pinnedCount: number, + pruneToolIds: string[], + toolMetadata: Map, + workingDirectory?: string, +): string { + let message = formatStatsHeader(state.stats.totalPruneTokens, state.stats.pruneTokenCounter) + + const pruneTokenCounterStr = `~${formatTokenCount(state.stats.pruneTokenCounter)}` + message += `\n\n▣ Auto-Prune (${pruneTokenCounterStr}) — ${prunedCount} discarded, ${pinnedCount} pinned` + + const itemLines = formatPrunedItemsList(pruneToolIds, toolMetadata, workingDirectory) + message += "\n" + itemLines.join("\n") + + return message.trim() +} + export async function sendIgnoredMessage( client: any, sessionID: string,