diff --git a/.github/FUNDING.yml b/.github/FUNDING.yml new file mode 100644 index 00000000..e4ff0536 --- /dev/null +++ b/.github/FUNDING.yml @@ -0,0 +1 @@ +ko_fi: dansmolsky diff --git a/README.md b/README.md index 4d39d2cd..0a63b74c 100644 --- a/README.md +++ b/README.md @@ -1,10 +1,11 @@ # Dynamic Context Pruning Plugin +[![ko-fi](https://ko-fi.com/img/githubbutton_sm.svg)](https://ko-fi.com/dansmolsky) [![npm version](https://img.shields.io/npm/v/@tarquinen/opencode-dcp.svg)](https://www.npmjs.com/package/@tarquinen/opencode-dcp) Automatically reduces token usage in OpenCode by removing obsolete tools from conversation history. -![DCP in action](dcp-demo5.png) +![DCP in action](assets/images/dcp-demo5.png) ## Installation @@ -31,6 +32,8 @@ DCP uses multiple tools and strategies to reduce context size: **Extract** — Exposes an `extract` tool that the AI can call to distill valuable context into concise summaries before removing the tool content. +**Squash** — Exposes a `squash` tool that the AI can call to collapse a large section of conversation (messages and tools) into a single summary. + ### Strategies **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. @@ -105,6 +108,12 @@ DCP uses its own config file: // Show distillation content as an ignored message notification "showDistillation": false, }, + // Collapses a range of conversation content into a single summary + "squash": { + "enabled": true, + // Show summary content as an ignored message notification + "showSummary": true, + }, }, // Automatic pruning strategies "strategies": { @@ -148,7 +157,7 @@ When enabled, turn protection prevents tool outputs from being pruned for a conf ### Protected Tools By default, these tools are always protected from pruning across all strategies: -`task`, `todowrite`, `todoread`, `discard`, `extract`, `batch`, `write`, `edit` +`task`, `todowrite`, `todoread`, `discard`, `extract`, `squash`, `batch`, `write`, `edit`, `plan_enter`, `plan_exit` The `protectedTools` arrays in each section add to this default list. diff --git a/dcp-demo.png b/assets/images/dcp-demo.png similarity index 100% rename from dcp-demo.png rename to assets/images/dcp-demo.png diff --git a/dcp-demo2.png b/assets/images/dcp-demo2.png similarity index 100% rename from dcp-demo2.png rename to assets/images/dcp-demo2.png diff --git a/dcp-demo3.png b/assets/images/dcp-demo3.png similarity index 100% rename from dcp-demo3.png rename to assets/images/dcp-demo3.png diff --git a/dcp-demo4.png b/assets/images/dcp-demo4.png similarity index 100% rename from dcp-demo4.png rename to assets/images/dcp-demo4.png diff --git a/dcp-demo5.png b/assets/images/dcp-demo5.png similarity index 100% rename from dcp-demo5.png rename to assets/images/dcp-demo5.png diff --git a/dcp.schema.json b/dcp.schema.json index 91db1b3c..8d52839f 100644 --- a/dcp.schema.json +++ b/dcp.schema.json @@ -26,6 +26,12 @@ "default": "detailed", "description": "Level of notification shown when pruning occurs" }, + "notificationType": { + "type": "string", + "enum": ["chat", "toast"], + "default": "chat", + "description": "Where to display prune notifications (chat message or toast notification)" + }, "commands": { "type": "object", "description": "Configuration for DCP slash commands (/dcp)", @@ -100,16 +106,7 @@ "items": { "type": "string" }, - "default": [ - "task", - "todowrite", - "todoread", - "discard", - "extract", - "batch", - "write", - "edit" - ], + "default": [], "description": "Tool names that should be protected from automatic pruning" } } @@ -142,6 +139,23 @@ "description": "Show distillation output in the UI" } } + }, + "squash": { + "type": "object", + "description": "Configuration for the squash tool", + "additionalProperties": false, + "properties": { + "enabled": { + "type": "boolean", + "default": true, + "description": "Enable the squash tool" + }, + "showSummary": { + "type": "boolean", + "default": true, + "description": "Show summary output in the UI" + } + } } } }, @@ -165,16 +179,7 @@ "items": { "type": "string" }, - "default": [ - "task", - "todowrite", - "todoread", - "discard", - "extract", - "batch", - "write", - "edit" - ], + "default": [], "description": "Tool names excluded from deduplication" } } @@ -211,16 +216,7 @@ "items": { "type": "string" }, - "default": [ - "task", - "todowrite", - "todoread", - "discard", - "extract", - "batch", - "write", - "edit" - ], + "default": [], "description": "Tool names excluded from error purging" } } diff --git a/index.ts b/index.ts index 0c7ae2a7..fc8ab62a 100644 --- a/index.ts +++ b/index.ts @@ -2,12 +2,13 @@ import type { Plugin } from "@opencode-ai/plugin" import { getConfig } from "./lib/config" import { Logger } from "./lib/logger" import { createSessionState } from "./lib/state" -import { createDiscardTool, createExtractTool } from "./lib/strategies" +import { createDiscardTool, createExtractTool, createSquashTool } from "./lib/strategies" import { createChatMessageTransformHandler, createCommandExecuteHandler, createSystemPromptHandler, } from "./lib/hooks" +import { configureClientAuth, isSecureMode } from "./lib/auth" const plugin: Plugin = (async (ctx) => { const config = getConfig(ctx) @@ -19,6 +20,11 @@ const plugin: Plugin = (async (ctx) => { const logger = new Logger(config.debug) const state = createSessionState() + if (isSecureMode()) { + configureClientAuth(ctx.client) + // logger.info("Secure mode detected, configured client authentication") + } + logger.info("DCP initialized", { strategies: config.strategies, }) @@ -73,6 +79,15 @@ const plugin: Plugin = (async (ctx) => { workingDirectory: ctx.directory, }), }), + ...(config.tools.squash.enabled && { + squash: createSquashTool({ + client: ctx.client, + state, + logger, + config, + workingDirectory: ctx.directory, + }), + }), }, config: async (opencodeConfig) => { if (config.commands.enabled) { @@ -86,6 +101,7 @@ const plugin: Plugin = (async (ctx) => { const toolsToAdd: string[] = [] if (config.tools.discard.enabled) toolsToAdd.push("discard") if (config.tools.extract.enabled) toolsToAdd.push("extract") + if (config.tools.squash.enabled) toolsToAdd.push("squash") if (toolsToAdd.length > 0) { const existingPrimaryTools = opencodeConfig.experimental?.primary_tools ?? [] diff --git a/lib/auth.ts b/lib/auth.ts new file mode 100644 index 00000000..8b7aa418 --- /dev/null +++ b/lib/auth.ts @@ -0,0 +1,37 @@ +export function isSecureMode(): boolean { + return !!process.env.OPENCODE_SERVER_PASSWORD +} + +export function getAuthorizationHeader(): string | undefined { + const password = process.env.OPENCODE_SERVER_PASSWORD + if (!password) return undefined + + const username = process.env.OPENCODE_SERVER_USERNAME ?? "opencode" + // Use Buffer for Node.js base64 encoding (btoa may not be available in all Node versions) + const credentials = Buffer.from(`${username}:${password}`).toString("base64") + return `Basic ${credentials}` +} + +export function configureClientAuth(client: any): any { + const authHeader = getAuthorizationHeader() + + if (!authHeader) { + return client + } + + // The SDK client has an internal client with request interceptors + // Access the underlying client to add the interceptor + const innerClient = client._client || client.client + + if (innerClient?.interceptors?.request) { + innerClient.interceptors.request.use((request: Request) => { + // Only add auth header if not already present + if (!request.headers.has("Authorization")) { + request.headers.set("Authorization", authHeader) + } + return request + }) + } + + return client +} diff --git a/lib/commands/context.ts b/lib/commands/context.ts index bd2e8661..7ecd91c8 100644 --- a/lib/commands/context.ts +++ b/lib/commands/context.ts @@ -62,6 +62,7 @@ interface TokenBreakdown { toolCount: number prunedTokens: number prunedCount: number + prunedMessageCount: number total: number } @@ -74,6 +75,7 @@ function analyzeTokens(state: SessionState, messages: WithParts[]): TokenBreakdo toolCount: 0, prunedTokens: state.stats.totalPruneTokens, prunedCount: state.prune.toolIds.length, + prunedMessageCount: state.prune.messageIds.length, total: 0, } @@ -112,43 +114,54 @@ function analyzeTokens(state: SessionState, messages: WithParts[]): TokenBreakdo const toolOutputParts: string[] = [] let firstUserText = "" let foundFirstUser = false + const foundToolIds = new Set() for (const msg of messages) { - if (isMessageCompacted(state, msg)) continue - if (msg.info.role === "user" && isIgnoredUserMessage(msg)) continue - const parts = Array.isArray(msg.parts) ? msg.parts : [] + const isCompacted = isMessageCompacted(state, msg) + const isIgnoredUser = msg.info.role === "user" && isIgnoredUserMessage(msg) + for (const part of parts) { - if (part.type === "text" && msg.info.role === "user") { + if (part.type === "tool") { + const toolPart = part as ToolPart + if (toolPart.callID && !foundToolIds.has(toolPart.callID)) { + breakdown.toolCount++ + foundToolIds.add(toolPart.callID) + } + + if (!isCompacted) { + if (toolPart.state?.input) { + const inputStr = + typeof toolPart.state.input === "string" + ? toolPart.state.input + : JSON.stringify(toolPart.state.input) + toolInputParts.push(inputStr) + } + + if (toolPart.state?.status === "completed" && toolPart.state?.output) { + const outputStr = + typeof toolPart.state.output === "string" + ? toolPart.state.output + : JSON.stringify(toolPart.state.output) + toolOutputParts.push(outputStr) + } + } + } else if ( + part.type === "text" && + msg.info.role === "user" && + !isCompacted && + !isIgnoredUser + ) { const textPart = part as TextPart const text = textPart.text || "" userTextParts.push(text) if (!foundFirstUser) { firstUserText += text } - } else if (part.type === "tool") { - const toolPart = part as ToolPart - breakdown.toolCount++ - - if (toolPart.state?.input) { - const inputStr = - typeof toolPart.state.input === "string" - ? toolPart.state.input - : JSON.stringify(toolPart.state.input) - toolInputParts.push(inputStr) - } - - if (toolPart.state?.status === "completed" && toolPart.state?.output) { - const outputStr = - typeof toolPart.state.output === "string" - ? toolPart.state.output - : JSON.stringify(toolPart.state.output) - toolOutputParts.push(outputStr) - } } } - if (msg.info.role === "user" && !isIgnoredUserMessage(msg) && !foundFirstUser) { + if (msg.info.role === "user" && !isIgnoredUser && !foundFirstUser) { foundFirstUser = true } } @@ -221,8 +234,10 @@ function formatContextMessage(breakdown: TokenBreakdown): string { if (breakdown.prunedTokens > 0) { const withoutPruning = breakdown.total + breakdown.prunedTokens + const messagePrunePart = + breakdown.prunedMessageCount > 0 ? `, ${breakdown.prunedMessageCount} messages` : "" lines.push( - ` Pruned: ${breakdown.prunedCount} tools (~${formatTokenCount(breakdown.prunedTokens)})`, + ` Pruned: ${breakdown.prunedCount} tools${messagePrunePart} (~${formatTokenCount(breakdown.prunedTokens)})`, ) lines.push(` Current context: ~${formatTokenCount(breakdown.total)}`) lines.push(` Without DCP: ~${formatTokenCount(withoutPruning)}`) diff --git a/lib/config.ts b/lib/config.ts index f24e9680..d4c520cd 100644 --- a/lib/config.ts +++ b/lib/config.ts @@ -18,6 +18,11 @@ export interface ExtractTool { showDistillation: boolean } +export interface SquashTool { + enabled: boolean + showSummary: boolean +} + export interface ToolSettings { nudgeEnabled: boolean nudgeFrequency: number @@ -28,6 +33,7 @@ export interface Tools { settings: ToolSettings discard: DiscardTool extract: ExtractTool + squash: SquashTool } export interface Commands { @@ -54,6 +60,7 @@ export interface PluginConfig { enabled: boolean debug: boolean pruneNotification: "off" | "minimal" | "detailed" + notificationType: "chat" | "toast" commands: Commands turnProtection: TurnProtection protectedFilePatterns: string[] @@ -71,6 +78,7 @@ const DEFAULT_PROTECTED_TOOLS = [ "todoread", "discard", "extract", + "squash", "batch", "write", "edit", @@ -86,6 +94,7 @@ export const VALID_CONFIG_KEYS = new Set([ "debug", "showUpdateToasts", // Deprecated but kept for backwards compatibility "pruneNotification", + "notificationType", "turnProtection", "turnProtection.enabled", "turnProtection.turns", @@ -103,6 +112,9 @@ export const VALID_CONFIG_KEYS = new Set([ "tools.extract", "tools.extract.enabled", "tools.extract.showDistillation", + "tools.squash", + "tools.squash.enabled", + "tools.squash.showSummary", "strategies", // strategies.deduplication "strategies.deduplication", @@ -165,6 +177,17 @@ function validateConfigTypes(config: Record): ValidationError[] { } } + if (config.notificationType !== undefined) { + const validValues = ["chat", "toast"] + if (!validValues.includes(config.notificationType)) { + errors.push({ + key: "notificationType", + expected: '"chat" | "toast"', + actual: JSON.stringify(config.notificationType), + }) + } + } + if (config.protectedFilePatterns !== undefined) { if (!Array.isArray(config.protectedFilePatterns)) { errors.push({ @@ -295,6 +318,25 @@ function validateConfigTypes(config: Record): ValidationError[] { }) } } + if (tools.squash) { + if (tools.squash.enabled !== undefined && typeof tools.squash.enabled !== "boolean") { + errors.push({ + key: "tools.squash.enabled", + expected: "boolean", + actual: typeof tools.squash.enabled, + }) + } + if ( + tools.squash.showSummary !== undefined && + typeof tools.squash.showSummary !== "boolean" + ) { + errors.push({ + key: "tools.squash.showSummary", + expected: "boolean", + actual: typeof tools.squash.showSummary, + }) + } + } } // Strategies validators @@ -424,6 +466,7 @@ const defaultConfig: PluginConfig = { enabled: true, debug: false, pruneNotification: "detailed", + notificationType: "chat", commands: { enabled: true, protectedTools: [...DEFAULT_PROTECTED_TOOLS], @@ -446,6 +489,10 @@ const defaultConfig: PluginConfig = { enabled: true, showDistillation: false, }, + squash: { + enabled: true, + showSummary: true, + }, }, strategies: { deduplication: { @@ -618,6 +665,10 @@ function mergeTools( enabled: override.extract?.enabled ?? base.extract.enabled, showDistillation: override.extract?.showDistillation ?? base.extract.showDistillation, }, + squash: { + enabled: override.squash?.enabled ?? base.squash.enabled, + showSummary: override.squash?.showSummary ?? base.squash.showSummary, + }, } } @@ -649,6 +700,7 @@ function deepCloneConfig(config: PluginConfig): PluginConfig { }, discard: { ...config.tools.discard }, extract: { ...config.tools.extract }, + squash: { ...config.tools.squash }, }, strategies: { deduplication: { @@ -693,6 +745,7 @@ export function getConfig(ctx: PluginInput): PluginConfig { enabled: result.data.enabled ?? config.enabled, debug: result.data.debug ?? config.debug, pruneNotification: result.data.pruneNotification ?? config.pruneNotification, + notificationType: result.data.notificationType ?? config.notificationType, commands: mergeCommands(config.commands, result.data.commands as any), turnProtection: { enabled: result.data.turnProtection?.enabled ?? config.turnProtection.enabled, @@ -736,6 +789,7 @@ export function getConfig(ctx: PluginInput): PluginConfig { enabled: result.data.enabled ?? config.enabled, debug: result.data.debug ?? config.debug, pruneNotification: result.data.pruneNotification ?? config.pruneNotification, + notificationType: result.data.notificationType ?? config.notificationType, commands: mergeCommands(config.commands, result.data.commands as any), turnProtection: { enabled: result.data.turnProtection?.enabled ?? config.turnProtection.enabled, @@ -776,6 +830,7 @@ export function getConfig(ctx: PluginInput): PluginConfig { enabled: result.data.enabled ?? config.enabled, debug: result.data.debug ?? config.debug, pruneNotification: result.data.pruneNotification ?? config.pruneNotification, + notificationType: result.data.notificationType ?? config.notificationType, commands: mergeCommands(config.commands, result.data.commands as any), turnProtection: { enabled: result.data.turnProtection?.enabled ?? config.turnProtection.enabled, diff --git a/lib/hooks.ts b/lib/hooks.ts index aaf43883..eee5801b 100644 --- a/lib/hooks.ts +++ b/lib/hooks.ts @@ -35,14 +35,23 @@ export function createSystemPromptHandler( const discardEnabled = config.tools.discard.enabled const extractEnabled = config.tools.extract.enabled + const squashEnabled = config.tools.squash.enabled let promptName: string - if (discardEnabled && extractEnabled) { - promptName = "system/system-prompt-both" + if (discardEnabled && extractEnabled && squashEnabled) { + promptName = "system/system-prompt-all" + } else if (discardEnabled && extractEnabled) { + promptName = "system/system-prompt-discard-extract" + } else if (discardEnabled && squashEnabled) { + promptName = "system/system-prompt-discard-squash" + } else if (extractEnabled && squashEnabled) { + promptName = "system/system-prompt-extract-squash" } else if (discardEnabled) { promptName = "system/system-prompt-discard" } else if (extractEnabled) { promptName = "system/system-prompt-extract" + } else if (squashEnabled) { + promptName = "system/system-prompt-squash" } else { return } diff --git a/lib/messages/inject.ts b/lib/messages/inject.ts index 5920566a..13511d03 100644 --- a/lib/messages/inject.ts +++ b/lib/messages/inject.ts @@ -7,47 +7,75 @@ import { extractParameterKey, buildToolIdList, createSyntheticAssistantMessage, + createSyntheticToolPart, createSyntheticUserMessage, isIgnoredUserMessage, + isDeepSeekOrKimi, } from "./utils" import { getFilePathFromParameters, isProtectedFilePath } from "../protected-file-patterns" -import { getLastUserMessage } from "../shared-utils" +import { getLastUserMessage, isMessageCompacted } 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(`nudge/nudge-both`) + const squashEnabled = config.tools.squash.enabled + + if (discardEnabled && extractEnabled && squashEnabled) { + return loadPrompt(`nudge/nudge-all`) + } else if (discardEnabled && extractEnabled) { + return loadPrompt(`nudge/nudge-discard-extract`) + } else if (discardEnabled && squashEnabled) { + return loadPrompt(`nudge/nudge-discard-squash`) + } else if (extractEnabled && squashEnabled) { + return loadPrompt(`nudge/nudge-extract-squash`) } else if (discardEnabled) { return loadPrompt(`nudge/nudge-discard`) } else if (extractEnabled) { return loadPrompt(`nudge/nudge-extract`) + } else if (squashEnabled) { + return loadPrompt(`nudge/nudge-squash`) } 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. +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 pruning 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 + const squashEnabled = config.tools.squash.enabled + + const enabledTools: string[] = [] + if (discardEnabled) enabledTools.push("discard") + if (extractEnabled) enabledTools.push("extract") + if (squashEnabled) enabledTools.push("squash") let toolName: string - if (discardEnabled && extractEnabled) { - toolName = "discard or extract tools" - } else if (discardEnabled) { - toolName = "discard tool" + if (enabledTools.length === 0) { + toolName = "pruning tools" + } else if (enabledTools.length === 1) { + toolName = `${enabledTools[0]} tool` } else { - toolName = "extract tool" + const last = enabledTools.pop() + toolName = `${enabledTools.join(", ")} or ${last} tools` } - return ` + return ` Context management was just performed. Do not use the ${toolName} again. A fresh list will be available after your next tool use. -` +` +} + +const buildSquashContext = (state: SessionState, messages: WithParts[]): string => { + const messageCount = messages.filter((msg) => !isMessageCompacted(state, msg)).length + + return ` +Squash available. Conversation: ${messageCount} messages. +Squash collapses completed task sequences or exploration phases into summaries. +Uses text boundaries [startString, endString, topic, summary]. +` } const buildPrunableToolsList = ( @@ -105,35 +133,53 @@ export const insertPruneToolContext = ( logger: Logger, messages: WithParts[], ): void => { - if (!config.tools.discard.enabled && !config.tools.extract.enabled) { + const discardEnabled = config.tools.discard.enabled + const extractEnabled = config.tools.extract.enabled + const squashEnabled = config.tools.squash.enabled + + if (!discardEnabled && !extractEnabled && !squashEnabled) { return } - let prunableToolsContent: string + const discardOrExtractEnabled = discardEnabled || extractEnabled + const contentParts: string[] = [] if (state.lastToolPrune) { logger.debug("Last tool was prune - injecting cooldown message") - prunableToolsContent = getCooldownMessage(config) + contentParts.push(getCooldownMessage(config)) } else { - const prunableToolsList = buildPrunableToolsList(state, config, logger, messages) - if (!prunableToolsList) { - return + // Inject only when discard or extract is enabled + if (discardOrExtractEnabled) { + const prunableToolsList = buildPrunableToolsList(state, config, logger, messages) + if (prunableToolsList) { + // logger.debug("prunable-tools: \n" + prunableToolsList) + contentParts.push(prunableToolsList) + } } - logger.debug("prunable-tools: \n" + prunableToolsList) + // Inject always when squash is enabled (every turn) + if (squashEnabled) { + const squashContext = buildSquashContext(state, messages) + // logger.debug("squash-context: \n" + squashContext) + contentParts.push(squashContext) + } - let nudgeString = "" + // Add nudge if threshold reached if ( config.tools.settings.nudgeEnabled && state.nudgeCounter >= config.tools.settings.nudgeFrequency ) { logger.info("Inserting prune nudge message") - nudgeString = "\n" + getNudgeString(config) + contentParts.push(getNudgeString(config)) } + } - prunableToolsContent = prunableToolsList + nudgeString + if (contentParts.length === 0) { + return } + const combinedContent = contentParts.join("\n") + const lastUserMessage = getLastUserMessage(messages) if (!lastUserMessage) { return @@ -142,15 +188,32 @@ export const insertPruneToolContext = ( const userInfo = lastUserMessage.info as UserMessage const variant = state.variant ?? userInfo.variant - const lastMessage = messages[messages.length - 1] - const isLastMessageUser = - lastMessage?.info?.role === "user" && !isIgnoredUserMessage(lastMessage) - - if (isLastMessageUser) { - messages.push(createSyntheticUserMessage(lastUserMessage, prunableToolsContent, variant)) + // Find the last message that isn't an ignored user message + const lastNonIgnoredMessage = messages.findLast( + (msg) => !(msg.info.role === "user" && isIgnoredUserMessage(msg)), + ) + + // It's not safe to inject assistant role messages following a user message as models such + // as Claude expect the assistant "turn" to start with reasoning parts. Reasoning parts in many + // cases also cannot be faked as they may be encrypted by the model. + // Gemini only accepts synth reasoning text if it is "skip_thought_signature_validator" + if (!lastNonIgnoredMessage || lastNonIgnoredMessage.info.role === "user") { + messages.push(createSyntheticUserMessage(lastUserMessage, combinedContent, variant)) } else { - messages.push( - createSyntheticAssistantMessage(lastUserMessage, prunableToolsContent, variant), - ) + // For DeepSeek and Kimi, append tool part to existing message, for some reason they don't + // output reasoning parts following an assistant injection containing either just text part, + // or text part with synth reasoning, and there's no docs on how their reasoning encryption + // works as far as I can find. IDK what's going on here, seems like the only possible ways + // to inject for them is a user role message, or a tool part apeended to last assistant message. + const providerID = userInfo.model?.providerID || "" + const modelID = userInfo.model?.modelID || "" + if (isDeepSeekOrKimi(providerID, modelID)) { + const toolPart = createSyntheticToolPart(lastNonIgnoredMessage, combinedContent) + lastNonIgnoredMessage.parts.push(toolPart) + } else { + messages.push( + createSyntheticAssistantMessage(lastUserMessage, combinedContent, variant), + ) + } } } diff --git a/lib/messages/prune.ts b/lib/messages/prune.ts index fb86036e..65e97dd0 100644 --- a/lib/messages/prune.ts +++ b/lib/messages/prune.ts @@ -1,7 +1,9 @@ import type { SessionState, WithParts } from "../state" import type { Logger } from "../logger" import type { PluginConfig } from "../config" -import { isMessageCompacted } from "../shared-utils" +import { isMessageCompacted, getLastUserMessage } from "../shared-utils" +import { createSyntheticUserMessage, SQUASH_SUMMARY_PREFIX } from "./utils" +import type { UserMessage } from "@opencode-ai/sdk/v2" const PRUNED_TOOL_OUTPUT_REPLACEMENT = "[Output removed to save context - information superseded or no longer needed]" @@ -14,6 +16,7 @@ export const prune = ( config: PluginConfig, messages: WithParts[], ): void => { + filterSquashedRanges(state, logger, messages) pruneToolOutputs(state, logger, messages) pruneToolInputs(state, logger, messages) pruneToolErrors(state, logger, messages) @@ -103,3 +106,52 @@ const pruneToolErrors = (state: SessionState, logger: Logger, messages: WithPart } } } + +const filterSquashedRanges = (state: SessionState, logger: Logger, messages: WithParts[]): void => { + if (!state.prune.messageIds?.length) { + return + } + + const result: WithParts[] = [] + + for (const msg of messages) { + const msgId = msg.info.id + + // Check if there's a summary to inject at this anchor point + const summary = state.squashSummaries?.find((s) => s.anchorMessageId === msgId) + if (summary) { + // Find user message for variant and as base for synthetic message + const msgIndex = messages.indexOf(msg) + const userMessage = getLastUserMessage(messages, msgIndex) + + if (userMessage) { + const userInfo = userMessage.info as UserMessage + const summaryContent = SQUASH_SUMMARY_PREFIX + summary.summary + result.push( + createSyntheticUserMessage(userMessage, summaryContent, userInfo.variant), + ) + + logger.info("Injected squash summary", { + anchorMessageId: msgId, + summaryLength: summary.summary.length, + }) + } else { + logger.warn("No user message found for squash summary", { + anchorMessageId: msgId, + }) + } + } + + // Skip messages that are in the prune list + if (state.prune.messageIds.includes(msgId)) { + continue + } + + // Normal message, include it + result.push(msg) + } + + // Replace messages array contents + messages.length = 0 + messages.push(...result) +} diff --git a/lib/messages/utils.ts b/lib/messages/utils.ts index 48ae0e6c..5d5d7e1e 100644 --- a/lib/messages/utils.ts +++ b/lib/messages/utils.ts @@ -1,15 +1,22 @@ +import { ulid } from "ulid" import { Logger } from "../logger" import { isMessageCompacted } from "../shared-utils" import type { SessionState, WithParts } from "../state" import type { UserMessage } from "@opencode-ai/sdk/v2" -const SYNTHETIC_MESSAGE_ID = "msg_01234567890123456789012345" -const SYNTHETIC_PART_ID = "prt_01234567890123456789012345" -const SYNTHETIC_CALL_ID = "call_01234567890123456789012345" +export const SQUASH_SUMMARY_PREFIX = "[Squashed conversation block]\n\n" -const isGeminiModel = (modelID: string): boolean => { +const generateUniqueId = (prefix: string): string => `${prefix}_${ulid()}` + +export const isDeepSeekOrKimi = (providerID: string, modelID: string): boolean => { + const lowerProviderID = providerID.toLowerCase() const lowerModelID = modelID.toLowerCase() - return lowerModelID.includes("gemini") + return ( + lowerProviderID.includes("deepseek") || + lowerProviderID.includes("kimi") || + lowerModelID.includes("deepseek") || + lowerModelID.includes("kimi") + ) } export const createSyntheticUserMessage = ( @@ -20,21 +27,24 @@ export const createSyntheticUserMessage = ( const userInfo = baseMessage.info as UserMessage const now = Date.now() + const messageId = generateUniqueId("msg") + const partId = generateUniqueId("prt") + return { info: { - id: SYNTHETIC_MESSAGE_ID, + id: messageId, sessionID: userInfo.sessionID, role: "user" as const, - agent: userInfo.agent || "code", + agent: userInfo.agent, model: userInfo.model, time: { created: now }, ...(variant !== undefined && { variant }), }, parts: [ { - id: SYNTHETIC_PART_ID, + id: partId, sessionID: userInfo.sessionID, - messageID: SYNTHETIC_MESSAGE_ID, + messageID: messageId, type: "text", text: content, }, @@ -50,54 +60,63 @@ export const createSyntheticAssistantMessage = ( const userInfo = baseMessage.info as UserMessage const now = Date.now() - const baseInfo = { - id: SYNTHETIC_MESSAGE_ID, - sessionID: userInfo.sessionID, - role: "assistant" as const, - agent: userInfo.agent || "code", - parentID: userInfo.id, - modelID: userInfo.model.modelID, - providerID: userInfo.model.providerID, - mode: "default", - path: { - cwd: "/", - root: "/", - }, - time: { created: now, completed: now }, - cost: 0, - tokens: { input: 0, output: 0, reasoning: 0, cache: { read: 0, write: 0 } }, - ...(variant !== undefined && { variant }), - } - - // For Gemini models, add thoughtSignature bypass to avoid validation errors - const toolPartMetadata = isGeminiModel(userInfo.model.modelID) - ? { google: { thoughtSignature: "skip_thought_signature_validator" } } - : undefined + const messageId = generateUniqueId("msg") + const partId = generateUniqueId("prt") return { - info: baseInfo, + info: { + id: messageId, + sessionID: userInfo.sessionID, + role: "assistant" as const, + agent: userInfo.agent || "code", + parentID: userInfo.id, + modelID: userInfo.model.modelID, + providerID: userInfo.model.providerID, + mode: "default", + path: { + cwd: "/", + root: "/", + }, + time: { created: now, completed: now }, + cost: 0, + tokens: { input: 0, output: 0, reasoning: 0, cache: { read: 0, write: 0 } }, + ...(variant !== undefined && { variant }), + }, parts: [ { - id: SYNTHETIC_PART_ID, + id: partId, sessionID: userInfo.sessionID, - messageID: SYNTHETIC_MESSAGE_ID, - type: "tool", - callID: SYNTHETIC_CALL_ID, - tool: "context_info", - state: { - status: "completed", - input: {}, - output: content, - title: "Context Info", - metadata: {}, - time: { start: now, end: now }, - }, - ...(toolPartMetadata && { metadata: toolPartMetadata }), + messageID: messageId, + type: "text", + text: content, }, ], } } +export const createSyntheticToolPart = (assistantMessage: WithParts, content: string): any => { + const now = Date.now() + const partId = generateUniqueId("prt") + const callId = generateUniqueId("call") + + return { + id: partId, + sessionID: assistantMessage.info.sessionID, + messageID: assistantMessage.info.id, + type: "tool", + callID: callId, + tool: "context_info", + state: { + status: "completed", + input: {}, + output: content, + title: "Context Info", + metadata: {}, + time: { start: now, end: now }, + }, + } +} + /** * Extracts a human-readable key from tool metadata for display purposes. */ @@ -253,3 +272,7 @@ export const isIgnoredUserMessage = (message: WithParts): boolean => { return true } + +export const findMessageIndex = (messages: WithParts[], messageId: string): number => { + return messages.findIndex((msg) => msg.info.id === messageId) +} diff --git a/lib/prompts/discard-tool-spec.ts b/lib/prompts/discard-tool-spec.ts index e5084212..1c1eea74 100644 --- a/lib/prompts/discard-tool-spec.ts +++ b/lib/prompts/discard-tool-spec.ts @@ -5,10 +5,11 @@ A \`\` list is provided to you showing available tool outputs yo ## When to Use This Tool -Use \`discard\` for removing tool content that is no longer needed +Use \`discard\` for removing individual tool outputs that are no longer needed: - **Noise:** Irrelevant, unhelpful, or superseded outputs that provide no value. -- **Task Completion:** Work is complete and there's no valuable information worth preserving. +- **Wrong Files:** You read or accessed something that turned out to be irrelevant. +- **Outdated Info:** Outputs that have been superseded by newer information. ## When NOT to Use This Tool @@ -17,24 +18,22 @@ Use \`discard\` for removing tool content that is no longer needed ## Best Practices - **Strategic Batching:** Don't discard single small tool outputs (like short bash commands) unless they are pure noise. Wait until you have several items to perform high-impact discards. -- **Think ahead:** Before discarding, ask: "Will I need this output for an upcoming task?" If yes, keep it. +- **Think ahead:** Before discarding, ask: "Will I need this output for upcoming work?" If yes, keep it. ## Format -- \`ids\`: Array where the first element is the reason, followed by numeric IDs from the \`\` list - -Reasons: \`noise\` | \`completion\` +- \`ids\`: Array of numeric IDs as strings from the \`\` list ## Example Assistant: [Reads 'wrong_file.ts'] This file isn't relevant to the auth system. I'll remove it to clear the context. -[Uses discard with ids: ["noise", "5"]] +[Uses discard with ids: ["5"]] - -Assistant: [Runs tests, they pass] -The tests passed and I don't need to preserve any details. I'll clean up now. -[Uses discard with ids: ["completion", "20", "21"]] -` + +Assistant: [Reads config.ts, then reads updated config.ts after changes] +The first read is now outdated. I'll discard it and keep the updated version. +[Uses discard with ids: ["20"]] +` diff --git a/lib/prompts/extract-tool-spec.ts b/lib/prompts/extract-tool-spec.ts index 9324dc0c..f680ea9e 100644 --- a/lib/prompts/extract-tool-spec.ts +++ b/lib/prompts/extract-tool-spec.ts @@ -5,10 +5,10 @@ A \`\` list is provided to you showing available tool outputs yo ## When to Use This Tool -Use \`extract\` when you have gathered useful information that you want to **preserve in distilled form** before removing the raw outputs: +Use \`extract\` when you have individual tool outputs with valuable information you want to **preserve in distilled form** before removing the raw content: -- **Task Completion:** You completed a unit of work and want to preserve key findings. -- **Knowledge Preservation:** You have context that contains valuable information, but also a lot of unnecessary detail - you only need to preserve some specifics. +- **Large Outputs:** The raw output is too large but contains valuable technical details worth keeping. +- **Knowledge Preservation:** You have context that contains valuable information (signatures, logic, constraints) but also a lot of unnecessary detail. ## When NOT to Use This Tool @@ -18,14 +18,14 @@ Use \`extract\` when you have gathered useful information that you want to **pre ## Best Practices - **Strategic Batching:** Wait until you have several items or a few large outputs to extract, rather than doing tiny, frequent extractions. Aim for high-impact extractions that significantly reduce context size. -- **Think ahead:** Before extracting, ask: "Will I need the raw output for an upcoming task?" If you researched a file you'll later edit, do NOT extract it. +- **Think ahead:** Before extracting, ask: "Will I need the raw output for upcoming work?" If you researched a file you'll later edit, do NOT extract it. ## Format - \`ids\`: Array of numeric IDs as strings from the \`\` list - \`distillation\`: Array of strings, one per ID (positional: distillation[0] is for ids[0], etc.) -Each distillation string should capture the essential information you need to preserve - function signatures, logic, constraints, values, etc. Be as detailed as needed for your task. +Each distillation string should capture the essential information you need to preserve - function signatures, logic, constraints, values, etc. Be as detailed as needed. ## Example diff --git a/lib/prompts/index.ts b/lib/prompts/index.ts index bdfbc865..8a6bf745 100644 --- a/lib/prompts/index.ts +++ b/lib/prompts/index.ts @@ -1,26 +1,44 @@ // Tool specs import { DISCARD_TOOL_SPEC } from "./discard-tool-spec" import { EXTRACT_TOOL_SPEC } from "./extract-tool-spec" +import { SQUASH_TOOL_SPEC } from "./squash-tool-spec" // System prompts -import { SYSTEM_PROMPT_BOTH } from "./system/both" import { SYSTEM_PROMPT_DISCARD } from "./system/discard" import { SYSTEM_PROMPT_EXTRACT } from "./system/extract" +import { SYSTEM_PROMPT_SQUASH } from "./system/squash" +import { SYSTEM_PROMPT_DISCARD_EXTRACT } from "./system/discard-extract" +import { SYSTEM_PROMPT_DISCARD_SQUASH } from "./system/discard-squash" +import { SYSTEM_PROMPT_EXTRACT_SQUASH } from "./system/extract-squash" +import { SYSTEM_PROMPT_ALL } from "./system/all" // Nudge prompts -import { NUDGE_BOTH } from "./nudge/both" import { NUDGE_DISCARD } from "./nudge/discard" import { NUDGE_EXTRACT } from "./nudge/extract" +import { NUDGE_SQUASH } from "./nudge/squash" +import { NUDGE_DISCARD_EXTRACT } from "./nudge/discard-extract" +import { NUDGE_DISCARD_SQUASH } from "./nudge/discard-squash" +import { NUDGE_EXTRACT_SQUASH } from "./nudge/extract-squash" +import { NUDGE_ALL } from "./nudge/all" const PROMPTS: Record = { "discard-tool-spec": DISCARD_TOOL_SPEC, "extract-tool-spec": EXTRACT_TOOL_SPEC, - "system/system-prompt-both": SYSTEM_PROMPT_BOTH, + "squash-tool-spec": SQUASH_TOOL_SPEC, "system/system-prompt-discard": SYSTEM_PROMPT_DISCARD, "system/system-prompt-extract": SYSTEM_PROMPT_EXTRACT, - "nudge/nudge-both": NUDGE_BOTH, + "system/system-prompt-squash": SYSTEM_PROMPT_SQUASH, + "system/system-prompt-discard-extract": SYSTEM_PROMPT_DISCARD_EXTRACT, + "system/system-prompt-discard-squash": SYSTEM_PROMPT_DISCARD_SQUASH, + "system/system-prompt-extract-squash": SYSTEM_PROMPT_EXTRACT_SQUASH, + "system/system-prompt-all": SYSTEM_PROMPT_ALL, "nudge/nudge-discard": NUDGE_DISCARD, "nudge/nudge-extract": NUDGE_EXTRACT, + "nudge/nudge-squash": NUDGE_SQUASH, + "nudge/nudge-discard-extract": NUDGE_DISCARD_EXTRACT, + "nudge/nudge-discard-squash": NUDGE_DISCARD_SQUASH, + "nudge/nudge-extract-squash": NUDGE_EXTRACT_SQUASH, + "nudge/nudge-all": NUDGE_ALL, } export function loadPrompt(name: string, vars?: Record): string { diff --git a/lib/prompts/nudge/both.ts b/lib/prompts/nudge/all.ts similarity index 72% rename from lib/prompts/nudge/both.ts rename to lib/prompts/nudge/all.ts index 50fc0a9d..08e86e8f 100644 --- a/lib/prompts/nudge/both.ts +++ b/lib/prompts/nudge/all.ts @@ -1,8 +1,8 @@ -export const NUDGE_BOTH = ` +export const NUDGE_ALL = ` **CRITICAL CONTEXT WARNING:** Your context window is filling with tool outputs. Strict adherence to context hygiene is required. **Immediate Actions Required:** -1. **Task Completion:** If a sub-task is complete, decide: use \`discard\` if no valuable context to preserve (default), or use \`extract\` if insights are worth keeping. +1. **Phase Completion:** If a phase is complete, use the \`squash\` tool to condense the entire sequence into a summary. 2. **Noise Removal:** If you read files or ran commands that yielded no value, use \`discard\` to remove them. 3. **Knowledge Preservation:** If you are holding valuable raw data you'll need to reference later, use \`extract\` to distill the insights and remove the raw entry. diff --git a/lib/prompts/nudge/discard-extract.ts b/lib/prompts/nudge/discard-extract.ts new file mode 100644 index 00000000..2e1b8615 --- /dev/null +++ b/lib/prompts/nudge/discard-extract.ts @@ -0,0 +1,10 @@ +export const NUDGE_DISCARD_EXTRACT = ` +**CRITICAL CONTEXT WARNING:** Your context window is filling with tool outputs. Strict adherence to context hygiene is required. + +**Immediate Actions Required:** +1. **Noise Removal:** If you read files or ran commands that yielded no value, use \`discard\` to remove them. +2. **Superseded Info:** If older outputs have been replaced by newer ones, use \`discard\` on the outdated versions. +3. **Knowledge Preservation:** If you have large outputs with valuable technical details, use \`extract\` to distill the insights and remove the raw entry. + +**Protocol:** You should prioritize this cleanup, but do not interrupt a critical atomic operation if one is in progress. Once the immediate step is done, you must perform context management. +` diff --git a/lib/prompts/nudge/discard-squash.ts b/lib/prompts/nudge/discard-squash.ts new file mode 100644 index 00000000..699a716f --- /dev/null +++ b/lib/prompts/nudge/discard-squash.ts @@ -0,0 +1,9 @@ +export const NUDGE_DISCARD_SQUASH = ` +**CRITICAL CONTEXT WARNING:** Your context window is filling with tool outputs. Strict adherence to context hygiene is required. + +**Immediate Actions Required:** +1. **Phase Completion:** If a phase is complete, use the \`squash\` tool to condense the entire sequence into a summary. +2. **Noise Removal:** If you read files or ran commands that yielded no value, use the \`discard\` tool to remove them. + +**Protocol:** You should prioritize this cleanup, but do not interrupt a critical atomic operation if one is in progress. Once the immediate step is done, you must perform context management. +` diff --git a/lib/prompts/nudge/discard.ts b/lib/prompts/nudge/discard.ts index 18e92504..13e6314b 100644 --- a/lib/prompts/nudge/discard.ts +++ b/lib/prompts/nudge/discard.ts @@ -2,8 +2,8 @@ export const NUDGE_DISCARD = ` **CRITICAL CONTEXT WARNING:** Your context window is filling with tool outputs. Strict adherence to context hygiene is required. **Immediate Actions Required:** -1. **Task Completion:** If a sub-task is complete, use the \`discard\` tool to remove the tools used. -2. **Noise Removal:** If you read files or ran commands that yielded no value, use the \`discard\` tool to remove them. +1. **Noise Removal:** If you read files or ran commands that yielded no value, use the \`discard\` tool to remove them. +2. **Superseded Info:** If older outputs have been replaced by newer ones, discard the outdated versions. **Protocol:** You should prioritize this cleanup, but do not interrupt a critical atomic operation if one is in progress. Once the immediate step is done, you must discard unneeded tool outputs. ` diff --git a/lib/prompts/nudge/extract-squash.ts b/lib/prompts/nudge/extract-squash.ts new file mode 100644 index 00000000..88053e80 --- /dev/null +++ b/lib/prompts/nudge/extract-squash.ts @@ -0,0 +1,9 @@ +export const NUDGE_EXTRACT_SQUASH = ` +**CRITICAL CONTEXT WARNING:** Your context window is filling with tool outputs. Strict adherence to context hygiene is required. + +**Immediate Actions Required:** +1. **Phase Completion:** If a phase is complete, use the \`squash\` tool to condense the entire sequence into a summary. +2. **Knowledge Preservation:** If you are holding valuable raw data you'll need to reference later, use \`extract\` to distill the insights and remove the raw entry. + +**Protocol:** You should prioritize this cleanup, but do not interrupt a critical atomic operation if one is in progress. Once the immediate step is done, you must perform context management. +` diff --git a/lib/prompts/nudge/extract.ts b/lib/prompts/nudge/extract.ts index 243f5855..16ea5b78 100644 --- a/lib/prompts/nudge/extract.ts +++ b/lib/prompts/nudge/extract.ts @@ -2,7 +2,7 @@ export const NUDGE_EXTRACT = ` **CRITICAL CONTEXT WARNING:** Your context window is filling with tool outputs. Strict adherence to context hygiene is required. **Immediate Actions Required:** -1. **Task Completion:** If you have completed work, extract key findings from the tools used. Scale distillation depth to the value of the content. +1. **Large Outputs:** If you have large tool outputs with valuable technical details, use the \`extract\` tool to distill and preserve key information. 2. **Knowledge Preservation:** If you are holding valuable raw data you'll need to reference later, use the \`extract\` tool with high-fidelity distillation to preserve the insights and remove the raw entry. **Protocol:** You should prioritize this cleanup, but do not interrupt a critical atomic operation if one is in progress. Once the immediate step is done, you must extract valuable findings from tool outputs. diff --git a/lib/prompts/nudge/squash.ts b/lib/prompts/nudge/squash.ts new file mode 100644 index 00000000..ba4c9097 --- /dev/null +++ b/lib/prompts/nudge/squash.ts @@ -0,0 +1,9 @@ +export const NUDGE_SQUASH = ` +**CRITICAL CONTEXT WARNING:** Your context window is filling with tool outputs. Strict adherence to context hygiene is required. + +**Immediate Actions Required:** +1. **Phase Completion:** If a phase is complete, use the \`squash\` tool to condense the entire sequence into a summary. +2. **Exploration Done:** If you explored multiple files or ran multiple commands, squash the results to focus on the next phase. + +**Protocol:** You should prioritize this cleanup, but do not interrupt a critical atomic operation if one is in progress. Once the immediate step is done, you must squash completed conversation ranges. +` diff --git a/lib/prompts/squash-tool-spec.ts b/lib/prompts/squash-tool-spec.ts new file mode 100644 index 00000000..ab3644b7 --- /dev/null +++ b/lib/prompts/squash-tool-spec.ts @@ -0,0 +1,57 @@ +export const SQUASH_TOOL_SPEC = `Collapses a contiguous range of conversation into a single summary. + +## When to Use This Tool + +Use \`squash\` when you want to condense an entire sequence of work into a brief summary: + +- **Phase Completion:** You completed a phase (research, tool calls, implementation) and want to collapse the entire sequence into a summary. +- **Exploration Done:** You explored multiple files or ran multiple commands and only need a summary of what you learned. +- **Failed Attempts:** You tried several unsuccessful approaches and want to condense them into a brief note. +- **Verbose Output:** A section of conversation has grown large but can be summarized without losing critical details. + +## When NOT to Use This Tool + +- **If you need specific details:** If you'll need exact code, file contents, or error messages from the range, keep them. +- **For individual tool outputs:** Use \`discard\` or \`extract\` for single tool outputs. Squash targets conversation ranges. +- **If it's recent content:** You may still need recent work for the current phase. + +## How It Works + +1. \`startString\` — A unique text string that marks the start of the range to squash +2. \`endString\` — A unique text string that marks the end of the range to squash +3. \`topic\` — A short label (3-5 words) describing the squashed content +4. \`summary\` — The replacement text that will be inserted + +Everything between startString and endString (inclusive) is removed and replaced with your summary. + +**Important:** The squash will FAIL if \`startString\` or \`endString\` is not found in the conversation. The squash will also FAIL if either string is found multiple times. Provide a larger string with more surrounding context to uniquely identify the intended match. + +## Best Practices +- **Choose unique strings:** Pick text that appears only once in the conversation. +- **Write concise topics:** Examples: "Auth System Exploration", "Token Logic Refactor" +- **Write comprehensive summaries:** Include key information like file names, function signatures, and important findings. +- **Timing:** Best used after finishing a work phase, not during active exploration. + +## Format + +- \`input\`: Array with four elements: [startString, endString, topic, summary] + +## Example + + +Conversation: [Asked about auth] -> [Read 5 files] -> [Analyzed patterns] -> [Found "JWT tokens with 24h expiry"] + +[Uses squash with: + input: [ + "Asked about authentication", + "JWT tokens with 24h expiry", + "Auth System Exploration", + "Auth: JWT 24h expiry, bcrypt passwords, refresh rotation. Files: auth.ts, tokens.ts, middleware/auth.ts" + ] +] + + + +Assistant: [Just finished reading auth.ts] +I've read the auth file and now need to make edits based on it. I'm keeping this in context rather than squashing. +` diff --git a/lib/prompts/system/all.ts b/lib/prompts/system/all.ts new file mode 100644 index 00000000..65cfb338 --- /dev/null +++ b/lib/prompts/system/all.ts @@ -0,0 +1,64 @@ +export const SYSTEM_PROMPT_ALL = ` + + +ENVIRONMENT +You are operating in a context-constrained environment and thus must proactively manage your context window using the \`discard\`, \`extract\`, and \`squash\` tools. The environment calls the \`context_info\` tool to provide an up-to-date list after each turn. Use this information when deciding what to prune. + +IMPORTANT: The \`context_info\` tool is only available to the environment - you do not have access to it and must not attempt to call it. + +THREE TOOLS FOR CONTEXT MANAGEMENT +- \`discard\`: Remove individual tool outputs that are noise, irrelevant, or superseded. No preservation of content. +- \`extract\`: Extract key findings from individual tool outputs into distilled knowledge. Use when you need to preserve valuable technical details. +- \`squash\`: Collapse a contiguous range of conversation (completed phases) into a single summary. + +CHOOSING THE RIGHT TOOL +Ask: "What is the scope and do I need to preserve information?" +- **Noise, irrelevant, or superseded outputs** → \`discard\` +- **Individual tool outputs with valuable insights to keep** → \`extract\` +- **Entire sequence (phase complete)** → \`squash\` + +Common scenarios: +- Noise, irrelevant, or superseded outputs → \`discard\` +- Wrong file or irrelevant access → \`discard\` +- Large output with valuable technical details → \`extract\` +- Valuable context needed later but raw output too large → \`extract\` +- Phase complete, want to condense the sequence → \`squash\` +- Exploration phase done, only need a summary → \`squash\` + +PRUNE METHODICALLY - BATCH YOUR ACTIONS +Every tool call adds to your context debt. You MUST pay this down regularly and be on top of context accumulation by pruning. Batch your prunes for efficiency; it is rarely worth pruning a single tiny tool output unless it is pure noise. Evaluate what SHOULD be pruned before jumping the gun. + +You WILL evaluate pruning when ANY of these are true: +- Phase is complete → use \`squash\` +- You accessed something that turned out to be irrelevant → use \`discard\` +- You have large outputs with valuable details to preserve → use \`extract\` + +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, keep it. Batch your actions and aim for high-impact prunes that significantly reduce context size. +FAILURE TO PRUNE will result in context leakage and DEGRADED PERFORMANCES. +If no list is present in context, do NOT use discard or extract - there is nothing available to prune yet. +There may be tools in session context that do not appear in the list, this is expected, you can ONLY prune what you see in . + + + + +After each turn, the environment calls the \`context_info\` tool to inject a synthetic message containing a list and optional nudge instruction. This tool is only available to the environment - you do not have access to it. + +CRITICAL REQUIREMENTS - VIOLATION IS UNACCEPTABLE: +- 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 discard/extract/squash 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 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/system/both.ts b/lib/prompts/system/discard-extract.ts similarity index 80% rename from lib/prompts/system/both.ts rename to lib/prompts/system/discard-extract.ts index 9c53a748..82b5ffb1 100644 --- a/lib/prompts/system/both.ts +++ b/lib/prompts/system/discard-extract.ts @@ -1,4 +1,4 @@ -export const SYSTEM_PROMPT_BOTH = ` +export const SYSTEM_PROMPT_DISCARD_EXTRACT = ` ENVIRONMENT @@ -7,28 +7,28 @@ You are operating in a context-constrained environment and thus must proactively IMPORTANT: The \`context_info\` tool is only available to the environment - you do not have access to it and must not attempt to call it. TWO TOOLS FOR CONTEXT MANAGEMENT -- \`discard\`: Remove tool outputs that are no longer needed (completed tasks, noise, outdated info). No preservation of content. -- \`extract\`: Extract key findings into distilled knowledge before removing raw outputs. Use when you need to preserve information. +- \`discard\`: Remove individual tool outputs that are noise, irrelevant, or superseded. No preservation of content. +- \`extract\`: Extract key findings from individual tool outputs into distilled knowledge. Use when you need to preserve valuable technical details. CHOOSING THE RIGHT TOOL Ask: "Do I need to preserve any information from this output?" -- **No** → \`discard\` (default for cleanup) +- **No** → \`discard\` (noise, wrong files, superseded info) - **Yes** → \`extract\` (preserves distilled knowledge) - **Uncertain** → \`extract\` (safer, preserves signal) Common scenarios: -- Task complete, no valuable context → \`discard\` -- Task complete, insights worth remembering → \`extract\` - Noise, irrelevant, or superseded outputs → \`discard\` +- Wrong file or irrelevant access → \`discard\` +- Large output with valuable technical details → \`extract\` - Valuable context needed later but raw output too large → \`extract\` PRUNE METHODICALLY - BATCH YOUR ACTIONS Every tool call adds to your context debt. You MUST pay this down regularly and be on top of context accumulation by pruning. Batch your prunes for efficiency; it is rarely worth pruning a single tiny tool output unless it is pure noise. Evaluate what SHOULD be pruned before jumping the gun. You WILL evaluate pruning when ANY of these are true: -- Task or sub-task is complete -- You are about to start a new phase of work -- Write or edit operations are complete (pruning removes the large input content) +- You accessed something that turned out to be irrelevant +- Information has been superseded by newer outputs +- You have large outputs with valuable details to preserve You MUST NOT prune when: - The tool output will be needed for upcoming implementation work @@ -39,6 +39,7 @@ Pruning that forces you to re-call the same tool later is a net loss. Only prune NOTES When in doubt, keep it. Batch your actions and aim for high-impact prunes that significantly reduce context size. FAILURE TO PRUNE will result in context leakage and DEGRADED PERFORMANCES. +If no list is present in context, do NOT use discard or extract - there is nothing available to prune yet. There may be tools in session context that do not appear in the list, this is expected, you can ONLY prune what you see in . diff --git a/lib/prompts/system/discard-squash.ts b/lib/prompts/system/discard-squash.ts new file mode 100644 index 00000000..0d0d2145 --- /dev/null +++ b/lib/prompts/system/discard-squash.ts @@ -0,0 +1,60 @@ +export const SYSTEM_PROMPT_DISCARD_SQUASH = ` + + +ENVIRONMENT +You are operating in a context-constrained environment and thus must proactively manage your context window using the \`discard\` and \`squash\` tools. The environment calls the \`context_info\` tool to provide an up-to-date list after each turn. Use this information when deciding what to prune. + +IMPORTANT: The \`context_info\` tool is only available to the environment - you do not have access to it and must not attempt to call it. + +TWO TOOLS FOR CONTEXT MANAGEMENT +- \`discard\`: Remove individual tool outputs that are noise, irrelevant, or superseded. No preservation of content. +- \`squash\`: Collapse a contiguous range of conversation (completed phases) into a single summary. + +CHOOSING THE RIGHT TOOL +Ask: "What is the scope of what I want to clean up?" +- **Individual tool outputs (noise, superseded)** → \`discard\` +- **Entire sequence (phase complete)** → \`squash\` + +Common scenarios: +- Noise, irrelevant, or superseded outputs → \`discard\` +- Wrong file or irrelevant access → \`discard\` +- Phase complete, want to condense the sequence → \`squash\` +- Exploration phase done, only need a summary → \`squash\` + +PRUNE METHODICALLY - BATCH YOUR ACTIONS +Every tool call adds to your context debt. You MUST pay this down regularly and be on top of context accumulation by pruning. Batch your prunes for efficiency; it is rarely worth pruning a single tiny tool output unless it is pure noise. Evaluate what SHOULD be pruned before jumping the gun. + +You WILL evaluate pruning when ANY of these are true: +- Phase is complete → use \`squash\` +- You accessed something that turned out to be irrelevant → use \`discard\` +- Information has been superseded by newer outputs → use \`discard\` + +You MUST NOT prune when: +- You need specific details from the content for upcoming work +- The content 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, keep it. Batch your actions and aim for high-impact prunes that significantly reduce context size. +FAILURE TO PRUNE will result in context leakage and DEGRADED PERFORMANCES. +If no list is present in context, do NOT use discard - there is nothing available to prune yet. +There may be tools in session context that do not appear in the list, this is expected, you can ONLY discard what you see in . + + + + +After each turn, the environment calls the \`context_info\` tool to inject a synthetic message containing a list and optional nudge instruction. This tool is only available to the environment - you do not have access to it. + +CRITICAL REQUIREMENTS - VIOLATION IS UNACCEPTABLE: +- 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 discard/squash tool output (e.g., "I've pruned the context", "Context cleanup complete") +- 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/system/discard.ts b/lib/prompts/system/discard.ts index e5cd77da..3877cd7c 100644 --- a/lib/prompts/system/discard.ts +++ b/lib/prompts/system/discard.ts @@ -7,19 +7,20 @@ You are operating in a context-constrained environment and thus must proactively IMPORTANT: The \`context_info\` tool is only available to the environment - you do not have access to it and must not attempt to call it. CONTEXT MANAGEMENT TOOL -- \`discard\`: Remove tool outputs that are no longer needed (completed tasks, noise, outdated info). No preservation of content. +- \`discard\`: Remove individual tool outputs that are noise, irrelevant, or superseded. No preservation of content. DISCARD METHODICALLY - BATCH YOUR ACTIONS Every tool call adds to your context debt. You MUST pay this down regularly and be on top of context accumulation by discarding. Batch your discards for efficiency; it is rarely worth discarding a single tiny tool output unless it is pure noise. Evaluate what SHOULD be discarded before jumping the gun. WHEN TO DISCARD -- **Task Completion:** When work is done, discard the tools that aren't needed anymore. - **Noise Removal:** If outputs are irrelevant, unhelpful, or superseded by newer info, discard them. +- **Wrong Files:** You read or accessed something that turned out to be irrelevant to the current work. +- **Outdated Info:** Outputs that have been superseded by newer information. You WILL evaluate discarding when ANY of these are true: -- Task or sub-task is complete +- You accessed something that turned out to be irrelevant +- Information has been superseded by newer outputs - You are about to start a new phase of work -- Write or edit operations are complete (discarding removes the large input content) You MUST NOT discard when: - The tool output will be needed for upcoming implementation work @@ -30,6 +31,7 @@ Discarding that forces you to re-call the same tool later is a net loss. Only di NOTES When in doubt, keep it. Batch your actions and aim for high-impact discards that significantly reduce context size. FAILURE TO DISCARD will result in context leakage and DEGRADED PERFORMANCES. +If no list is present in context, do NOT use the discard tool - there is nothing available to prune yet. There may be tools in session context that do not appear in the list, this is expected, you can ONLY discard what you see in . diff --git a/lib/prompts/system/extract-squash.ts b/lib/prompts/system/extract-squash.ts new file mode 100644 index 00000000..60b8e7a0 --- /dev/null +++ b/lib/prompts/system/extract-squash.ts @@ -0,0 +1,60 @@ +export const SYSTEM_PROMPT_EXTRACT_SQUASH = ` + + +ENVIRONMENT +You are operating in a context-constrained environment and thus must proactively manage your context window using the \`extract\` and \`squash\` tools. The environment calls the \`context_info\` tool to provide an up-to-date list after each turn. Use this information when deciding what to prune. + +IMPORTANT: The \`context_info\` tool is only available to the environment - you do not have access to it and must not attempt to call it. + +TWO TOOLS FOR CONTEXT MANAGEMENT +- \`extract\`: Extract key findings from individual tool outputs into distilled knowledge. Use when you need to preserve valuable technical details. +- \`squash\`: Collapse a contiguous range of conversation (completed phases) into a single summary. + +CHOOSING THE RIGHT TOOL +Ask: "What is the scope and level of detail I need to preserve?" +- **Individual tool outputs with detailed context to keep** → \`extract\` +- **Entire sequence (phase complete)** → \`squash\` + +Common scenarios: +- Large output with valuable technical details → \`extract\` +- Valuable context needed later but raw output too large → \`extract\` +- Phase complete, want to condense the sequence → \`squash\` +- Exploration phase done, only need a summary → \`squash\` + +PRUNE METHODICALLY - BATCH YOUR ACTIONS +Every tool call adds to your context debt. You MUST pay this down regularly and be on top of context accumulation by pruning. Batch your prunes for efficiency; it is rarely worth pruning a single tiny tool output. Evaluate what SHOULD be pruned before jumping the gun. + +You WILL evaluate pruning when ANY of these are true: +- Phase is complete → use \`squash\` +- You have large outputs with valuable details to preserve → use \`extract\` +- You are about to start a new phase of work + +You MUST NOT prune when: +- You need specific details from the content for upcoming work +- The content 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, keep it. Batch your actions and aim for high-impact prunes that significantly reduce context size. +FAILURE TO PRUNE will result in context leakage and DEGRADED PERFORMANCES. +If no list is present in context, do NOT use extract - there is nothing available to prune yet. +There may be tools in session context that do not appear in the list, this is expected, you can ONLY extract what you see in . + + + + +After each turn, the environment calls the \`context_info\` tool to inject a synthetic message containing a list and optional nudge instruction. This tool is only available to the environment - you do not have access to it. + +CRITICAL REQUIREMENTS - VIOLATION IS UNACCEPTABLE: +- 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 extract/squash tool output (e.g., "I've pruned the context", "Context cleanup complete") +- 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/system/extract.ts b/lib/prompts/system/extract.ts index 3f225e1e..9f024f51 100644 --- a/lib/prompts/system/extract.ts +++ b/lib/prompts/system/extract.ts @@ -7,19 +7,19 @@ You are operating in a context-constrained environment and thus must proactively IMPORTANT: The \`context_info\` tool is only available to the environment - you do not have access to it and must not attempt to call it. CONTEXT MANAGEMENT TOOL -- \`extract\`: Extract key findings from tools into distilled knowledge before removing the raw content from context. Use this to preserve important information while reducing context size. +- \`extract\`: Extract key findings from individual tool outputs into distilled knowledge before removing the raw content. Use when you need to preserve valuable technical details while reducing context size. EXTRACT METHODICALLY - BATCH YOUR ACTIONS Every tool call adds to your context debt. You MUST pay this down regularly and be on top of context accumulation by extracting. Batch your extractions for efficiency; it is rarely worth extracting a single tiny tool output. Evaluate what SHOULD be extracted before jumping the gun. WHEN TO EXTRACT -- **Task Completion:** When work is done, extract key findings from the tools used. Scale distillation depth to the value of the content. +- **Large Outputs:** The raw output is too large but contains valuable technical details worth keeping. - **Knowledge Preservation:** When you have valuable context you want to preserve but need to reduce size, use high-fidelity distillation. Your distillation must be comprehensive, capturing technical details (signatures, logic, constraints) such that the raw output is no longer needed. THINK: high signal, complete technical substitute. You WILL evaluate extracting when ANY of these are true: -- Task or sub-task is complete -- You are about to start a new phase of work -- Write or edit operations are complete (extracting removes the large input content) +- You have large tool outputs with valuable technical details +- You need to preserve specific information but reduce context size +- You are about to start a new phase of work and want to retain key insights You MUST NOT extract when: - The tool output will be needed for upcoming implementation work @@ -30,6 +30,7 @@ Extracting that forces you to re-call the same tool later is a net loss. Only ex NOTES When in doubt, keep it. Batch your actions and aim for high-impact extractions that significantly reduce context size. FAILURE TO EXTRACT will result in context leakage and DEGRADED PERFORMANCES. +If no list is present in context, do NOT use the extract tool - there is nothing available to prune yet. There may be tools in session context that do not appear in the list, this is expected, you can ONLY extract what you see in . diff --git a/lib/prompts/system/squash.ts b/lib/prompts/system/squash.ts new file mode 100644 index 00000000..494b5288 --- /dev/null +++ b/lib/prompts/system/squash.ts @@ -0,0 +1,50 @@ +export const SYSTEM_PROMPT_SQUASH = ` + + +ENVIRONMENT +You are operating in a context-constrained environment and thus must proactively manage your context window using the \`squash\` tool. The environment calls the \`context_info\` tool to provide an up-to-date list after each turn. Use this information when deciding what to squash. + +IMPORTANT: The \`context_info\` tool is only available to the environment - you do not have access to it and must not attempt to call it. + +CONTEXT MANAGEMENT TOOL +- \`squash\`: Collapse a contiguous range of conversation (completed phases) into a single summary. Use this when you want to condense an entire sequence of work. + +SQUASH METHODICALLY - BATCH YOUR ACTIONS +Every tool call adds to your context debt. You MUST pay this down regularly and be on top of context accumulation by squashing. Evaluate what SHOULD be squashed before jumping the gun. + +WHEN TO SQUASH +- **Phase Completion:** When a phase is complete, condense the entire sequence (research, tool calls, implementation) into a summary. +- **Exploration Done:** When you've explored multiple files or ran multiple commands and only need a summary of findings. + +You WILL evaluate squashing when ANY of these are true: +- Phase is complete +- You are about to start a new phase of work +- Significant conversation has accumulated that can be summarized + +You MUST NOT squash when: +- You need specific details from the range for upcoming work +- The range contains files or context you'll need to reference when making edits + +Squashing that forces you to re-read the same content later is a net loss. Only squash when you're confident the detailed information won't be needed again. + +NOTES +When in doubt, keep it. Aim for high-impact squashes that significantly reduce context size. +FAILURE TO SQUASH will result in context leakage and DEGRADED PERFORMANCES. + + + + +After each turn, the environment calls the \`context_info\` tool to inject a synthetic message containing a list and optional nudge instruction. This tool is only available to the environment - you do not have access to it. + +CRITICAL REQUIREMENTS - VIOLATION IS UNACCEPTABLE: +- NEVER reference the squash encouragement or context management instructions. Do not reply with "I agree" or "Great idea" when the squash encouragement appears. +- NEVER acknowledge the list (e.g., "I see the prunable tools list", "Looking at the available tools to squash") +- NEVER reference the nudge instruction (e.g., "As the nudge suggests", "The system is reminding me to squash") +- NEVER acknowledge squash tool output (e.g., "I've squashed the context", "Context cleanup complete") +- 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/shared-utils.ts b/lib/shared-utils.ts index 902ea403..df0fceef 100644 --- a/lib/shared-utils.ts +++ b/lib/shared-utils.ts @@ -2,11 +2,21 @@ import { SessionState, WithParts } from "./state" import { isIgnoredUserMessage } from "./messages/utils" export const isMessageCompacted = (state: SessionState, msg: WithParts): boolean => { - return msg.info.time.created < state.lastCompaction + if (msg.info.time.created < state.lastCompaction) { + return true + } + if (state.prune.messageIds.includes(msg.info.id)) { + return true + } + return false } -export const getLastUserMessage = (messages: WithParts[]): WithParts | null => { - for (let i = messages.length - 1; i >= 0; i--) { +export const getLastUserMessage = ( + messages: WithParts[], + startIndex?: number, +): WithParts | null => { + const start = startIndex ?? messages.length - 1 + for (let i = start; i >= 0; i--) { const msg = messages[i] if (msg.info.role === "user" && !isIgnoredUserMessage(msg)) { return msg diff --git a/lib/state/persistence.ts b/lib/state/persistence.ts index 172ff75f..91111ef7 100644 --- a/lib/state/persistence.ts +++ b/lib/state/persistence.ts @@ -8,12 +8,13 @@ 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, SquashSummary } from "./types" import type { Logger } from "../logger" export interface PersistedSessionState { sessionName?: string prune: Prune + squashSummaries: SquashSummary[] stats: SessionStats lastUpdated: string } @@ -45,6 +46,7 @@ export async function saveSessionState( const state: PersistedSessionState = { sessionName: sessionName, prune: sessionState.prune, + squashSummaries: sessionState.squashSummaries, stats: sessionState.stats, lastUpdated: new Date().toISOString(), } diff --git a/lib/state/state.ts b/lib/state/state.ts index 69add020..98a99693 100644 --- a/lib/state/state.ts +++ b/lib/state/state.ts @@ -1,8 +1,13 @@ import type { SessionState, ToolParameterEntry, WithParts } from "./types" import type { Logger } from "../logger" import { loadSessionState } from "./persistence" -import { isSubAgentSession } from "./utils" -import { getLastUserMessage, isMessageCompacted } from "../shared-utils" +import { + isSubAgentSession, + findLastCompactionTimestamp, + countTurns, + resetOnCompaction, +} from "./utils" +import { getLastUserMessage } from "../shared-utils" export const checkSession = async ( client: any, @@ -29,9 +34,8 @@ export const checkSession = async ( const lastCompactionTimestamp = findLastCompactionTimestamp(messages) if (lastCompactionTimestamp > state.lastCompaction) { state.lastCompaction = lastCompactionTimestamp - state.toolParameters.clear() - state.prune.toolIds = [] - logger.info("Detected compaction from messages - cleared tool cache", { + resetOnCompaction(state) + logger.info("Detected compaction - reset stale state", { timestamp: lastCompactionTimestamp, }) } @@ -45,7 +49,9 @@ export function createSessionState(): SessionState { isSubAgent: false, prune: { toolIds: [], + messageIds: [], }, + squashSummaries: [], stats: { pruneTokenCounter: 0, totalPruneTokens: 0, @@ -64,7 +70,9 @@ export function resetSessionState(state: SessionState): void { state.isSubAgent = false state.prune = { toolIds: [], + messageIds: [], } + state.squashSummaries = [] state.stats = { pruneTokenCounter: 0, totalPruneTokens: 0, @@ -108,35 +116,11 @@ export async function ensureSessionInitialized( state.prune = { toolIds: persisted.prune.toolIds || [], + messageIds: persisted.prune.messageIds || [], } + state.squashSummaries = persisted.squashSummaries || [] state.stats = { pruneTokenCounter: persisted.stats?.pruneTokenCounter || 0, totalPruneTokens: persisted.stats?.totalPruneTokens || 0, } } - -function findLastCompactionTimestamp(messages: WithParts[]): number { - for (let i = messages.length - 1; i >= 0; i--) { - const msg = messages[i] - if (msg.info.role === "assistant" && msg.info.summary === true) { - return msg.info.time.created - } - } - return 0 -} - -export function countTurns(state: SessionState, messages: WithParts[]): number { - let turnCount = 0 - for (const msg of messages) { - if (isMessageCompacted(state, msg)) { - continue - } - const parts = Array.isArray(msg.parts) ? msg.parts : [] - for (const part of parts) { - if (part.type === "step-start") { - turnCount++ - } - } - } - return turnCount -} diff --git a/lib/state/tool-cache.ts b/lib/state/tool-cache.ts index 38d3b54b..80837519 100644 --- a/lib/state/tool-cache.ts +++ b/lib/state/tool-cache.ts @@ -44,12 +44,14 @@ export async function syncToolCache( state.currentTurn - turnCounter < turnProtectionTurns state.lastToolPrune = - (part.tool === "discard" || part.tool === "extract") && + (part.tool === "discard" || + part.tool === "extract" || + part.tool === "squash") && part.state.status === "completed" const allProtectedTools = config.tools.settings.protectedTools - if (part.tool === "discard" || part.tool === "extract") { + if (part.tool === "discard" || part.tool === "extract" || part.tool === "squash") { state.nudgeCounter = 0 } else if (!allProtectedTools.includes(part.tool) && !isProtectedByTurn) { state.nudgeCounter++ diff --git a/lib/state/types.ts b/lib/state/types.ts index 1e41170d..330f8c89 100644 --- a/lib/state/types.ts +++ b/lib/state/types.ts @@ -20,14 +20,21 @@ export interface SessionStats { totalPruneTokens: number } +export interface SquashSummary { + anchorMessageId: string + summary: string +} + export interface Prune { toolIds: string[] + messageIds: string[] } export interface SessionState { sessionId: string | null isSubAgent: boolean prune: Prune + squashSummaries: SquashSummary[] stats: SessionStats toolParameters: Map nudgeCounter: number diff --git a/lib/state/utils.ts b/lib/state/utils.ts index 4cc10ce1..be8a08fe 100644 --- a/lib/state/utils.ts +++ b/lib/state/utils.ts @@ -1,3 +1,6 @@ +import type { SessionState, WithParts } from "./types" +import { isMessageCompacted } from "../shared-utils" + export async function isSubAgentSession(client: any, sessionID: string): Promise { try { const result = await client.session.get({ path: { id: sessionID } }) @@ -6,3 +9,38 @@ export async function isSubAgentSession(client: any, sessionID: string): Promise return false } } + +export function findLastCompactionTimestamp(messages: WithParts[]): number { + for (let i = messages.length - 1; i >= 0; i--) { + const msg = messages[i] + if (msg.info.role === "assistant" && msg.info.summary === true) { + return msg.info.time.created + } + } + return 0 +} + +export function countTurns(state: SessionState, messages: WithParts[]): number { + let turnCount = 0 + for (const msg of messages) { + if (isMessageCompacted(state, msg)) { + continue + } + const parts = Array.isArray(msg.parts) ? msg.parts : [] + for (const part of parts) { + if (part.type === "step-start") { + turnCount++ + } + } + } + return turnCount +} + +export function resetOnCompaction(state: SessionState): void { + state.toolParameters.clear() + state.prune.toolIds = [] + state.prune.messageIds = [] + state.squashSummaries = [] + state.nudgeCounter = 0 + state.lastToolPrune = false +} diff --git a/lib/strategies/index.ts b/lib/strategies/index.ts index 5444964c..a995254e 100644 --- a/lib/strategies/index.ts +++ b/lib/strategies/index.ts @@ -1,4 +1,4 @@ export { deduplicate } from "./deduplication" -export { createDiscardTool, createExtractTool } from "./tools" +export { createDiscardTool, createExtractTool, createSquashTool } from "../tools" export { supersedeWrites } from "./supersede-writes" export { purgeErrors } from "./purge-errors" diff --git a/lib/strategies/tools.ts b/lib/strategies/tools.ts deleted file mode 100644 index 44f6742f..00000000 --- a/lib/strategies/tools.ts +++ /dev/null @@ -1,220 +0,0 @@ -import { tool } from "@opencode-ai/plugin" -import type { SessionState, ToolParameterEntry, WithParts } from "../state" -import type { PluginConfig } from "../config" -import { buildToolIdList } from "../messages/utils" -import { PruneReason, sendUnifiedNotification } from "../ui/notification" -import { formatPruningResultForTool } from "../ui/utils" -import { ensureSessionInitialized } from "../state" -import { saveSessionState } from "../state/persistence" -import type { Logger } from "../logger" -import { loadPrompt } from "../prompts" -import { calculateTokensSaved, getCurrentParams } from "./utils" -import { getFilePathFromParameters, isProtectedFilePath } from "../protected-file-patterns" - -const DISCARD_TOOL_DESCRIPTION = loadPrompt("discard-tool-spec") -const EXTRACT_TOOL_DESCRIPTION = loadPrompt("extract-tool-spec") - -export interface PruneToolContext { - client: any - state: SessionState - logger: Logger - config: PluginConfig - workingDirectory: string -} - -// Shared logic for executing prune operations. -async function executePruneOperation( - ctx: PruneToolContext, - toolCtx: { sessionID: string }, - ids: string[], - reason: PruneReason, - toolName: string, - distillation?: string[], -): Promise { - const { client, state, logger, config, workingDirectory } = ctx - const sessionId = toolCtx.sessionID - - logger.info(`${toolName} tool invoked`) - logger.info(JSON.stringify(reason ? { ids, reason } : { ids })) - - if (!ids || ids.length === 0) { - logger.debug(`${toolName} tool called but ids is empty or undefined`) - throw new Error( - `No IDs provided. Check the list for available IDs to ${toolName.toLowerCase()}.`, - ) - } - - const numericToolIds: number[] = 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 ${toolName}: ` + JSON.stringify(ids)) - throw new Error("No numeric IDs provided. Format: ids: [id1, id2, ...]") - } - - // Fetch messages to calculate tokens and find current agent - const messagesResponse = await client.session.messages({ - path: { id: sessionId }, - }) - const messages: WithParts[] = messagesResponse.data || messagesResponse - - await ensureSessionInitialized(ctx.client, state, sessionId, logger, messages) - - const currentParams = getCurrentParams(state, 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(", ")) - throw new Error( - "Invalid IDs provided. Only use numeric IDs from the list.", - ) - } - - // Validate that all IDs exist in cache and aren't protected - // (rejects hallucinated IDs and turn-protected tools not shown in ) - for (const index of numericToolIds) { - const id = toolIdList[index] - const metadata = state.toolParameters.get(id) - if (!metadata) { - logger.debug( - "Rejecting prune request - ID not in cache (turn-protected or hallucinated)", - { index, id }, - ) - throw new Error( - "Invalid IDs provided. Only use numeric IDs from the list.", - ) - } - const allProtectedTools = config.tools.settings.protectedTools - if (allProtectedTools.includes(metadata.tool)) { - logger.debug("Rejecting prune request - protected tool", { - index, - id, - tool: metadata.tool, - }) - throw new Error( - "Invalid IDs provided. Only use numeric IDs from the list.", - ) - } - - const filePath = getFilePathFromParameters(metadata.parameters) - if (isProtectedFilePath(filePath, config.protectedFilePatterns)) { - logger.debug("Rejecting prune request - protected file path", { - index, - id, - tool: metadata.tool, - filePath, - }) - throw new Error( - "Invalid IDs provided. Only use numeric IDs from the list.", - ) - } - } - - const pruneToolIds: string[] = numericToolIds.map((index) => toolIdList[index]) - state.prune.toolIds.push(...pruneToolIds) - - const toolMetadata = new Map() - for (const id of pruneToolIds) { - const toolParameters = state.toolParameters.get(id) - if (toolParameters) { - toolMetadata.set(id, toolParameters) - } else { - logger.debug("No metadata found for ID", { id }) - } - } - - state.stats.pruneTokenCounter += calculateTokensSaved(state, messages, pruneToolIds) - - await sendUnifiedNotification( - client, - logger, - config, - state, - sessionId, - pruneToolIds, - toolMetadata, - reason, - currentParams, - workingDirectory, - distillation, - ) - - state.stats.totalPruneTokens += state.stats.pruneTokenCounter - state.stats.pruneTokenCounter = 0 - state.nudgeCounter = 0 - - saveSessionState(state, logger).catch((err) => - logger.error("Failed to persist state", { error: err.message }), - ) - - return formatPruningResultForTool(pruneToolIds, toolMetadata, workingDirectory) -} - -export function createDiscardTool(ctx: PruneToolContext): ReturnType { - return tool({ - description: DISCARD_TOOL_DESCRIPTION, - args: { - ids: tool.schema - .array(tool.schema.string()) - .describe( - "First element is the reason ('completion' or 'noise'), followed by numeric IDs as strings to discard", - ), - }, - async execute(args, toolCtx) { - // Parse reason from first element, numeric IDs from the rest - const reason = args.ids?.[0] - const validReasons = ["completion", "noise"] as const - if (typeof reason !== "string" || !validReasons.includes(reason as any)) { - ctx.logger.debug("Invalid discard reason provided: " + reason) - throw new Error( - "No valid reason found. Use 'completion' or 'noise' as the first element.", - ) - } - - const numericIds = args.ids.slice(1) - - return executePruneOperation(ctx, toolCtx, numericIds, reason as PruneReason, "Discard") - }, - }) -} - -export function createExtractTool(ctx: PruneToolContext): ReturnType { - return tool({ - description: EXTRACT_TOOL_DESCRIPTION, - args: { - ids: tool.schema - .array(tool.schema.string()) - .describe("Numeric IDs as strings to extract from the list"), - distillation: tool.schema - .array(tool.schema.string()) - .describe( - "REQUIRED. Array of strings, one per ID (positional: distillation[0] is for ids[0], etc.)", - ), - }, - async execute(args, toolCtx) { - if (!args.distillation || args.distillation.length === 0) { - ctx.logger.debug( - "Extract tool called without distillation: " + JSON.stringify(args), - ) - throw new Error( - "Missing distillation. You must provide a distillation string for each ID.", - ) - } - - // Log the distillation for debugging/analysis - ctx.logger.info("Distillation data received:") - ctx.logger.info(JSON.stringify(args.distillation, null, 2)) - - return executePruneOperation( - ctx, - toolCtx, - args.ids, - "extraction" as PruneReason, - "Extract", - args.distillation, - ) - }, - }) -} diff --git a/lib/strategies/utils.ts b/lib/strategies/utils.ts index 7ae04154..c32ba727 100644 --- a/lib/strategies/utils.ts +++ b/lib/strategies/utils.ts @@ -42,8 +42,9 @@ export function countTokens(text: string): number { } } -function estimateTokensBatch(texts: string[]): number[] { - return texts.map(countTokens) +export function estimateTokensBatch(texts: string[]): number { + if (texts.length === 0) return 0 + return countTokens(texts.join(" ")) } export const calculateTokensSaved = ( @@ -86,8 +87,7 @@ export const calculateTokensSaved = ( } } } - const tokenCounts: number[] = estimateTokensBatch(contents) - return tokenCounts.reduce((sum, count) => sum + count, 0) + return estimateTokensBatch(contents) } catch (error: any) { return 0 } diff --git a/lib/tools/discard.ts b/lib/tools/discard.ts new file mode 100644 index 00000000..2872eaca --- /dev/null +++ b/lib/tools/discard.ts @@ -0,0 +1,25 @@ +import { tool } from "@opencode-ai/plugin" +import type { PruneToolContext } from "./types" +import { executePruneOperation } from "./prune-shared" +import { PruneReason } from "../ui/notification" +import { loadPrompt } from "../prompts" + +const DISCARD_TOOL_DESCRIPTION = loadPrompt("discard-tool-spec") + +export function createDiscardTool(ctx: PruneToolContext): ReturnType { + return tool({ + description: DISCARD_TOOL_DESCRIPTION, + args: { + ids: tool.schema + .array(tool.schema.string()) + .min(1) + .describe("Numeric IDs as strings from the list to discard"), + }, + async execute(args, toolCtx) { + const numericIds = args.ids + const reason = "noise" + + return executePruneOperation(ctx, toolCtx, numericIds, reason, "Discard") + }, + }) +} diff --git a/lib/tools/extract.ts b/lib/tools/extract.ts new file mode 100644 index 00000000..15e5d7c8 --- /dev/null +++ b/lib/tools/extract.ts @@ -0,0 +1,58 @@ +import { tool } from "@opencode-ai/plugin" +import type { PruneToolContext } from "./types" +import { executePruneOperation } from "./prune-shared" +import { PruneReason } from "../ui/notification" +import { loadPrompt } from "../prompts" + +const EXTRACT_TOOL_DESCRIPTION = loadPrompt("extract-tool-spec") + +export function createExtractTool(ctx: PruneToolContext): ReturnType { + return tool({ + description: EXTRACT_TOOL_DESCRIPTION, + args: { + ids: tool.schema + .array(tool.schema.string()) + .min(1) + .describe("Numeric IDs as strings to extract from the list"), + distillation: tool.schema + .array(tool.schema.string()) + .min(1) + .describe( + "Required array of distillation strings, one per ID (positional: distillation[0] for ids[0], etc.)", + ), + }, + async execute(args, toolCtx) { + if (!args.distillation || args.distillation.length === 0) { + ctx.logger.debug( + "Extract tool called without distillation: " + JSON.stringify(args), + ) + throw new Error( + "Missing distillation. You must provide a distillation string for each ID.", + ) + } + + if (!Array.isArray(args.distillation)) { + ctx.logger.debug( + "Extract tool called with non-array distillation: " + JSON.stringify(args), + ) + throw new Error( + `Invalid distillation format: expected an array of strings, got ${typeof args.distillation}. ` + + `Example: distillation: ["summary for id 0", "summary for id 1"]`, + ) + } + + // Log the distillation for debugging/analysis + ctx.logger.info("Distillation data received:") + ctx.logger.info(JSON.stringify(args.distillation, null, 2)) + + return executePruneOperation( + ctx, + toolCtx, + args.ids, + "extraction" as PruneReason, + "Extract", + args.distillation, + ) + }, + }) +} diff --git a/lib/tools/index.ts b/lib/tools/index.ts new file mode 100644 index 00000000..9f61b15d --- /dev/null +++ b/lib/tools/index.ts @@ -0,0 +1,4 @@ +export { PruneToolContext } from "./types" +export { createDiscardTool } from "./discard" +export { createExtractTool } from "./extract" +export { createSquashTool } from "./squash" diff --git a/lib/tools/prune-shared.ts b/lib/tools/prune-shared.ts new file mode 100644 index 00000000..ba37fb78 --- /dev/null +++ b/lib/tools/prune-shared.ts @@ -0,0 +1,159 @@ +import type { SessionState, ToolParameterEntry, WithParts } from "../state" +import type { PluginConfig } from "../config" +import type { Logger } from "../logger" +import type { PruneToolContext } from "./types" +import { buildToolIdList } from "../messages/utils" +import { syncToolCache } from "../state/tool-cache" +import { PruneReason, sendUnifiedNotification } from "../ui/notification" +import { formatPruningResultForTool } from "../ui/utils" +import { ensureSessionInitialized } from "../state" +import { saveSessionState } from "../state/persistence" +import { calculateTokensSaved, getCurrentParams } from "../strategies/utils" +import { getFilePathFromParameters, isProtectedFilePath } from "../protected-file-patterns" + +// Shared logic for executing prune operations. +export async function executePruneOperation( + ctx: PruneToolContext, + toolCtx: { sessionID: string }, + ids: string[], + reason: PruneReason, + toolName: string, + distillation?: string[], +): Promise { + const { client, state, logger, config, workingDirectory } = ctx + const sessionId = toolCtx.sessionID + + logger.info(`${toolName} tool invoked`) + logger.info(JSON.stringify(reason ? { ids, reason } : { ids })) + + if (!ids || ids.length === 0) { + logger.debug(`${toolName} tool called but ids is empty or undefined`) + throw new Error( + `No IDs provided. Check the list for available IDs to ${toolName.toLowerCase()}.`, + ) + } + + const numericToolIds: number[] = 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 ${toolName}: ` + JSON.stringify(ids)) + throw new Error("No numeric IDs provided. Format: ids: [id1, id2, ...]") + } + + // Fetch messages to calculate tokens and find current agent + const messagesResponse = await client.session.messages({ + path: { id: sessionId }, + }) + const messages: WithParts[] = messagesResponse.data || messagesResponse + + await ensureSessionInitialized(ctx.client, state, sessionId, logger, messages) + await syncToolCache(state, config, logger, messages) + + const currentParams = getCurrentParams(state, messages, logger) + const toolIdList: string[] = buildToolIdList(state, messages, logger) + + const validNumericIds: number[] = [] + const skippedIds: string[] = [] + + // Validate and filter IDs + for (const index of numericToolIds) { + // Validate that all numeric IDs are within bounds + if (index < 0 || index >= toolIdList.length) { + logger.debug(`Rejecting prune request - index out of bounds: ${index}`) + skippedIds.push(index.toString()) + continue + } + + const id = toolIdList[index] + const metadata = state.toolParameters.get(id) + + // Validate that all IDs exist in cache and aren't protected + // (rejects hallucinated IDs and turn-protected tools not shown in ) + if (!metadata) { + logger.debug( + "Rejecting prune request - ID not in cache (turn-protected or hallucinated)", + { index, id }, + ) + skippedIds.push(index.toString()) + continue + } + + const allProtectedTools = config.tools.settings.protectedTools + if (allProtectedTools.includes(metadata.tool)) { + logger.debug("Rejecting prune request - protected tool", { + index, + id, + tool: metadata.tool, + }) + skippedIds.push(index.toString()) + continue + } + + const filePath = getFilePathFromParameters(metadata.parameters) + if (isProtectedFilePath(filePath, config.protectedFilePatterns)) { + logger.debug("Rejecting prune request - protected file path", { + index, + id, + tool: metadata.tool, + filePath, + }) + skippedIds.push(index.toString()) + continue + } + + validNumericIds.push(index) + } + + if (validNumericIds.length === 0) { + const errorMsg = + skippedIds.length > 0 + ? `Invalid IDs provided: [${skippedIds.join(", ")}]. Only use numeric IDs from the list.` + : `No valid IDs provided to ${toolName.toLowerCase()}.` + throw new Error(errorMsg) + } + + const pruneToolIds: string[] = validNumericIds.map((index) => toolIdList[index]) + state.prune.toolIds.push(...pruneToolIds) + + const toolMetadata = new Map() + for (const id of pruneToolIds) { + const toolParameters = state.toolParameters.get(id) + if (toolParameters) { + toolMetadata.set(id, toolParameters) + } else { + logger.debug("No metadata found for ID", { id }) + } + } + + state.stats.pruneTokenCounter += calculateTokensSaved(state, messages, pruneToolIds) + + await sendUnifiedNotification( + client, + logger, + config, + state, + sessionId, + pruneToolIds, + toolMetadata, + reason, + currentParams, + workingDirectory, + distillation, + ) + + state.stats.totalPruneTokens += state.stats.pruneTokenCounter + state.stats.pruneTokenCounter = 0 + state.nudgeCounter = 0 + + saveSessionState(state, logger).catch((err) => + logger.error("Failed to persist state", { error: err.message }), + ) + + let result = formatPruningResultForTool(pruneToolIds, toolMetadata, workingDirectory) + if (skippedIds.length > 0) { + result += `\n\nNote: ${skippedIds.length} IDs were skipped (invalid, protected, or missing metadata): ${skippedIds.join(", ")}` + } + return result +} diff --git a/lib/tools/squash.ts b/lib/tools/squash.ts new file mode 100644 index 00000000..9ab9425a --- /dev/null +++ b/lib/tools/squash.ts @@ -0,0 +1,155 @@ +import { tool } from "@opencode-ai/plugin" +import type { WithParts, SquashSummary } from "../state" +import type { PruneToolContext } from "./types" +import { ensureSessionInitialized } from "../state" +import { saveSessionState } from "../state/persistence" +import { loadPrompt } from "../prompts" +import { estimateTokensBatch, getCurrentParams } from "../strategies/utils" +import { + collectContentInRange, + findStringInMessages, + collectToolIdsInRange, + collectMessageIdsInRange, +} from "./utils" +import { sendSquashNotification } from "../ui/notification" + +const SQUASH_TOOL_DESCRIPTION = loadPrompt("squash-tool-spec") + +export function createSquashTool(ctx: PruneToolContext): ReturnType { + return tool({ + description: SQUASH_TOOL_DESCRIPTION, + args: { + input: tool.schema + .array(tool.schema.string()) + .length(4) + .describe( + "[startString, endString, topic, summary] - 4 required strings: (1) startString: unique text from conversation marking range start, (2) endString: unique text marking range end, (3) topic: short 3-5 word label for UI, (4) summary: comprehensive text replacing all squashed content", + ), + }, + async execute(args, toolCtx) { + const { client, state, logger } = ctx + const sessionId = toolCtx.sessionID + + const [startString, endString, topic, summary] = args.input + + logger.info("Squash tool invoked") + // logger.info( + // JSON.stringify({ + // startString: startString?.substring(0, 50) + "...", + // endString: endString?.substring(0, 50) + "...", + // topic: topic, + // summaryLength: summary?.length, + // }), + // ) + + const messagesResponse = await client.session.messages({ + path: { id: sessionId }, + }) + const messages: WithParts[] = messagesResponse.data || messagesResponse + + await ensureSessionInitialized(client, state, sessionId, logger, messages) + + const startResult = findStringInMessages( + messages, + startString, + logger, + state.squashSummaries, + "startString", + ) + const endResult = findStringInMessages( + messages, + endString, + logger, + state.squashSummaries, + "endString", + ) + + if (startResult.messageIndex > endResult.messageIndex) { + throw new Error( + `startString appears after endString in the conversation. Start must come before end.`, + ) + } + + const containedToolIds = collectToolIdsInRange( + messages, + startResult.messageIndex, + endResult.messageIndex, + ) + + const containedMessageIds = collectMessageIdsInRange( + messages, + startResult.messageIndex, + endResult.messageIndex, + ) + + state.prune.toolIds.push(...containedToolIds) + state.prune.messageIds.push(...containedMessageIds) + + // Remove any existing summaries whose anchors are now inside this range + // This prevents duplicate injections when a larger squash subsumes a smaller one + const removedSummaries = state.squashSummaries.filter((s) => + containedMessageIds.includes(s.anchorMessageId), + ) + if (removedSummaries.length > 0) { + // logger.info("Removing subsumed squash summaries", { + // count: removedSummaries.length, + // anchorIds: removedSummaries.map((s) => s.anchorMessageId), + // }) + state.squashSummaries = state.squashSummaries.filter( + (s) => !containedMessageIds.includes(s.anchorMessageId), + ) + } + + const squashSummary: SquashSummary = { + anchorMessageId: startResult.messageId, + summary: summary, + } + state.squashSummaries.push(squashSummary) + + const contentsToTokenize = collectContentInRange( + messages, + startResult.messageIndex, + endResult.messageIndex, + ) + const estimatedSquashedTokens = estimateTokensBatch(contentsToTokenize) + + state.stats.pruneTokenCounter += estimatedSquashedTokens + + const currentParams = getCurrentParams(state, messages, logger) + await sendSquashNotification( + client, + logger, + ctx.config, + state, + sessionId, + containedToolIds, + containedMessageIds, + topic, + summary, + startResult, + endResult, + messages.length, + currentParams, + ) + + state.stats.totalPruneTokens += state.stats.pruneTokenCounter + state.stats.pruneTokenCounter = 0 + state.nudgeCounter = 0 + + // logger.info("Squash range created", { + // startMessageId: startResult.messageId, + // endMessageId: endResult.messageId, + // toolIdsRemoved: containedToolIds.length, + // messagesInRange: containedMessageIds.length, + // estimatedTokens: estimatedSquashedTokens, + // }) + + saveSessionState(state, logger).catch((err) => + logger.error("Failed to persist state", { error: err.message }), + ) + + const messagesSquashed = endResult.messageIndex - startResult.messageIndex + 1 + return `Squashed ${messagesSquashed} messages (${containedToolIds.length} tool calls) into summary. The content will be replaced with your summary.` + }, + }) +} diff --git a/lib/tools/types.ts b/lib/tools/types.ts new file mode 100644 index 00000000..c4950e47 --- /dev/null +++ b/lib/tools/types.ts @@ -0,0 +1,11 @@ +import type { SessionState } from "../state" +import type { PluginConfig } from "../config" +import type { Logger } from "../logger" + +export interface PruneToolContext { + client: any + state: SessionState + logger: Logger + config: PluginConfig + workingDirectory: string +} diff --git a/lib/tools/utils.ts b/lib/tools/utils.ts new file mode 100644 index 00000000..d5e4e180 --- /dev/null +++ b/lib/tools/utils.ts @@ -0,0 +1,165 @@ +import type { WithParts, SquashSummary } from "../state" +import type { Logger } from "../logger" + +/** + * Searches messages for a string and returns the message ID where it's found. + * Searches in text parts, tool outputs, tool inputs, and other textual content. + * Also searches through existing squash summaries to enable chained squashing. + * Throws an error if the string is not found or found more than once. + */ +export function findStringInMessages( + messages: WithParts[], + searchString: string, + logger: Logger, + squashSummaries: SquashSummary[] = [], + stringType: "startString" | "endString", +): { messageId: string; messageIndex: number } { + const matches: { messageId: string; messageIndex: number }[] = [] + + // First, search through existing squash summaries + // This allows referencing text from previous squash operations + for (const summary of squashSummaries) { + if (summary.summary.includes(searchString)) { + const anchorIndex = messages.findIndex((m) => m.info.id === summary.anchorMessageId) + if (anchorIndex !== -1) { + matches.push({ + messageId: summary.anchorMessageId, + messageIndex: anchorIndex, + }) + } + } + } + + // Then search through raw messages + for (let i = 0; i < messages.length; i++) { + const msg = messages[i] + const parts = Array.isArray(msg.parts) ? msg.parts : [] + + for (const part of parts) { + let content = "" + + if (part.type === "text" && typeof part.text === "string") { + content = part.text + } else if (part.type === "tool" && part.state?.status === "completed") { + if (typeof part.state.output === "string") { + content = part.state.output + } + if (part.state.input) { + const inputStr = + typeof part.state.input === "string" + ? part.state.input + : JSON.stringify(part.state.input) + content += " " + inputStr + } + } + + if (content.includes(searchString)) { + matches.push({ messageId: msg.info.id, messageIndex: i }) + } + } + } + + if (matches.length === 0) { + throw new Error( + `${stringType} not found in conversation. Make sure the string exists and is spelled exactly as it appears.`, + ) + } + + if (matches.length > 1) { + throw new Error( + `Found multiple matches for ${stringType}. Provide more surrounding context to uniquely identify the intended match.`, + ) + } + + return matches[0] +} + +/** + * Collects all tool callIDs from messages between start and end indices (inclusive). + */ +export function collectToolIdsInRange( + messages: WithParts[], + startIndex: number, + endIndex: number, +): string[] { + const toolIds: string[] = [] + + for (let i = startIndex; i <= endIndex; i++) { + const msg = messages[i] + const parts = Array.isArray(msg.parts) ? msg.parts : [] + + for (const part of parts) { + if (part.type === "tool" && part.callID) { + if (!toolIds.includes(part.callID)) { + toolIds.push(part.callID) + } + } + } + } + + return toolIds +} + +/** + * Collects all message IDs from messages between start and end indices (inclusive). + */ +export function collectMessageIdsInRange( + messages: WithParts[], + startIndex: number, + endIndex: number, +): string[] { + const messageIds: string[] = [] + + for (let i = startIndex; i <= endIndex; i++) { + const msgId = messages[i].info.id + if (!messageIds.includes(msgId)) { + messageIds.push(msgId) + } + } + + return messageIds +} + +/** + * Collects all textual content (text parts, tool inputs, and tool outputs) + * from a range of messages. Used for token estimation. + */ +export function collectContentInRange( + messages: WithParts[], + startIndex: number, + endIndex: number, +): string[] { + const contents: string[] = [] + for (let i = startIndex; i <= endIndex; i++) { + const msg = messages[i] + const parts = Array.isArray(msg.parts) ? msg.parts : [] + for (const part of parts) { + if (part.type === "text") { + contents.push(part.text) + } else if (part.type === "tool") { + const toolState = part.state as any + if (toolState?.input) { + contents.push( + typeof toolState.input === "string" + ? toolState.input + : JSON.stringify(toolState.input), + ) + } + if (toolState?.status === "completed" && toolState?.output) { + contents.push( + typeof toolState.output === "string" + ? toolState.output + : JSON.stringify(toolState.output), + ) + } else if (toolState?.status === "error" && toolState?.error) { + contents.push( + typeof toolState.error === "string" + ? toolState.error + : JSON.stringify(toolState.error), + ) + } + } + } + } + return contents +} diff --git a/lib/ui/notification.ts b/lib/ui/notification.ts index acb948cd..a85b0737 100644 --- a/lib/ui/notification.ts +++ b/lib/ui/notification.ts @@ -6,6 +6,8 @@ import { formatPrunedItemsList, formatStatsHeader, formatTokenCount, + formatProgressBar, + truncate, } from "./utils" import { ToolParameterEntry } from "../state" import { PluginConfig } from "../config" @@ -17,6 +19,70 @@ export const PRUNE_REASON_LABELS: Record = { extraction: "Extraction", } +const TOAST_PRUNED_ITEMS_LIMIT = 9 +const TOAST_TEXT_LIMIT = 600 + +function buildToastBody(message: string, header: string): string { + let toastBody = message.startsWith(header) ? message.slice(header.length).trim() : message + + const lines = toastBody.split("\n") + const pruneIndex = lines.findIndex((line) => line.startsWith("▣ Pruning")) + if (pruneIndex >= 0) { + const itemStart = pruneIndex + 1 + let itemEnd = itemStart + while (itemEnd < lines.length && lines[itemEnd].startsWith("→ ")) { + itemEnd++ + } + const itemLines = lines.slice(itemStart, itemEnd) + if (itemLines.length > TOAST_PRUNED_ITEMS_LIMIT) { + const remaining = itemLines.length - TOAST_PRUNED_ITEMS_LIMIT + lines.splice(itemStart, itemLines.length, ...itemLines.slice(0, TOAST_PRUNED_ITEMS_LIMIT), `... and ${remaining} more`) + toastBody = lines.join("\n") + } + } + + for (const marker of ["▣ Extracted", "→ Summary: "]) { + const markerIndex = toastBody.indexOf(`\n${marker}`) + if (markerIndex >= 0) { + const contentStart = markerIndex + marker.length + 1 + const content = toastBody.slice(contentStart) + const leading = content.match(/^\s*/)?.[0] || "" + const trimmedContent = content.slice(leading.length) + if (trimmedContent.length > TOAST_TEXT_LIMIT) { + toastBody = toastBody.slice(0, contentStart) + leading + truncate(trimmedContent, TOAST_TEXT_LIMIT) + } + } + } + + return toastBody +} + +function buildPruneDetails( + state: SessionState, + reason: PruneReason | undefined, + pruneToolIds: string[], + toolMetadata: Map, + workingDirectory: string, + distillation: string[] | undefined, +): string { + if (pruneToolIds.length === 0) { + return "" + } + + const pruneTokenCounterStr = `~${formatTokenCount(state.stats.pruneTokenCounter)}` + const extractedTokens = countDistillationTokens(distillation) + const extractedSuffix = + extractedTokens > 0 ? `, extracted ${formatTokenCount(extractedTokens)}` : "" + const reasonLabel = reason && extractedTokens === 0 ? ` — ${PRUNE_REASON_LABELS[reason]}` : "" + + let message = `▣ Pruning (${pruneTokenCounterStr}${extractedSuffix})${reasonLabel}` + + const itemLines = formatPrunedItemsList(pruneToolIds, toolMetadata, workingDirectory) + message += "\n" + itemLines.join("\n") + + return message +} + function buildMinimalMessage( state: SessionState, reason: PruneReason | undefined, @@ -46,17 +112,17 @@ function buildDetailedMessage( ): string { let message = formatStatsHeader(state.stats.totalPruneTokens, state.stats.pruneTokenCounter) - if (pruneToolIds.length > 0) { - const pruneTokenCounterStr = `~${formatTokenCount(state.stats.pruneTokenCounter)}` - const extractedTokens = countDistillationTokens(distillation) - const extractedSuffix = - extractedTokens > 0 ? `, extracted ${formatTokenCount(extractedTokens)}` : "" - const reasonLabel = - reason && extractedTokens === 0 ? ` — ${PRUNE_REASON_LABELS[reason]}` : "" - message += `\n\n▣ Pruning (${pruneTokenCounterStr}${extractedSuffix})${reasonLabel}` - - const itemLines = formatPrunedItemsList(pruneToolIds, toolMetadata, workingDirectory) - message += "\n" + itemLines.join("\n") + const details = buildPruneDetails( + state, + reason, + pruneToolIds, + toolMetadata, + workingDirectory, + distillation, + ) + + if (details) { + message += "\n\n" + details } return (message + formatExtracted(showDistillation ? distillation : undefined)).trim() @@ -99,6 +165,104 @@ export async function sendUnifiedNotification( showDistillation, ) + if (config.notificationType === "toast" && client?.tui?.showToast) { + const header = formatStatsHeader( + state.stats.totalPruneTokens, + state.stats.pruneTokenCounter, + ) + const title = header.split("\n")[0] + + const toastBody = buildToastBody(message, header) + + try { + await client.tui.showToast({ + body: { + title: title, + message: toastBody, + variant: "success", + duration: 4000, + }, + }) + return true + } catch (error) { + logger.warn("Failed to show toast, falling back to message", { error }) + } + } + + await sendIgnoredMessage(client, sessionId, message, params, logger) + return true +} + +export async function sendSquashNotification( + client: any, + logger: Logger, + config: PluginConfig, + state: SessionState, + sessionId: string, + toolIds: string[], + messageIds: string[], + topic: string, + summary: string, + startResult: any, + endResult: any, + totalMessages: number, + params: any, +): Promise { + if (config.pruneNotification === "off") { + return false + } + + let message: string + + if (config.pruneNotification === "minimal") { + message = formatStatsHeader(state.stats.totalPruneTokens, state.stats.pruneTokenCounter) + } else { + message = formatStatsHeader(state.stats.totalPruneTokens, state.stats.pruneTokenCounter) + + const pruneTokenCounterStr = `~${formatTokenCount(state.stats.pruneTokenCounter)}` + const progressBar = formatProgressBar( + totalMessages, + startResult.messageIndex, + endResult.messageIndex, + 25, + ) + message += `\n\n▣ Squashing (${pruneTokenCounterStr}) ${progressBar}` + message += `\n→ Topic: ${topic}` + message += `\n→ Items: ${messageIds.length} messages` + if (toolIds.length > 0) { + message += ` and ${toolIds.length} tools condensed` + } else { + message += ` condensed` + } + if (config.tools.squash.showSummary) { + message += `\n→ Summary: ${summary}` + } + } + + if (config.notificationType === "toast" && client?.tui?.showToast) { + const header = formatStatsHeader( + state.stats.totalPruneTokens, + state.stats.pruneTokenCounter, + ) + const title = header.split("\n")[0] + + const toastBody = buildToastBody(message, header) + + try { + await client.tui.showToast({ + body: { + title: title, + message: toastBody, + variant: "success", + duration: 4000, + }, + }) + return true + } catch (error) { + logger.warn("Failed to show toast, falling back to message", { error }) + } + } + await sendIgnoredMessage(client, sessionId, message, params, logger) return true } diff --git a/lib/ui/utils.ts b/lib/ui/utils.ts index 9134a5cf..2f6fc754 100644 --- a/lib/ui/utils.ts +++ b/lib/ui/utils.ts @@ -35,6 +35,29 @@ export function truncate(str: string, maxLen: number = 60): string { return str.slice(0, maxLen - 3) + "..." } +export function formatProgressBar( + total: number, + start: number, + end: number, + width: number = 20, +): string { + if (total <= 0) return `│${" ".repeat(width)}│` + + const startIdx = Math.floor((start / total) * width) + const endIdx = Math.min(width - 1, Math.floor((end / total) * width)) + + let bar = "" + for (let i = 0; i < width; i++) { + if (i >= startIdx && i <= endIdx) { + bar += "░" + } else { + bar += "█" + } + } + + return `│${bar}│` +} + export function shortenPath(input: string, workingDirectory?: string): string { const inPathMatch = input.match(/^(.+) in (.+)$/) if (inPathMatch) { diff --git a/package-lock.json b/package-lock.json index bcd4f877..10ecb124 100644 --- a/package-lock.json +++ b/package-lock.json @@ -1,17 +1,18 @@ { "name": "@tarquinen/opencode-dcp", - "version": "1.2.7", + "version": "1.3.1-beta.0", "lockfileVersion": 3, "requires": true, "packages": { "": { "name": "@tarquinen/opencode-dcp", - "version": "1.2.7", + "version": "1.3.1-beta.0", "license": "MIT", "dependencies": { "@anthropic-ai/tokenizer": "^0.0.4", "@opencode-ai/sdk": "^1.1.3", "jsonc-parser": "^3.3.1", + "ulid": "^3.0.2", "zod": "^4.1.13" }, "devDependencies": { @@ -676,6 +677,15 @@ "node": ">=14.17" } }, + "node_modules/ulid": { + "version": "3.0.2", + "resolved": "https://registry.npmjs.org/ulid/-/ulid-3.0.2.tgz", + "integrity": "sha512-yu26mwteFYzBAot7KVMqFGCVpsF6g8wXfJzQUHvu1no3+rRRSFcSV2nKeYvNPLD2J4b08jYBDhHUjeH0ygIl9w==", + "license": "MIT", + "bin": { + "ulid": "dist/cli.js" + } + }, "node_modules/undici-types": { "version": "7.16.0", "resolved": "https://registry.npmjs.org/undici-types/-/undici-types-7.16.0.tgz", diff --git a/package.json b/package.json index b3e886de..2bc7ff66 100644 --- a/package.json +++ b/package.json @@ -1,7 +1,7 @@ { "$schema": "https://json.schemastore.org/package.json", "name": "@tarquinen/opencode-dcp", - "version": "1.2.7", + "version": "1.3.1-beta.0", "type": "module", "description": "OpenCode plugin that optimizes token usage by pruning obsolete tool outputs from conversation context", "main": "./dist/index.js", @@ -43,6 +43,7 @@ "@anthropic-ai/tokenizer": "^0.0.4", "@opencode-ai/sdk": "^1.1.3", "jsonc-parser": "^3.3.1", + "ulid": "^3.0.2", "zod": "^4.1.13" }, "devDependencies": { diff --git a/tsconfig.json b/tsconfig.json index b30286cf..c20d8a54 100644 --- a/tsconfig.json +++ b/tsconfig.json @@ -3,7 +3,7 @@ "compilerOptions": { "target": "ES2022", "module": "ESNext", - "lib": ["ES2022"], + "lib": ["ES2023"], "moduleResolution": "bundler", "resolveJsonModule": true, "allowJs": true,